[Paper Review] Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

David Raposo, Sam Ritter+4 more|arXiv (Cornell University)|Apr 2, 2024

Natural Language Processing Techniques4 citations

TL;DR

This paper introduces Mixture-of-Depths (MoD), a method that dynamically allocates compute in transformer language models by allowing tokens to skip layers via a learned top-k routing mechanism, capping total FLOPs per layer. MoD models achieve equivalent or better performance than baseline transformers with up to 50% fewer FLOPs per forward pass and up to 50% faster inference, while maintaining training efficiency and hardware compatibility through a static computation graph.

ABSTRACT

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.

Motivation & Objective

• To address the inefficiency of uniform compute distribution in transformers, where all tokens receive equal FLOPs regardless of computational need.
• To enable dynamic, context-sensitive compute allocation across both token and layer dimensions without introducing dynamic computation graphs.
• To maintain hardware efficiency and training speed while reducing FLOPs per forward pass through intelligent routing of tokens through varying depths of the model.
• To explore whether learned routing can improve model efficiency and performance without sacrificing training time or accuracy.
• To demonstrate that compute can be saved at inference time without performance degradation, enabling faster inference or model scaling.

Proposed method

• A top-k routing mechanism selects which tokens participate in self-attention and MLP computations at each layer, based on a router head, with k fixed a priori.
• Tokens not selected by the router are passed through a residual connection, effectively skipping the transformer block computation and saving FLOPs.
• The method enforces a static computation graph with known tensor sizes, ensuring hardware efficiency and compatibility with existing accelerators.
• The routing decisions are learned end-to-end during training, with no auxiliary balancing loss required, leveraging top-k routing for stability.
• During autoregressive sampling, the router is distilled using a simple auxiliary classifier or loss to mimic training-time routing decisions without future token information.
• The approach allows trade-offs between speed and performance by adjusting the compute budget (k) per layer, enabling faster inference or model scaling.

Figure 1: Mixture-of-Depths Transformer. As in mixture-of-experts (MoE) transformers we use a router to choose among potential computational paths. But unlike in MoE transformers the possible choices are a standard block’s computation (i.e., self-attention and MLP) or a residual connection. Since so

Experimental results

Research questions

• RQ1 Can transformers learn to dynamically allocate compute across tokens and layers while maintaining a static computation graph?
• RQ2 Does dynamic compute allocation via routing lead to significant FLOP savings without performance degradation?
• RQ3 Can MoD models achieve better performance or faster inference compared to standard transformers under equivalent training FLOP budgets?
• RQ4 How does the absence of future token information during inference affect routing performance, and can it be mitigated?
• RQ5 Can the routing mechanism be extended to support different types of computation beyond standard blocks and residual connections?

Key findings

• MoD models achieve equivalent or better performance than baseline transformers when trained with the same FLOPs and wall-clock time, demonstrating improved compute efficiency.
• MoD models reduce FLOPs per forward pass by up to 50% compared to standard transformers, enabling faster inference without performance loss.
• The method enables up to 50% faster inference steps during post-training sampling, even when training FLOPs and time are kept constant.
• The top-k routing mechanism during training generalizes well to autoregressive inference, with minimal performance degradation, even when future token identities are unavailable.
• MoD models learn to route intelligently, skipping unnecessary computations for easier predictions while preserving long-range dependencies through strategic attention routing.
• The framework is extensible and can support routing to diverse computation types (e.g., memory lookup, tool use), enabling new architectural designs with controlled FLOP costs.

Figure 2: Routing schemes. Tokens are funnelled to the computational path of their choice when using token-choice routing (left). If a given path exceeds its capacity (e.g., more than two tokens in this example) then surplus tokens must be dropped (purple token). The exact token that is ultimately d