Understanding Mixture-of-Depths Routing in Transformers

Mixture-of-Depths (MoD) routing is an architectural optimization for Transformer artificial intelligence models that dynamically allocates computing power to individual tokens based on their complexity. This article explores what Mixture-of-Depths is, why traditional Transformers waste compute, and the exact mechanisms MoD uses to route tokens through or around specific model layers to drastically improve training and inference efficiency.

The Flaw in Standard Transformers

Traditional Transformer models use a uniform computation approach. Whether a token represents a simple, predictable word like "the" or a complex, context-heavy concept, it passes through every single layer of the neural network. This static allocation of compute is inefficient. Just as humans do not spend the same amount of mental energy on every word in a sentence, AI models do not need to apply maximum processing power to every token.

What is Mixture-of-Depths (MoD)?

Mixture-of-Depths is a technique developed to solve the inefficiency of uniform computation. While a related concept, Mixture-of-Experts (MoE), routes tokens to different specialized sub-networks within a layer, MoD decides whether a token needs to be processed by a layer at all. By allowing tokens to skip certain layers entirely, the model varies the "depth" of the network on a per-token basis. It forces the Transformer to focus its computational budget only on the tokens that require deep processing.

How MoD Allocates Compute Dynamically

Mixture-of-Depths relies on a strict, predefined compute budget and a lightweight routing mechanism to process tokens dynamically. The process occurs at each MoD-enabled layer in the following steps:

1. Setting the Capacity Cap

Before processing begins, engineers set a static capacity limit for the layer—for example, 50% of the total tokens in a sequence. This guarantees that the hardware will only ever process half of the tokens at that specific layer, locking in a predictable reduction in floating-point operations (FLOPs) and memory usage.

2. Scoring via the Router

As a sequence of tokens approaches a Transformer block, a small linear projection layer—called the router—evaluates each token's embedding. The router assigns a scalar weight (a single score) to every token. This score represents the router’s prediction of how much that specific token will benefit from being processed by the upcoming self-attention and feed-forward blocks.

3. Top-k Selection

The layer then ranks all the tokens in the sequence based on their router scores. Using a "top-k" selection algorithm, the model selects the highest-scoring tokens up to the predefined capacity limit (k).

4. Dual-Path Processing

Once the tokens are divided, they take one of two paths:

• The Compute Path: Tokens that score in the top-k enter the Transformer block. They undergo heavy calculations, including multi-head attention and dense matrix multiplications, allowing them to gather context and update their representations.
• The Residual Path: Tokens that fail to make the cut bypass the layer entirely. They travel along the residual connection (a highway that skips the layer) and require zero computational processing for that step.

After the top-k tokens finish their calculations, they are merged back with the skipped tokens, and the entire sequence moves to the next layer, where a new router repeats the process.

Solving the Autoregressive Challenge

While top-k routing works easily during training—because the model can see the entire sequence of tokens at once and compare their scores—it creates a problem during live inference (text generation). During inference, models generate text one token at a time. The model cannot know if the current token is in the "top 50%" if the future tokens do not exist yet.

To solve this, MoD trains a secondary, tiny predictor alongside the router. During training, this predictor learns to mimic the top-k threshold. During inference, instead of comparing a token against future tokens, the model simply asks the predictor: "Would this token likely make the top-k cut?" If the answer is yes, the token enters the layer; if no, it skips it.

The Benefits of Mixture-of-Depths

By dynamically allocating compute per token, Mixture-of-Depths achieves massive efficiency gains. Models utilizing MoD can match the performance and accuracy of standard, dense Transformers while using up to 50% fewer FLOPs during training. Furthermore, because fewer tokens are processed at each layer, the model's key-value (KV) cache is smaller, resulting in significantly faster generation speeds and lower memory requirements during deployment.