Mixture-of-Depths

Why It Matters for Modern Models

Mixture-of-depths (MoD) uses a learned router to send selected tokens through a transformer block while the other tokens skip it. MoD models matched baseline transformers trained with the same FLOPs and wall-clock time while needing only a fraction of the FLOPs per forward pass, and were upwards of 50% faster per step when sampling (Raposo et al., 2024).
It is related to mixture-of-experts: both use learned top-k routing, but MoE chooses which expert processes a token, while MoD chooses whether a block processes it at all.
Unlike earlier adaptive-computation schemes, in which the amount of computation varies with the input (Graves, 2016), MoD fixes the total compute per block in advance, which keeps the computation graph static and the cost predictable.

Questions to Think About

Why does multiplying the block output by the router weight matter for training, when the top-k choice itself is not differentiable?
Why is choosing the top k tokens over a whole sequence a problem for generating text one token at a time, and how does predicting the choice get around it?

Core Math

Key Equation

h^{\ell+1}_t = \begin{cases} h^\ell_t + r_t\, [f_\ell(H^\ell_{S_\ell})]_t & t \in S_\ell \\ h^\ell_t & t \notin S_\ell \end{cases}

Mixture-of-depths (MoD) lets each block of a transformer process only some of the tokens in a sequence; the others skip the block through the residual connection (Raposo et al., 2024).

At block $\ell$ , a learned router gives each of the $n$ tokens a scalar weight $r_t$ . Only the $k$ tokens with the largest weights, the set $S_\ell$ , go through the block's attention and MLP, written $f_\ell$ . Let $H^\ell_{S_\ell}$ be the selected token sequence; $[f_\ell(H^\ell_{S_\ell})]_t$ is its output at selected position $t$ :

h^{\ell+1}_t = \begin{cases} h^\ell_t + r_t\, [f_\ell(H^\ell_{S_\ell})]_t & t \in S_\ell \\ h^\ell_t & t \notin S_\ell \end{cases}

Multiplying the block output by $r_t$ puts the router on the gradient path, so it can be trained by ordinary backpropagation.

Because $k$ is fixed in advance, the compute per block and the tensor shapes are known before the input is seen. With a capacity of $k/n = 12.5\%$ on every other block, the setting Raposo et al. found best, 87.5% of tokens skip each routed block.

Choosing the top $k$ over a whole sequence looks at future tokens, which are not available when generating one token at a time. Raposo et al. therefore train either the router itself, with an auxiliary loss, or a small separate predictor to guess whether each token would have been in the top $k$ , and use that guess during sampling.

Key Papers