Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
This addresses the inefficiency in compute usage for language models, offering a practical speed-up for inference, though it is incremental as it builds on existing transformer architectures.
The paper tackles the uniform compute allocation in transformer-based language models by introducing a method that dynamically allocates FLOPs to specific positions in sequences, optimizing across layers. The result is models that match baseline performance with equivalent training FLOPs and times but require fewer FLOPs per forward pass and can be up to 50% faster during post-training sampling.
Transformer-based language models spread FLOPs uniformly across input sequences. In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth. Our method enforces a total compute budget by capping the number of tokens ($k$) that can participate in the self-attention and MLP computations at a given layer. The tokens to be processed are determined by the network using a top-$k$ routing mechanism. Since $k$ is defined a priori, this simple procedure uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. Nevertheless, since the identities of the $k$ tokens are fluid, this method can expend FLOPs non-uniformly across the time and model depth dimensions. Thus, compute expenditure is entirely predictable in sum total, but dynamic and context-sensitive at the token-level. Not only do models trained in this way learn to dynamically allocate compute, they do so efficiently. These models match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass, and can be upwards of 50\% faster to step during post-training sampling.