Mixture-of-Depths: How Dynamic Compute Allocation and Layer Skipping Scale LLM Efficiency

Standard transformer architectures allocate a uniform computational budget to every token in a sequence. Regardless of whether a model is processing a predictable punctuation mark, a common grammatical connective, or a mathematically dense reasoning step, every token undergoes an identical sequence of matrix multiplications across every multi-head attention and multilayer perceptron (MLP) block throughout the network's depth.

This static compute distribution is computationally inefficient. While human language and formal reasoning exhibit non-uniform information density, standard dense transformers expend identical FLOP budgets across positions. Conditional computation methods have historically attempted to resolve this inefficiency, but early techniques frequently introduced variable tensor dimensions or dynamic execution graphs that degrade hardware utilization on modern GPU and TPU accelerators.

In April 2024, researchers at Google DeepMind introduced Mixture-of-Depths (MoD), an architectural paradigm that enables transformers to dynamically route individual tokens through or around specific computational blocks while maintaining a strictly static compute budget and fixed tensor shapes. By selecting a dynamic subset of tokens to participate in self-attention and MLP computations at each layer, MoD models achieve baseline loss parity while requiring up to 50% fewer FLOPs per forward pass and stepping upwards of 60% faster during inference.

The Uniform Compute Bottleneck

In a conventional autoregressive transformer of L layers and hidden dimension d_model, an input sequence of length S passes through every layer sequentially. Every token position participates fully in the O(S^2 * d_model) attention computations and the O(S * d_model * d_ff) feed-forward projections at every depth.

This uniform allocation encounters two structural limitations:

• Token Difficulty Variance: Predicting the next token following a deterministic phrase prefix requires substantially less representational capacity than resolving long-range algorithmic dependencies or lexical ambiguities.
• Compute-Loss Inefficiencies: Training a larger parameter model on a fixed FLOP budget (as established in Chinchilla scaling laws) yields lower loss than over-allocating compute to redundant operations across all tokens.

Previous approaches to conditional computation attempted to resolve this through early-exit architectures, such as Confident Adaptive Language Modeling (CALM) or Depth-Adaptive Transformers. However, early exiting is monotonic and irreversible: once a token exits at layer k, it cannot participate in self-attention at layer k+m, preventing deeper layers from integrating early representations into subsequent multi-token reasoning.

Core Architecture and Learned Routing

Mixture-of-Depths addresses this limitation by introducing token-level routing across depth, allowing tokens to selectively bypass intermediate layers via residual connections and re-engage with the computation graph at subsequent layers.

Static Capacity Formulation

Rather than permitting a variable number of tokens to execute per layer, MoD defines a static capacity C in (0, 1] prior to training. For a sequence length S, the capacity defines the exact number of tokens k = floor(C * S) permitted to participate in the block's computation.

Because k is fixed before training begins:

• Tensor shapes inside the attention and MLP kernels remain completely static.
• Hardware scheduling on matrix-multiplication units (TPU MXUs and GPU Tensor Cores) avoids dynamic memory re-allocation.
• Total model FLOPs per forward pass are deterministically bounded.

Expert-Choice Routing Mechanism

For each routable transformer block at layer l, a router assigns a scalar affinity weight r_i to each token embedding x_i via a learned linear projection:

r_i = w_r^T * x_i

The model evaluates router weights across the sequence and identifies the k-th largest value, defining a threshold tau = TopK(R, k)_k, where R is the set of all router weights in the sequence.

The set of active tokens S_active consists of all positions satisfying r_i >= tau:

S_active = {i in {1, ..., S} | r_i >= tau}, where |S_active| = k

For each token i, the block output y_i is computed as:

• If token i is in S_active: y_i = r_i * f(X_active)_i + x_i
• If token i is not in S_active: y_i = x_i

Here, f(X_active) denotes the combined Self-Attention and MLP computation evaluated strictly on the gathered subset of active tokens. Multiplying the active computation output by the scalar router weight r_i ensures that router parameters receive non-zero gradients directly from the primary language modeling loss via backpropagation.

Tokens outside S_active pass through the identity residual path without incurring matrix multiplication FLOPs.

The Non-Causal Autoregressive Challenge

A key structural challenge in expert-choice routing is its non-causal nature during sequence-level top-k selection.

In standard training (prefilling), the router compares r_i against all tokens in the context window. However, during autoregressive token-by-token generation, future tokens are unknown. The model cannot determine whether the current token ranks within the top-k percentile of the ungenerated sequence.

The DeepMind research team evaluated two distinct mechanisms to preserve causal autoregressive execution during inference:

1. Auxiliary Binary Classification Loss

The router output r_i is trained with an auxiliary binary cross-entropy objective where the target is 1 if the token was selected in the sequence top-k and 0 otherwise:

L_router = - (1/S) * sum [ target_i * log(sigmoid(r_i)) + (1 - target_i) * log(1 - sigmoid(r_i)) ]

This objective shifts the sigmoid of router weights such that top-k tokens produce values sigmoid(r_i) > 0.5, while bypassed tokens produce sigmoid(r_i) < 0.5. During autoregressive generation, routing decisions are made causally by checking sigmoid(r_t) > 0.5 without requiring access to future tokens.

2. Causal Router Predictor

Alternatively, a small auxiliary MLP classifier is attached to the layer input with a stop-gradient operator. The classifier predicts whether token x_t will belong to the top-k set. Empirical results showed this auxiliary predictor rapidly reaches over 97% classification accuracy, resulting in negligible downstream perplexity degradation when switching from non-causal top-k to causal per-token thresholding during sampling.

Architectural Configurations: Alternation and MoDE

Empirical evaluations across isoFLOP regimes revealed critical architectural design rules for optimal compute routing.

Block Alternation

Routing every single transformer block degrades language modeling performance because information propagation across the sequence requires regular, unconstrained self-attention. DeepMind found that the optimal configuration routes every other block (alternating between a standard full-capacity transformer layer and an MoD routing layer).

With alternating blocks, reducing the routable block capacity down to C = 12.5% (meaning 87.5% of tokens bypass the block) achieved superior isoFLOP efficiency compared to standard dense baselines.

Input Tokens: [T1, T2, T3, T4, T5, T6, T7, T8]
     |
┌────▼────────────────────────────────────────┐
│ Layer 1: Standard Full Self-Attention + MLP │ (100% Capacity)
└────┬────────────────────────────────────────┘
     │
┌────▼────────────────────────────────────────┐
│ Layer 2: MoD Block (Top-k Router: C = 25%)  │
│   ├── Active:  [T2, T6] ──> [ Attn + MLP ] ─┤ (25% Compute)
│   └── Bypass:  [T1, T3, T4, T5, T7, T8] ────┤ (Residual Only)
└────┬────────────────────────────────────────┘
     │
┌────▼────────────────────────────────────────┐
│ Layer 3: Standard Full Self-Attention + MLP │ (100% Capacity)
└────┬────────────────────────────────────────┘
     │
┌────▼────────────────────────────────────────┐
│ Layer 4: MoD Block (Top-k Router: C = 25%)  │
│   ├── Active:  [T1, T5] ──> [ Attn + MLP ] ─┤ (25% Compute)
│   └── Bypass:  [T2, T3, T4, T6, T7, T8] ────┤ (Residual Only)
└────┬────────────────────────────────────────┘
     │
     ▼
Next Token Prediction

Mixture-of-Depths-and-Experts (MoDE)

MoD integrates naturally with Mixture-of-Experts (MoE) architectures, yielding Mixture-of-Depths-and-Experts (MoDE):

• Staged MoDE: Tokens first pass through an MoD router that determines whether they participate in self-attention. Active tokens then pass to an MoE router that assigns them to specific expert MLPs.
• Integrated MoDE: MoE routing integrates "no-op" residual paths directly into the expert candidate pool. Tokens can select either a domain-specific expert MLP or the identity residual path.

Integrated MoDE outperforms conventional MoE capacity reduction with token dropping because tokens explicitly optimize router weights toward the residual path rather than suffering uncoordinated drops.

Empirical Performance and Scaling Implications

IsoFLOP analyses conducted across compute budgets ranging from 6e18 to 1e20 FLOPs demonstrated consistent efficiency advantages for MoD models:

• Active FLOPs per Forward Pass: IsoFLOP-optimal MoD maintains total compute efficiency while fast MoD variants reduce forward-pass FLOPs by up to 50%.
• Pre-training Step Throughput: Smaller MoD variants that achieve baseline loss parity step up to 66% faster during pre-training due to reduced FLOP intensity.
• Inference Latency: By skipping self-attention and MLP layers for unselected tokens, autoregressive step times decrease by upwards of 50%.
• KV Cache Footprint: Tokens bypassing self-attention blocks avoid Key-Value projections and updates at intermediate layers, reducing memory bandwidth pressure.

When evaluating routing patterns, router weights strongly correlate with prediction entropy: tokens with high lexical uncertainty engage all available depth blocks, whereas predictable tokens consistently route around intermediate layers.

Implementation Trade-Offs and Summary

Mixture-of-Depths introduces a predictable framework for non-uniform compute allocation in language models:

• Static Compute and Tensor Geometry: Unlike dynamic halting mechanisms, MoD maintains static tensor sizes during training and execution.
• Complementary to Sparsity: MoD operates along the depth dimension, functioning orthogonally to width-wise conditional computation (MoE) and sequence-wise context compression.
• Training and Serving Efficiency: MoD models achieve lower validation loss at equivalent training compute budgets or equivalent loss with up to 50% fewer forward pass FLOPs.

As frontier language models scale, dynamic depth allocation provides a viable architectural mechanism to detach total parameter capacity from per-token inference expenditure.

Sources

• Raposo, D., Ritter, S., Richards, B., Lillicrap, T., Humphreys, P. C., & Santoro, A. (2024). Mixture-of-Depths: Dynamically allocating compute in transformer-based language models. arXiv:2404.02258.
• Hoffmann, J., Borgeaud, S., Mensch, A., et al. (2022). Training Compute-Optimal Large Language Models. arXiv:2203.15556.
• Schuster, T., Fisch, A., Gupta, J., et al. (2022). Confident Adaptive Language Modeling. arXiv:2207.07061.
• Shazeer, N., Mirhoseini, A., Maziarz, K., et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv:1701.06538.
• Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.