Mixture-of-Depths: Architectural Foundations, Routing Dynamics, and Variable Token Compute in Large Language Models
An authoritative engineering deep-dive into Mixture-of-Depths (MoD), sequence-level top-k routing mechanics, KV-cache dynamics, MoDE architectures, and variable token compute.
Introduction
Standard autoregressive Transformer architectures enforce a rigid, static computational graph. For an input sequence $X = [x_1, x_2, \dots, x_T] \in \mathbb{R}^{T \times d_{\text{model}}}$, every token is subjected to an identical sequence of $L$ sequential layers comprising Multi-Head Self-Attention (MHSA) and Feed-Forward Networks (FFN). This uniform compute allocation imposes a homogeneous floating-point operation (FLOP) cost of roughly $\mathcal{O}(24 L d_{\text{model}}^2 + 4 L T d_{\text{model}})$ per token, irrespective of the underlying token's information entropy, semantic density, or structural predictability.
In natural language, the predictive difficulty across a sequence is non-uniform. Punctuation, syntactic boilerplate, and common determiners require minimal contextual transformation to predict next-token distributions, whereas rare lexical entities, multi-hop reasoning steps, and disambiguation boundaries require deeper non-linear projections. Allocating identical compute across all positions yields severe computational inefficiency.
Standard Transformer:
[Token x_1] ───► [ Layer 1 ] ───► [ Layer 2 ] ───► ... ───► [ Layer L ] ───► Output
[Token x_2] ───► [ Layer 1 ] ───► [ Layer 2 ] ───► ... ───► [ Layer L ] ───► Output
[Token x_3] ───► [ Layer 1 ] ───► [ Layer 2 ] ───► ... ───► [ Layer L ] ───► Output
Mixture-of-Depths (MoD):
[Token x_1] ───► [ Layer 1 ] ───► ──── SKIP ───► ... ───► [ Layer L ] ───► Output
[Token x_2] ───► ──── SKIP ────► [ Layer 2 ] ───► ... ───► ──── SKIP ────► Output
[Token x_3] ───► [ Layer 1 ] ───► [ Layer 2 ] ───► ... ───► [ Layer L ] ───► Output
▲ ▲
Router Top-k Router Top-k
Mixture-of-Depths (MoD) fundamentally alters this paradigm. Rather than routing tokens across parallel specialized experts along the model's width—as seen in classic Mixture-of-Experts (MoE)—MoD routes tokens dynamically across the model's depth. By establishing a fixed per-block sequence capacity budget, MoD permits only a selected subset of tokens to participate in the computational sub-blocks of a given layer, while the remaining tokens bypass the layer via residual identity paths. This enables fine-grained, dynamic FLOP allocation per token without degrading downstream modeling capacity.
Theoretical Formulation: Top-k Sequence Routing and Dynamic Capacity
The core mechanic of Mixture-of-Depths revolves around conditioning layer execution on a learned routing scalar evaluated across the sequence dimension. Let $X^l \in \mathbb{R}^{T \times d}$ denote the hidden representations at layer $l$.
The Routing Function
Each MoD-instrumented layer defines a lightweight router parameterized by a projection vector $w_r^l \in \mathbb{R}^{d}$. For each token $x_i^l \in \mathbb{R}^d$ (where $i \in {1, \dots, T}$), the router computes an unnormalized scalar score $s_i^l$:
$$s_i^l = w_r^l \cdot x_i^l$$
The collection of scores across the sequence yields $S^l = [s_1^l, s_2^l, \dots, s_T^l]^\top \in \mathbb{R}^T$.
Hidden Tensor X^l: [T x d]
│
├───► [ Router: w_r^l ] ───► Scores S^l: [T]
│ │
│ Top-k Selection
│ (Capacity C = floor(α * T))
│ │
▼ ▼
Gather active tokens ◄───────────── Index Set S^l
tilde(X)^l: [C x d]
│
[ Transformer Block: f(.) ]
│
tilde(Y)^l: [C x d]
│
▼
Scatter-add with Residual ────────► Output Tensor X^(l+1): [T x d]
(Identity for non-selected tokens)
Sequence-Level Capacity Constraint
Unlike MoE models that enforce token-choice routing across $E$ parallel experts (where each token selects its top-$k$ experts), MoD implements expert-choice routing along the depth axis. The "expert" in this context is the single Transformer sub-block $f^l(\cdot)$ (which may encapsulate the MHSA, the MLP, or a combined block), and the router chooses which tokens among the sequence of length $T$ may enter it.
A static capacity factor $\alpha \in (0, 1]$ dictates the maximum token capacity $C$ of the layer:
$$C = \lfloor \alpha \cdot T \rfloor$$
The router identifies the top-$C$ tokens across the sequence:
$$\mathcal{S}^l = \operatorname{TopK}\left(S^l, C\right) = \left{ i \in {1, \dots, T} \mid s_i^l \ge \tau^l \right}$$
where $\tau^l$ is the adaptive score threshold defining the $C$-th largest element in $S^l$.
Layer Transformation and Residual Formulation
Let $\tilde{X}^l = \operatorname{Gather}(X^l, \mathcal{S}^l) \in \mathbb{R}^{C \times d}$ be the sub-tensor composed of tokens selected for computation. The layer transformation is evaluated strictly on this condensed tensor:
$$\tilde{Y}^l = f^l(\tilde{X}^l)$$
To preserve uninterrupted gradient propagation and maintain stable representation flow across depths, the computational output is re-scaled by the router weights and scattered back into the primary residual stream:
$$x_i^{l+1} = \begin{cases} x_i^l + \sigma(s_i^l) \cdot \tilde{y}_i^l, & \text{if } i \in \mathcal{S}^l \ x_i^l, & \text{if } i \notin \mathcal{S}^l \end{cases}$$
where $\sigma(\cdot)$ denotes the standard sigmoid activation function, providing a differentiable gating coefficient for active tokens.
Architectural Data Paths and KV-Cache Dynamics
Integrating dynamic depth routing introduces non-trivial architectural interactions with auto-regressive decoding, cross-token dependencies, and key-value (KV) caching.
Layer l KV-Cache Ingestion Scheme:
Token Indices: 1 2 3 4 5
Selected (Top-k): [YES] [NO] [YES] [NO] [YES]
│ │ │ │ │
Compute Q, K, V: ▼ │ ▼ │ ▼
(q,k,v) │ (q,k,v) │ (q,k,v)
│ │ │ │ │
KV-Cache Entry: [Slot 1] [PASSTHRU] [Slot 3] [PASSTHRU] [Slot 5]
Self-Attention Ingestion Under MoD
When self-attention is wrapped within an MoD boundary, two distinct implementation pathways emerge:
- Full Sequence Ingestion with Selective Query Projection: All $T$ tokens project keys ($K$) and values ($V$) to construct the complete attention context, but only the $C$ selected tokens execute query projections ($Q$) and attention aggregation.
- Selective Token Participation (Sparse Ingestion): Only selected tokens $\tilde{X}^l \in \mathbb{R}^{C \times d}$ compute $Q, K, V$. Non-selected tokens do not emit keys or values into the attention context of layer $l$.
The second variant delivers superior FLOP reduction. However, it requires structural mitigation during autoregressive decoding. If token $x_t$ skips layer $l$, downstream tokens at layer $l$ cannot attend directly to $x_t$'s representation at that specific depth. Information from $x_t$ must propagate exclusively via the residual stream until it reaches an active downstream layer $l + m$ where $x_t$ is selected or where other tokens read $x_t$'s residual updates.
Decoupled Routing: MHSA vs. MLP Blocks
Rather than treating the combined MHSA + FFN block as an atomic unit, MoD architectures achieve higher Pareto efficiency by decoupling routing across sub-modules:
- Router $\mathcal{R}_{\text{attn}}^l$: Selects capacity $C_{\text{attn}} = \lfloor \alpha_{\text{attn}} T \rfloor$ for Multi-Head Attention.
- Router $\mathcal{R}_{\text{mlp}}^l$: Selects capacity $C_{\text{mlp}} = \lfloor \alpha_{\text{mlp}} T \rfloor$ for the Feed-Forward Network.
Because the MLP accounts for approximately two-thirds of the non-embedding parameter count and per-token FLOPs in standard Transformer configurations ($8 d_{\text{model}}^2$ for standard MLPs, or $12 d_{\text{model}}^2$ for SwiGLU variants), allocating a smaller capacity factor to the MLP ($\alpha_{\text{mlp}} \approx 0.5$) while retaining higher capacity for attention ($\alpha_{\text{attn}} \approx 0.8$) preserves relational context while halving the bulk of matrix multiplication costs.
Unifying Depth and Width: The MoDE Architecture
Mixture-of-Depths and Mixture-of-Experts address orthogonal dimensions of conditional execution. Unifying them produces Mixture-of-Depths-and-Experts (MoDE), creating a compute paradigm that concurrently selects variable computational paths along both depth and width.
MoDE Routing Pipeline:
Token Vector x_i
│
▼
[ Depth Router ] ───► s_i < Threshold τ ───► [ Identity Residual Bypass ]
│
│ (s_i >= Threshold τ)
▼
[ Width Router ] ───► Softmax Gating over E Experts
│
┌─────┴─────────────────────┐
▼ ▼
[ Expert 1 ] ... [ Expert E ]
└─────┬─────────────────────┘
▼
[ Weighted Merge ]
│
▼
[ Residual Add ]
Mathematical Formulation of MoDE
In an MoDE layer, an incoming token $x_i$ is first evaluated by a top-$k$ depth router. If the token falls within the capacity allocation $C$, it is subsequently routed to one or more of $E$ parallel experts within the layer's FFN stage.
Let $\mathcal{G}_{\text{depth}}(x_i) \in {0, 1}$ represent the binary depth gating decision:
$$\mathcal{G}{\text{depth}}(x_i) = \mathbb{I}\left(w{\text{depth}} \cdot x_i \ge \tau\right)$$
Conditional on $\mathcal{G}{\text{depth}}(x_i) = 1$, the width routing gate $\mathcal{G}{\text{width}}(x_i) \in \mathbb{R}^E$ calculates classical top-$k_e$ routing over expert weights $W_e$:
$$g_e(x_i) = \operatorname{Softmax}\left(\operatorname{TopK_e}\left(W_{\text{width}} \cdot x_i, k_e\right)\right)_e$$
The consolidated layer forward pass is formalized as:
$$x_i^{l+1} = x_i^l + \mathcal{G}{\text{depth}}(x_i^l) \sum{e=1}^E g_e(x_i^l) \cdot \operatorname{Expert}_e\left(x_i^l\right)$$
Algorithmic Comparison
| Architectural Dimension | Dense Transformer | Mixture-of-Experts (MoE) | Mixture-of-Depths (MoD) | MoDE Unified |
|---|---|---|---|---|
| Compute Path | Static Depth / Static Width | Static Depth / Dynamic Width | Dynamic Depth / Static Width | Dynamic Depth / Dynamic Width |
| FLOP Allocation | Uniform per token | Uniform per token across active experts | Non-uniform per token | Non-uniform per token & expert |
| Active Params / Token | $100%$ | $\sim \frac{k_e}{E} \times \text{Params}_{\text{MLP}}$ | $\alpha \times \text{Params}_{\text{Block}}$ | $\alpha \times \frac{k_e}{E} \times \text{Params}_{\text{MLP}}$ |
| Routing Mechanism | None | Token-choice (typically Top-1/Top-2) | Sequence-choice (Top-$k$ over $T$) | Sequential: Top-$k$ Sequence $\to$ Top-$k_e$ Expert |
Training Dynamics, Differentiability, and Optimization Landscapes
Optimizing a model with discrete sequence-level $\operatorname{TopK}$ routing presents structural differentiability hurdles. The discrete selection operator $\mathcal{S} = \operatorname{TopK}(S, C)$ has zero derivative with respect to the router weights $w_r$ almost everywhere.
Gradient Propagation via Straight-Through and Soft Gating
To facilitate end-to-end backpropagation through the router weights without relying on high-variance reinforcement learning estimators (such as REINFORCE), MoD implements a continuous-discrete hybrid computational path.
- Forward Path: The discrete index set $\mathcal{S}$ is generated to extract active tokens $\tilde{X}$. The output of the sub-block $f(\tilde{x}_i)$ is multiplied by the continuous gating factor $\sigma(s_i)$:
$$y_i = x_i + \sigma(s_i) \cdot f(x_i) \quad \forall i \in \mathcal{S}$$
- Backward Path: While the inclusion in $\mathcal{S}$ is piecewise constant, the router parameter $w_r$ receives a gradient through the continuous scalar multiplier $\sigma(s_i^l)$:
$$\frac{\partial \mathcal{L}}{\partial w_r^l} = \sum_{i \in \mathcal{S}^l} \frac{\partial \mathcal{L}}{\partial x_i^{l+1}} \cdot f(x_i^l) \cdot \sigma'(s_i^l) \cdot x_i^l$$
Tokens that were close to the boundary threshold $\tau$ receive gradients that push their router scores higher or lower, dynamically updating routing boundaries across iterations.
Loss Function L
│
├───► (Through Residual Stream) ──────────────┐
│ │
└───► dL / d(x_i^(l+1)) │
│ ▼
┌───────┴────────┐ [ Layer Input ]
▼ ▼ ▲
dL / d(f(x_i)) dL / d(sigma(s_i)) │
│ │ │
│ ▼ │
│ sigma'(s_i) * x_i ──► dw_r^l ────┘
│ (Router Weight Update)
▼
(Block Weight Update)
Router Auxiliary Loss Formulations
Without explicit regularization, the router can destabilize into pathological equilibria where it consistently routes identical positional tokens regardless of semantic content, or displays extreme token churn across consecutive training steps. An auxiliary load-balancing loss $\mathcal{L}_{\text{router}}$ prevents optimization collapse:
$$\mathcal{L}{\text{router}} = \beta \cdot \frac{1}{L} \sum{l=1}^L D_{\text{KL}}\left(\bar{P}^l \parallel \mathcal{U}\right)$$
where $\bar{P}^l = \frac{1}{T}\sum_{i=1}^T \operatorname{Softmax}(s_i^l)$ represents the average routing distribution across the sequence, $\mathcal{U}$ represents a uniform distribution target, and $\beta$ is a balancing hyperparameter typically scaled to $0.01$.
Iso-FLOP Analysis
Empirical results demonstrate that MoD models trained under strict iso-FLOP constraints achieve lower cross-entropy validation loss compared to standard dense baselines.
Validation Loss
│
│ Dense Baseline (1.0x Compute Budget)
│ \
│ \ MoD Baseline (Same FLOP Budget, Deeper/Wider Model)
│ \ \
│ \ ▼
│ ▼ ● (Lower Perplexity Frontier)
│ ●
│
└──────────────────────────────────────────────── FLOPs
By reducing the per-step compute footprint by $50%$ ($\alpha = 0.5$), an MoD architecture can either:
- Train for $2\times$ the number of total tokens under an identical wall-clock and FLOP envelope.
- Scale the parameter capacity (depth $L$ and hidden dimension $d_{\text{model}}$) by roughly $1.6\times$ without exceeding the serving latency or training compute of the smaller dense baseline.
Hardware Realities, Inference Latency, and Systems Engineering
Translating theoretical FLOP reductions into realized wall-clock speedups requires careful hardware-level mapping. Modern tensor accelerators (GPUs, TPUs) favor static dense matrix multiplications (GEMMs) with uniform memory access patterns over irregular sparse memory operations.
Memory Layout and Kernel Execution Pipeline:
Prefill Phase (Sequence Length T):
Input Tensor [T x d] ──► Router ──► Compact via Scatter/Gather
│
▼
Dense Matrix GEMM [C x d] <── Peak Tensor Core
│ Utilization
▼
Unpack to Full Sequence [T x d]
Generation Phase (Autoregressive Step, T = 1):
Single Token x_t ──► Router Scalar s_t
│
┌────────────────┴────────────────┐
▼ ▼
(s_t >= Threshold) (s_t < Threshold)
Execute GEMM (Weights read) Bypass GEMM (Skip Weight Reads)
--> Compute-bound --> Memory-Bandwidth Preserved
Prefill Phase: Gather/Scatter Optimization
During the contextual prefill phase, the full sequence length $T$ is available simultaneously. Implementing naive conditional execution via branching causes thread divergence across GPU warps.
MoD addresses this through dense hardware gather operations:
- An index extraction kernel maps selected indices $\mathcal{S} \subset {1, \dots, T}$ to contiguous memory buffers:
$$\tilde{X} = \operatorname{Gather}(X, \mathcal{S}) \in \mathbb{R}^{C \times d}$$
- The GEMM kernel processes the condensed tensor $\tilde{X}$ at full hardware occupancy, leveraging tensor cores without sparsity-induced pipeline stalls.
- The resulting representations are re-scattered to the output array via a fused scatter-add kernel:
$$X_{\text{out}} = \operatorname{ScatterAdd}(X, \tilde{Y} \odot \sigma(S_{\mathcal{S}}), \mathcal{S})$$
Because $C = \lfloor \alpha T \rfloor$ is statically deterministic, memory layouts are fully pre-allocated at compile time, eliminating dynamic host-device synchronization barriers.
Decoding Phase: Memory Bandwidth Preservation
During single-token autoregressive generation ($T = 1$), the sequence-wide $\operatorname{TopK}$ selection operator cannot inspect future tokens. To execute conditional routing during generation, the model utilizes an absolute score threshold $\tau^l_{\text{learned}}$ determined during training or dynamically calibrated from the moving average of prefill thresholds.
When $s_t^l < \tau^l_{\text{learned}}$, the execution of layer $l$ is skipped entirely. This bypass provides critical systems advantages:
- Weight Eviction Avoidance: The high-bandwidth memory (HBM) read of parameter matrices $W_Q, W_K, W_V, W_{O}, W_{\text{gate}}, W_{\text{up}}, W_{\text{down}}$ for layer $l$ is completely bypassed.
- Bandwidth-Bound Speedups: Because single-batch token generation is strictly bounded by HBM read bandwidth rather than arithmetic compute throughput, skipping weight loading for $50%$ of layers nearly doubles token generation speed for memory-bound execution contexts.
Distributed Systems and Parallelism Compatibility
MoD natively composes with established distributed training paradigms:
- Tensor Parallelism (TP): The gather-compute-scatter operations execute within the local column/row parallel partitions without cross-device communication beyond standard
All-Reducecollectives at block boundaries. - Pipeline Parallelism (PP): Because sequence capacity $C$ is uniform across micro-batches, pipeline bubbles are not exacerbated by workload imbalances.
- Sequence Parallelism (SP): In setups utilizing sequence partitioning (e.g., Megatron-LM SP or RingAttention), gathering tokens across sequence partitions requires an all-to-all communication primitive to construct the active compute batch $C$, which introduces an interconnect latency overhead that must be balanced against the FLOP savings.
Conclusion
Mixture-of-Depths signals a fundamental shift from static, grid-like compute allocations toward non-uniform, content-adaptive execution graphs in autoregressive foundation models. By implementing sequence-level expert-choice routing across layer depths, MoD decouples raw model parameter capacity from the per-token computational cost.
Its capacity to seamlessly integrate with width-sparse routing paradigms (MoDE), maintain high hardware utilization via dense gather-scatter routines, and alleviate memory bandwidth bottlenecks during autoregressive generation positions dynamic depth routing as an essential architectural design pattern for the next generation of compute-optimal language models.
References
- Raposo, D., Ritter, C., Richards, B., Lillicrap, T., Humphreys, P. C., & Santoro, A. (2024). Mixture-of-Depths: Dynamically allocating compute in transformer-based language models. arXiv preprint. https://arxiv.org/abs/2404.02258
- Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research, 23(120), 1-39. https://arxiv.org/abs/2101.03961
- Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., & Kaiser, Ł. (2019). Universal Transformers. International Conference on Learning Representations (ICLR). https://arxiv.org/abs/1807.03819