Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

🎯 Pitch

Transformers waste FLOPs on tokens that don't need them—but Mixture-of-Depths models learn to skip 87.5% of tokens in alternating layers while outperforming isoFLOP baselines. This delivers 50% faster sampling with no performance loss, proving dynamic compute allocation can be both effective and hardware-friendly.

1. Executive Summary

This paper proposes Mixture-of-Depths (MoD) transformers, which learn to dynamically allocate FLOPs across token positions and model depth by using a learned top-k routing mechanism to select which tokens participate in self-attention and MLP computations at each layer—with unselected tokens bypassing the block entirely via residual connections—thereby enforcing a static, user-defined total compute budget per forward pass that is smaller than a vanilla transformer's. Trained on a language modeling objective using decoder-only transformers, MoD models match baseline performance at equivalent training FLOPs and wall-clock time while requiring a fraction of the FLOPs per forward pass, enabling step-time speedups upwards of 50% during post-training sampling. The optimal configuration aggressively routes 87.5% of tokens around every other block (using a 12.5% capacity on interleaved routing blocks) while still outperforming the isoFLOP-optimal baseline, establishing that transformers can recover or exceed baseline performance with substantially reduced per-forward-pass compute—provided the routing decisions are learned rather than stochastic.

2. Context and Motivation

The Core Problem: Uniform Compute Expenditure in Transformers

The fundamental inefficiency this paper tackles is deceptively simple: transformers expend the same amount of FLOPs on every token in every sequence, regardless of how much "thinking" that token actually requires. In a vanilla transformer, every token in a sequence participates fully in self-attention and MLP computations at every layer of the network. Whether a token represents a trivial definite article like "the" or a highly ambiguous word critical to parsing complex syntax, the computational cost is identical.

The authors state this directly in Section 1:

"Not all problems require the same amount of time or effort to solve. Analogously, in language modeling not all tokens and sequences require the same time or effort to accurately make a prediction. And yet, transformer models expend the same amount of compute per token in a forward pass."

This is not merely an aesthetic concern — it represents a genuine misallocation of computational resources. If some tokens are "easy" to predict (e.g., the untoken after a clear noun phrase, closed-class function words in deterministic syntactic contexts), then spending a full self-attention and MLP computation on them is wasteful. Conversely, if some tokens are "hard" (e.g., the first token of a novel named entity, a polysemous word in an ambiguous context, or the next token after a long-range dependency), they might genuinely benefit from deep processing. The ideal system would spend compute where it matters and save it where it doesn't.

Why This Matters: Training Cost, Inference Speed, and the Scaling Paradigm

This inefficiency has practical consequences that compound at scale. Transformer training runs are measured in exaFLOPs (the paper uses budgets of 6e18, 2e19, and 1e20 FLOPs), and each forward pass over billions of tokens involves massive matrix multiplications in self-attention and MLP layers. If a meaningful fraction of those FLOPs are wasted on tokens that don't require deep processing, then reducing that waste translates directly to:

Faster training: Fewer FLOPs per forward pass means shorter wall-clock training time on equivalent hardware — or the ability to train larger models within the same time budget.
Faster inference: For autoregressive generation (where tokens are produced one at a time and the model is stepped repeatedly), reducing per-step FLOPs by 50% means generating responses at roughly twice the speed. For deployed models serving millions of users, this has direct latency and cost implications.
More efficient use of scaling budgets: In the dominant paradigm of scaling laws (where model size and training duration are traded off under a fixed FLOP budget), any technique that reduces per-token compute without sacrificing performance effectively shifts the Pareto frontier — you can train a larger model or train for longer with the same total compute, or achieve the same performance with less compute.

There is also a theoretical significance tied to the nature of language understanding. If transformers can learn to route tokens intelligently — sending easy tokens around heavy computation and reserving deep processing for complex ones — this suggests something about the structure of language itself: that the difficulty of next-token prediction is heterogeneously distributed across positions, and that the network can discern this heterogeneity without explicit supervision. The routing decisions become a window into what the model finds "easy" versus "hard."

Prior Approaches and Where They Fall Short

The paper situates itself within the broader literature on conditional computation — the idea that neural networks should expend compute only when needed, as articulated by Bengio (2013) and explored extensively in subsequent work. The paper identifies several families of prior approaches, each with characteristic limitations:

Early-exit methods (Elbayad et al., 2019; Liu et al., 2021; Schuster et al., 2022) allow tokens to stop progressing through the transformer's depth once a sufficiently confident prediction is made, skipping all remaining layers. The critical limitation, as the authors highlight in Section 2, is that these methods are purely sequential and irreversible: once a token exits, it cannot re-engage with deeper layers, and later tokens cannot attend to it in those deeper layers. This means early-exit models have no ability to route a token around a middle block but then have it participate in deeper computations — something that might be important when the representation built in deeper layers is needed for certain tokens but not for others. The authors explicitly contrast MoD with early-exit:

"In MoD, unlike in early-exit methods, a token can skip middle layers, then be updated via self-attention with tokens that have gone through all the middle layers. We speculate that this might be a useful property."

Adaptive computation time and shared-weight methods (Graves, 2016; Dehghani et al., 2018; Simoulin and Crabbé, 2021) allow the network to dynamically decide how many computational steps to apply at each position, often by iterating through layers with shared weights. These approaches introduce dynamic computation graphs — the number of operations varies at runtime — which is incompatible with modern hardware accelerators (GPUs/TPUs) that are optimized for static, predictable tensor operations. The authors note this tension:

"general formulations of this challenging problem may not work well with existing hardware constraints since they tend to introduce dynamic computation graphs"

This is a central motivation for MoD's design: by fixing the number of tokens processed (capacity $k$ ) while letting the identity of those tokens be learned and context-dependent, MoD maintains static tensor sizes and computation graphs — the hardware always sees the same shape of operations, even though which tokens fill those shapes varies.

Mixture-of-Experts (MoE) transformers (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022) are the most direct technical ancestor of MoD. In MoE, a learned router sends tokens to one of several specialized MLP "experts," keeping total compute roughly constant (since all tokens go through some expert). MoD differs in a crucial way: rather than routing between functionally similar computations (different MLP experts), MoD routes between a full transformer block (self-attention + MLP) and a null operation (pure residual connection). This means MoD can achieve a genuine reduction in total FLOPs — not just a reallocation — because some tokens avoid the computationally expensive self-attention and MLP entirely. The authors frame this explicitly:

"Unlike other conditional computation approaches that try to conserve or expend additional compute, MoE transformers use conditional logic to route tokens to one of many expert MLPs while keeping total compute expenditure constant. Our mixture-of-depths method can be thought of as using the routing logic from MoE transformers, but rather than having multiple experts, MoD deploys a single expert which can be dynamically skipped."

CoLT5 (Ainslie et al., 2023) is perhaps the closest prior work, using conditional routing to select between heavy and light feedforward pathways and between full and sparse attention. However, CoLT5 was developed for an encoder-decoder setting, which sidesteps the causality problem that MoD must address in decoder-only transformers: in a non-causal encoder, you can look at the entire sequence to make routing decisions, but in autoregressive decoding, you cannot use future tokens to decide which current tokens to route. MoD's contribution of a predictive router — where an auxiliary network or loss learns to predict top-k membership using only past information — specifically addresses this gap, as the authors note:

"CoLT5 uses soft top-k for making routing decisions. However, CoLT5 focuses on a encoder-decoder setting, and thus does need to contend with the problem of efficient sequential decoding given the non-causal nature of the top-k operation. In contrast, our current work with MoD focuses on the decoder-only setting, and so we propose a predictive router to enable efficient inference for conditional computation in transformers."

Stochastic routing (analogous to dropout applied at the block level) is the simplest form of conditional compute but suffers from a critical weakness: it has no mechanism for distinguishing which tokens genuinely need processing. The paper includes this as a control and shows it "significantly under-performs relative to vanilla transformers" (Section 3.3), establishing that learned routing is essential — the network must be able to assign different processing levels to different tokens based on their content and context.

Where Existing Approaches Collectively Fall Short

Stepping back, the paper identifies a gap at the intersection of three desiderata that no prior method simultaneously satisfies:

Static computation graphs with predictable FLOPs (for hardware efficiency)
Genuine FLOP reduction (not just reallocation, as in MoE)
Decoder-only compatibility with autoregressive sampling (addressing the causality problem)

Early-exit methods satisfy (2) but fail on (1) and partially on (3) because tokens cannot re-engage after exiting. Adaptive computation time methods satisfy (2) and (3) but fail on (1) due to dynamic computation graphs. MoE satisfies (1) and (3) but fails on (2) because total FLOPs remain roughly constant — compute is reallocated across experts, not reduced. CoLT5 satisfies (1) and (2) but was designed for encoder-decoder models, leaving (3) unaddressed. Stochastic routing fails across the board — it doesn't effectively satisfy (2) because the performance degradation from arbitrary token dropping outweighs any FLOP savings.

MoD's explicit design goal is to hit all three: static compute budgets via per-block token capacities, genuine FLOP reduction by routing tokens around blocks, and decoder-only compatibility via the predictive router for autoregressive sampling.

How This Paper Positions Itself

The paper frames MoD not as a replacement for existing conditional computation methods but as a complementary tool in the efficiency toolkit, one that opens up a new axis for tuning the compute-performance tradeoff. The authors emphasize that MoD can be combined with MoE (creating "MoDE" models, Section 4.3) and with overtraining, and that the routing machinery is generic enough to support routing between arbitrary computation types — not just between a full block and a residual connection, but potentially between different specialized operations (memory lookups, tool use, etc.).

The empirical positioning is also notable. Rather than claiming absolute improvements over baseline transformers at any cost, the paper carefully anchors its comparisons to isoFLOP training budgets — meaning the total pre-training FLOPs are held constant between baseline and MoD models, and the question is whether MoD achieves better performance within that budget. This is a rigorous framing that avoids conflating gains from more compute with gains from smarter compute allocation. The paper's central empirical claim is that MoD shifts the isoFLOP curve "down and to the right" (Section 4.1): the optimal MoD model achieves lower loss and has more parameters than the optimal baseline at the same training FLOP budget, and smaller, faster MoD variants can match or exceed the optimal baseline's performance while being substantially faster to step.

3. Technical Approach

3.1 Reader Orientation

This is fundamentally a systems paper with an empirical analysis component — the authors propose, implement, and evaluate a specific architectural modification to transformer-based language models called Mixture-of-Depths (MoD). The core idea is: rather than having every token participate in every self-attention and MLP computation at every layer (as in vanilla transformers), the network uses a learned top-k routing mechanism at each block to select which tokens actually undergo the expensive computation, and which tokens simply bypass the block via a residual connection. The total number of tokens processed per block is fixed ahead of time (the capacity $k$ ), so the computation graph remains static and predictable, but which specific tokens are processed is determined dynamically and context-sensitively by the model's own learned routing weights. This enables genuine FLOP reduction — not just reallocation — while maintaining hardware-friendly static tensor shapes.

3.2 Big-Picture Architecture (Diagram in Words)

The MoD transformer has five interconnected mechanisms working together:

Token Embeddings ( $X^l$ ) — at each layer $l$ , we have a sequence of token representations of length $S$ . This is the standard input that any transformer block would receive.
Per-Block Router ( $w_\theta$ ) — a learned linear projection (a weight vector) that maps each token's embedding to a single scalar weight $r_i^l$ . This scalar expresses the router's preference for whether token $i$ should participate in the block's computations (higher value) or bypass them (lower value).
Top-k Selection ( $P_\beta(R^l)$ ) — a hard thresholding operation: given a user-specified capacity $C$ (an integer less than the sequence length $S$ ), the router weights across all tokens in the sequence are sorted, and only the top $C$ tokens with the highest router weights are selected to participate in the block's computations. The threshold is the $(1 - C/S)$ -th percentile of the router weights. This ensures exactly $C$ tokens are processed (static computation graph), but the identities of those $C$ tokens are fluid and learned.
Transformer Block Computation ( $f$ ) — the standard self-attention followed by MLP, but crucially operating only on the $C$ selected tokens rather than all $S$ tokens. This is where the FLOP savings come from: self-attention's query-key multiplication goes from $O(S^2)$ to $O(C^2)$ , and the MLP processes $C$ rather than $S$ token representations.
Routing Decision for Autoregressive Sampling — at inference time, the top-k operation is non-causal (it requires knowing all router weights across the sequence). To enable autoregressive generation, the paper trains either an auxiliary binary classifier (the "predictive router") or uses an auxiliary BCE loss that calibrates the router outputs so they can be thresholded independently without future information.

Information flows as follows: token embeddings enter layer $l$ → the router produces scalar weights for each token → the top $C$ tokens are selected → these $C$ tokens participate in self-attention and MLP, producing updated representations → the router weights for these selected tokens are multiplied onto the block's output (placing the router on the gradient path) → the updated representations are added via residual connection to the original embeddings → the $S - C$ tokens that were not selected pass through unchanged (pure residual connection) → the full sequence proceeds to the next layer.

3.3 Roadmap for the Deep Dive

First, the compute budget mechanism (Section 3.1 of the paper): how capacity $C$ defines a static, user-controlled FLOP reduction, and the intuitions behind why some tokens might not need processing.
Second, the routing around transformer blocks (Section 3.2): the two-path architecture — full block computation vs. null operation — and the mathematical formulation of how tokens are selectively processed.
Third, the routing schemes (Section 3.3): token-choice vs. expert-choice routing, why expert-choice was chosen, and the specific top-k mechanism that ensures exactly $C$ tokens are selected per block.
Fourth, the routing implementation (Section 3.4): the concrete mathematical operations — router weight computation, percentile-based thresholding, block output formulation — and the critical design choice of multiplying router weights onto the block output to place them on the gradient path.
Fifth, the autoregressive sampling solution (Section 3.5): the causality problem with top-k routing during inference, and the two solutions (auxiliary BCE loss and predictive router MLP) that enable independent per-token routing decisions.
Sixth, training configuration (Section 3.6): the hyperparameters, model sizes, and experimental design used for the isoFLOP analyses, including the key finding that routing every other block with 12.5% capacity is optimal.

3.4 Detailed, Sentence-Based Technical Breakdown

This section provides an exhaustive walkthrough of each mechanism in the MoD transformer, the mathematical formulation, and the design decisions that make it work.

The Compute Budget Mechanism (Capacity)

The fundamental abstraction that enables predictable FLOP reduction is the concept of capacity, borrowed from the Mixture-of-Experts literature but repurposed here for a different goal.

What capacity means in MoE: In Mixture-of-Experts transformers (Shazeer et al., 2017; Fedus et al., 2022), capacity defines the maximum number of tokens that can be routed to a given expert MLP. If more tokens try to route to an expert than its capacity permits, the excess tokens are "dropped" (their representations pass through the residual connection unchanged for that expert's computation). Since MoE transformers have multiple experts per layer and every token goes through some expert, the total capacity across all experts is approximately equal to the sequence length — meaning total FLOPs are roughly constant compared to a vanilla transformer. Capacity in MoE is primarily a load-balancing mechanism to prevent all tokens from routing to a single popular expert.

What capacity means in MoD: In MoD, capacity is used for the opposite purpose — not to reallocate fixed compute, but to genuinely reduce total compute. By setting a single block's capacity $C$ to some value less than the sequence length $S$ , the self-attention and MLP in that block process only $C$ tokens rather than $S$ tokens. The FLOP implications are concrete and dramatic:

For self-attention, the query-key dot-product matrix goes from $S \times S$ to $C \times C$ . The FLOP count scales quadratically with the number of tokens being attended over. If $C = S/2$ , then self-attention FLOPs are reduced to $(S/2)^2 / S^2 = 1/4$ of the vanilla transformer — a 75% reduction for that block's attention computation.
For the MLP, the reduction is linear rather than quadratic: the MLP processes $C$ token representations rather than $S$ , so FLOPs are reduced by a factor of $C/S$ . If $C = S/2$ , MLP FLOPs are halved.
The total per-block FLOP reduction is a combination of these two effects. In the specific configuration the paper finds optimal ( $C = 256$ out of $S = 2048$ total tokens, meaning $C/S = 0.125$ or 12.5% capacity), self-attention FLOPs are reduced to approximately $(0.125)^2 = 0.0156$ , or 1.56% of the original attention FLOPs — a 98.44% reduction for the routing blocks' attention. The MLP processes 12.5% as many tokens, so its FLOPs are 12.5% of the original.

Critically, because $C$ is defined a priori (before training begins), the computation graph has fixed, known tensor sizes. The hardware always allocates exactly enough memory for $C$ tokens to be processed, regardless of which specific tokens end up being selected. This is what makes MoD hardware-friendly: there are no dynamic resizes, no variable-length loops, no unpredictable memory allocations. The computation is static in shape but dynamic in content.

The intuition behind capacity reduction: The authors hypothesize that certain tokens in a sequence do not require as much processing as others, and that a trained network can learn to identify these tokens and route them around blocks. The key question is: how aggressively can you reduce capacity before the loss of processing for important tokens outweighs the benefits of reduced FLOPs? The paper's empirical finding — that 12.5% capacity on interleaved routing blocks is optimal — suggests that the overwhelming majority of tokens (87.5%) can safely skip computation at alternating layers without harming, and sometimes even improving, performance. This is a striking result: it implies that vanilla transformers are wasting a massive fraction of their compute on tokens that don't need it.

Relationship to stochastic routing: If capacity reduction were applied randomly — for example, by dropping tokens uniformly at random rather than learning which to drop — performance degrades significantly. The paper includes this as a control condition (Section 3.3 and Figure 3): MoD transformers that use stochastic routing (sampling router weights from a Gaussian distribution and then applying top-k) perform "drastically worse than both the baseline and normal MoD transformer." This establishes that capacity reduction alone is not sufficient; the learned selection of which tokens to process is essential.

Routing Around Transformer Blocks

The MoD transformer modifies each block (or every other block, depending on configuration) to offer tokens two mutually exclusive computational paths:

Path 1: Full block computation. The token participates in self-attention (attending over all other tokens that were also selected for this block) and the subsequent MLP. Its output representation is updated based on contextual information from other selected tokens. This path is computationally expensive but provides the transformer's core reasoning capability.

Path 2: Residual connection only (null operation). The token bypasses self-attention and the MLP entirely. Its output is identical to its input — the token's representation remains unchanged through this block. This is computationally cheap: a simple addition of zero (or near-zero) to the existing residual stream.

The total number of FLOPs per forward pass is fewer than a vanilla transformer whenever the capacity for Path 1 is less than the total number of tokens $S$ . The paper illustrates with a concrete example:

"if we were to set a block's capacity to $\frac{T}{2}$ (i.e., half the number of tokens as would be the case in a vanilla transformer) then query-times-key matrix multiplication during self-attention becomes 25% as FLOP-intensive as in a vanilla transformer $( (T/2)^2$ vs. $T^2$ )"

The key design tension is: shrinking capacity reduces FLOPs (faster), but risks degrading performance if important tokens are prevented from participating in computations. At the extremes: capacity equal to $S$ with all tokens routing to Path 1 recovers a vanilla transformer (no speedup, full performance); capacity equal to 0 with all tokens routing to Path 2 produces a model that never engages with the transformer's parameters (maximum speedup, near-zero performance). The empirical contribution is finding a sweet spot between these extremes where speed improves substantially and performance matches or exceeds the baseline.

Why this is different from dropout: Superficially, routing tokens around blocks might sound like structured dropout applied to entire transformer blocks. The crucial difference is conditioning: dropout randomly masks units independent of the input, while MoD's routing decisions are a learned function of the token's content and context. This means the network can learn systematic patterns — e.g., always process tokens at the end of a sequence (which might be more predictive of the next token), or always process tokens in certain syntactic positions — rather than randomly dropping tokens without regard for their importance. The stochastic routing control experiment (Figure 3) quantifies this difference: random routing performs drastically worse, demonstrating that the learned, content-dependent selection is responsible for MoD's ability to maintain performance under aggressive capacity reduction.

Routing Schemes: Token-Choice vs. Expert-Choice

The paper considers two families of learned routing mechanisms, both drawn from the Mixture-of-Experts literature, and selects one based on specific advantages for the MoD setting.

Token-choice routing (Section 3.3): In this scheme, a router produces a probability distribution for each token over the available computational paths. Each token is then assigned to the path with the highest probability for that token (or sampled from the distribution). This is the "natural" formulation — tokens "choose" where they want to go. However, token-choice routing has a fundamental load-balancing problem: there is no guarantee that tokens distribute themselves evenly across paths. If the router learns to send all tokens to the preferred path (e.g., all tokens want to participate in the block's computation), then capacity constraints are violated, and some tokens must be dropped arbitrarily (typically by sequence position or batch order). Conversely, if all tokens learn to avoid computation (all routing to the residual connection), the block becomes useless. To prevent these degenerate equilibria, token-choice routing requires an auxiliary balancing loss — an additional term in the training objective that penalizes uneven routing distributions. The authors note this explicitly: "auxiliary losses ensure that all tokens don't converge to the same path."

Expert-choice routing (Section 3.3): This scheme inverts the decision process. Instead of tokens choosing paths, each computational path chooses the top- $k$ tokens based on the tokens' router weights. Specifically, the router produces a scalar weight for each token expressing its "preference" for that path (in MoD's two-path case, there is effectively one weight per token, representing preference for the computation path over the residual path). The path then selects the $k$ tokens with the highest weights. This guarantees perfect load balance: exactly $k$ tokens are assigned to each path, by construction. The tradeoff is that some tokens may be "over-processed" (selected by multiple paths) or "under-processed" (selected by no path). In the standard MoE setting with multiple parallel experts, expert-choice routing risks sending the same token to multiple experts or losing tokens entirely. But in MoD's specific two-path case — one computation path, one null path — a single top- $k$ operation naturally partitions tokens into the top $k$ (route to computation) and the remaining $S - k$ (route to residual), with no possibility of overlap or coverage gaps.

Why expert-choice was selected for MoD: The paper provides three explicit justifications (Section 3.3):

No auxiliary balancing loss needed. Because the top- $k$ operation enforces exact capacity by construction, there is no risk of degenerate routing distributions, and thus no need for an auxiliary loss term that could interfere with the primary language modeling objective.
Relative routing weights can prioritize critical tokens. The top- $k$ operation depends on the magnitude of the router weights, meaning tokens compete with each other for limited computation slots. The router can "try to ensure that the most critical tokens are among the top- $k$ by setting their weight appropriately," which is qualitatively different from token-choice schemes where a token's routing decision is independent of other tokens' router outputs. This is particularly important for MoD since one path is a null operation — the router needs to ensure that genuinely important tokens are not among the tokens routed around the block.
Single top- $k$ for two paths eliminates over/under-processing. With exactly two paths (computation and residual) and a single top- $k$ operation with capacity $C$ , the $C$ tokens with the highest router weights go to the computation path, and the remaining $S - C$ tokens go to the residual path. This creates two mutually exclusive and exhaustive sets — no token is missed, and no token goes to both paths.

Routing Implementation: Mathematics and Mechanics

This is the core algorithmic section of the paper. We walk through the routing mechanism step by step, exactly as implemented.

Step 1: Router weight computation. Given a set of token embeddings at layer $l$ for a sequence of length $S$ :

$X^l = \{x_i^l \mid i \text{ is an integer, } 1 \leq i \leq S\}$

where $x_i^l \in \mathbb{R}^d$ is the $d$ -dimensional embedding for token $i$ at layer $l$ .

The router itself is a simple learned linear projection — a single weight vector $w_\theta \in \mathbb{R}^d$ — that maps each token embedding to a scalar:

$r_i^l = w_\theta^T x_i^l$

where $r_i^l \in \mathbb{R}$ is the router weight for token $i$ at layer $l$ .

What this computes: each token's $d$ -dimensional embedding is projected down to a single number via a dot product with a learned weight vector. This scalar expresses the router's preference for that token to participate in the block's computation — higher values mean "process this token," lower values mean "skip this token."

Why a simple linear projection: the router must be computationally cheap; if the router itself were expensive, it would defeat the purpose of reducing FLOPs. A single linear projection adds only $O(d)$ computation per token, which is negligible compared to the $O(d^2)$ or higher cost of self-attention and MLP.

Step 2: Percentile thresholding. Define $R^l = \{r_i^l \mid 1 \leq i \leq S\}$ as the set of all router weights for the sequence at layer $l$ . Let $P_\beta(R^l)$ be the $\beta$ -th percentile of this set, where:

$\beta = 1 - \frac{C}{S}$

with $C$ being the user-defined capacity — an integer less than $S$ specifying how many tokens will be processed by the block — and $S$ being the sequence length.

What $\beta$ represents: it is the fraction of tokens that will not be selected for computation. If $C = 256$ and $S = 2048$ , then $\beta = 1 - 256/2048 = 1 - 0.125 = 0.875$ . The threshold $P_{0.875}(R^l)$ is the router weight value at the 87.5th percentile — meaning 87.5% of tokens have router weights below this threshold, and 12.5% have router weights above it.

Why percentile-based thresholding: this ensures exactly $C$ tokens are selected, regardless of the absolute scale of the router weights. If the router learns to produce weights that are all very high or all very low, the percentile threshold adapts — it always selects exactly the top $C$ tokens relative to each other. This is crucial for maintaining a static computation graph: the hardware always processes exactly $C$ tokens, never more or less.

Step 3: Block output formulation. The output for each token $i$ at layer $l+1$ is defined as:

$x_i^{l+1} = \begin{cases} r_i^l \, f_i(\tilde{X}^l) + x_i^l, & \text{if } r_i^l > P_\beta(R^l) \\ x_i^l, & \text{if } r_i^l < P_\beta(R^l) \end{cases}$

where:

$\tilde{X}^l = \{x_j^l \mid r_j^l > P_\beta(R^l)\}$ is the set of tokens whose router weights exceed the threshold (the "selected" tokens), containing exactly $C$ elements.
$f$ is the standard transformer block computation — multi-head self-attention followed by MLP — but crucially operating only on $\tilde{X}^l$ .
$f_i(\tilde{X}^l)$ is the output of $f$ specifically for token $i$ (the $i$ -th position in the output).
$x_i^l$ is the residual connection: the token's original embedding before the block.

What this computes, case by case:

For tokens selected for computation ( $r_i^l > P_\beta(R^l)$ ): The token undergoes self-attention (attending over all other selected tokens in $\tilde{X}^l$ ) and the MLP. The resulting transformed representation $f_i(\tilde{X}^l)$ is then multiplied by the router weight $r_i^l$ before being added to the residual connection. This multiplication is a critical design choice — the router weight is on the gradient path for selected tokens, meaning gradient descent can adjust $w_\theta$ to increase or decrease this weight based on whether processing the token helped or hurt the language modeling objective. If multiplying by $r_i^l$ scales down the block's contribution for a token that didn't benefit from processing, the router can learn to reduce that token's weight in the future (potentially routing it around the block).

For tokens bypassing the block ( $r_i^l < P_\beta(R^l)$ ): The token's representation passes through unchanged. Its output is exactly its input — a pure residual connection with no transformation.

For tokens at the threshold ( $r_i^l = P_\beta(R^l)$ ): The paper's formulation uses strict inequalities ( $>$ and $<$ ). In practice, with floating-point router weights and a large sequence, exact equality is vanishingly rare, and the implementation resolves ties deterministically (e.g., by sequence order).

Why multiply the router weight onto the block output: The authors state that this "puts the router weights along the 'gradient path,' thus subjecting them to the forces of gradient descent through the course of the language modeling task." This means the router weight for a selected token directly modulates how much the block's computation affects that token's representation. If the router weight is small, the block's contribution is attenuated, and the router receives gradient signal to either increase the weight (if the block's computation was beneficial) or decrease it (if not). The authors experimented with versions where router weights were also included along the computational path for tokens that bypass the block, but found it "sufficient — and implementationally simpler — to only include the router weights along the computational path for those tokens that do not bypass the block's computations."

The self-attention nuance: A subtle but important point: the output $x_i^{l+1}$ for a token $i$ might depend on other tokens $x_{j \neq i}^l$ because of the self-attention operation within $f$ . Specifically, self-attention computes queries, keys, and values for all selected tokens in $\tilde{X}^l$ , and each selected token attends over all other selected tokens. This means that whether token $i$ itself is selected affects its output representation (it participates in self-attention as both query and key), but also that which other tokens are selected affects token $i$ 's output (it can only attend to selected tokens). The routing decisions create a dynamic attention pattern: at each routing block, the set of tokens that can be attended to is exactly the set of tokens that were selected for computation. This is a form of learned, content-dependent sparsity in the attention pattern, on top of the learned sparsity in which tokens get updated.

FLOP accounting: The cardinality of $\tilde{X}^l$ is exactly $C$ (the user-defined capacity). Therefore, self-attention's query-key multiplication involves matrices of size $C \times C$ rather than $S \times S$ , and the MLP processes $C$ token representations rather than $S$ . This is where the compute savings materialize. The paper notes: "the mixture-of-depths transformer accrues compute savings relative to the baseline because the input to the block's computations $f$ comprise fewer tokens than usual ( $C < S$ ), rendering the self-attention and MLP less expensive."

Routing frequency: The paper found that routing every other block (interleaving routing blocks with standard full-capacity blocks) is "crucial for strong performance" (Section 4.1). This means half the blocks in the transformer are standard (all $S$ tokens participate), and the other half are MoD routing blocks (only $C$ tokens participate). The full-capacity blocks provide periodic "full information" processing where every token can attend to every other token, while the routing blocks provide efficient partial processing. The paper does not provide a detailed theoretical justification for why interleaving is necessary, but plausibly: the full-capacity blocks ensure that the model never loses the ability to do global computation across the full sequence, while the routing blocks selectively deepen processing for the tokens that benefit from it.

Autoregressive Sampling: Solving the Causality Problem

The top-k routing mechanism described above has a fundamental problem during autoregressive generation: it is non-causal. To determine whether a given token $i$ is among the top $C$ tokens, the router needs access to all router weights $r_j^l$ for $j = 1, \dots, S$ in the sequence — including tokens that come after token $i$ . During training, this is fine because the entire sequence is available simultaneously (teacher forcing). During autoregressive sampling, however, future tokens don't exist yet — you cannot compute their router weights, so you cannot perform the top-k operation causally.

The problem in detail: When generating text token by token, at time step $t$ you have tokens $1, \dots, t$ available and are generating token $t+1$ . To route tokens at layer $l$ , you would need to know whether token $t$ (or any previous token) is among the top $C$ in the full sequence — but the full sequence doesn't exist yet. You could compute router weights for tokens $1, \dots, t$ and select the top $C$ among them, but this would mean capacity is always saturated (since $C$ could be larger than $t$ for short prefixes), and the routing decisions would change retroactively as the sequence grows, breaking the autoregressive property.

The paper proposes two solutions, both of which convert the non-causal top-k operation into a causal, per-token decision.

Solution 1: Auxiliary binary cross-entropy loss (BCE). During training, the router's outputs $r_i^l$ are used as logits for a binary classification task. The target is 1 if token $i$ was among the top- $C$ tokens in the full sequence (i.e., $r_i^l > P_\beta(R^l)$ ) and 0 otherwise. The auxiliary loss is standard binary cross-entropy:

$\mathcal{L}_{\text{aux}} = -\frac{1}{S} \sum_{i=1}^S \left[ t_i \log(\sigma(r_i^l)) + (1 - t_i) \log(1 - \sigma(r_i^l)) \right]$

where $t_i \in \{0, 1\}$ is the target (1 if selected, 0 if not) and $\sigma(r_i^l) = 1/(1 + e^{-r_i^l})$ is the sigmoid of the router weight.

What this computes: the standard binary cross-entropy between the sigmoid of each token's router weight and whether that token was actually selected by the non-causal top- $k$ operation. This loss pressures the router to produce weights such that $\sigma(r_i^l)$ is above 0.5 for tokens that would be in the top- $C$ and below 0.5 for tokens that would not.

Why this enables causal sampling: After training with this auxiliary loss, at inference time you can make per-token routing decisions independently: if $\sigma(r_i^l) > 0.5$ , route the token to computation; if $\sigma(r_i^l) < 0.5$ , route it to the residual. This decision depends only on token $i$ 's own router weight, not on any future tokens. The auxiliary loss has effectively "distilled" the global top- $k$ decision into a local thresholding rule. The authors report that "this auxiliary loss empirically affects the primary language modeling objective by approximately $0.2-0.3\%$ " — a small degradation, which they presumably measured as an increase in the primary loss compared to training without the auxiliary loss. The router outputs become centered such that approximately $\text{sigmoid}(r_i^l) \approx 0.5$ at the decision boundary, with selected tokens above and non-selected tokens below.

Solution 2: Auxiliary MLP predictor (predictive router). Instead of adding a loss term, a separate small MLP (a "second router") is trained to predict, from each token's representation (with gradient stopped), whether that token will be among the top- $C$ in the full sequence. This predictor receives the same input as the main router but has a stop-gradient applied, meaning its training does not affect the token representations or the primary language modeling objective.

What this computes: the predictive router is a binary classifier that approximates the top- $k$ selection function $t_i = \mathbb{1}[r_i^l > P_\beta(R^l)]$ using only per-token information. Like Solution 1, it produces a scalar that can be thresholded at inference time to make causal routing decisions.

Why this might be preferable: the authors note that this method "does not affect the language modeling objective, and empirically does not significantly impact the step speed." The predictive router is a small MLP (exact architecture unspecified in the paper) that adds minimal overhead at inference time.

Empirical validation of causal routing: The authors report (Section 4.2, Figure 6) that switching from the non-causal top- $k$ routing scheme to the predictor-based approach during autoregressive sampling "leads to minimal performance degradation." The auxiliary task is described as "relatively easy" and "quickly achieves 99% accuracy." This high accuracy suggests that the routing decisions are not heavily dependent on future tokens — the router weights for a given token contain enough information to predict whether that token will be selected without knowing the router weights of later tokens. This is somewhat surprising: one might expect that a token's importance in a sequence depends on what comes after it (e.g., a token that introduces a new topic might only be recognizable as important once the subsequent elaboration is seen). The high predictor accuracy suggests either that (a) the routing decisions are predominantly based on local, backward-looking features, or (b) the predictor learns to infer likely future content from past context (e.g., "this token looks like it starts a new paragraph, which usually means the subsequent tokens will be important, so I should route this token to computation now").

Training Configuration and Hyperparameters

The paper's experimental design is built around isoFLOP comparisons: holding total training FLOPs constant and comparing MoD variants against vanilla transformer baselines at the same computational budget.

Base hyperparameters (Section 3.6): All models use the same foundational configuration — cosine learning rate schedules equal to the number of training steps (no warmup specified), batch size 128, sequence length 2048 tokens. These are held constant across all comparisons; the only variations are in model architecture (number of layers, heads, embedding dimension) to create differently sized models for the isoFLOP analyses.

FLOP budgets: The paper reports on three training budgets: 6e18 FLOPs (for hyperparameter tuning), 2e19 FLOPs, and 1e20 FLOPs (for the full isoFLOP analysis). These are total pretraining FLOPs, meaning the cumulative computation across all forward and backward passes during training.

Model size range: For the isoFLOP analysis (Figure 4), models ranged from 60M to 3B parameters, with both baseline and MoD variants trained at multiple sizes to trace out the isoFLOP curve.

Optimal MoD configuration (determined from the 6e18 FLOP sweep, Figure 3):

Routing frequency: every other block (interleaved routing blocks and standard blocks). The paper states: "routing every other block was crucial for strong performance."
Capacity: 12.5% of the total sequence (256 tokens out of 2048 per batch element). The paper states: "we found that aggressive capacity reduction was best (gradual improvements were observed when reducing the capacity down to 12.5% of the total sequence, corresponding to 87.5% of tokens routing around blocks, with performance degrading beyond this point)."
Router architecture: a single learned linear projection (weight vector $w_\theta$ ) per routing block.
Routing scheme: expert-choice routing using the top- $k$ (top- $C$ ) mechanism, with the router weight multiplied onto the block output for selected tokens.

Stochastic routing baseline: For the control experiment, router weights were sampled from a Gaussian distribution rather than learned, and top- $k$ was applied to these random weights. This model "performed drastically worse than both the baseline and normal MoD transformer."

IsoFLOP analysis methodology (Figure 4): For each FLOP budget (6e18, 2e19, 1e20), multiple MoD variants of different sizes were trained using the fixed 12.5% capacity, every-other-block configuration. The same was done for vanilla transformer baselines. The isoFLOP curve shows training loss (presumably log perplexity on the training data, though the paper refers to it as "final log probability training objective") as a function of model size for a fixed total training FLOP budget. The "isoFLOP optimal" model is the one that minimizes this loss for that budget.

Key empirical findings from the training configuration:

The optimal MoD transformer "drags the baseline isoFLOP curve down and to the right" — achieving lower loss than the optimal baseline while also having more parameters.
Smaller MoD models exist that "are as- or better-performing than the optimal baseline model while being faster to step." The paper flags a specific 220M parameter MoD variant (model #3 in Figure 3) that slightly outperforms the 220M isoFLOP optimal baseline but is "upwards of 60% faster to step during training."
When run on equivalent hardware, these two models "take approximately the same amount of wall-clock time to train" — the faster per-step speed of the MoD variant (due to reduced FLOPs per forward pass) is offset by the need to train a model that is the same size (same number of parameters, so backpropagation costs are similar). This is a nuanced point: the speed gain is in the forward pass, but the backward pass and optimizer updates depend on the number of parameters, which is unchanged.

Wall-clock time vs. FLOPs: The paper notes that "FLOPs per forward pass" and "wall-clock step time" are "tightly correlated" in their experiments but are not identical. The forward pass FLOP reduction translates to genuinely faster step times because the computationally expensive operations (self-attention, MLP) are operating on fewer tokens. The authors provide normalized FLOPs per forward pass as a proxy for step time, noting that "a similar plot can be produced showing relative wall-clock step times and the same basic trend is present."

Scaling behavior: The paper finds that "it is better to add depth than to add width when adding FLOPs to the model" (Section 4.1). This means that when scaling up an MoD transformer to match the FLOPs per forward pass of the isoFLOP-optimal baseline (which the paper identifies as the condition for optimal MoD performance), one should preferentially add layers rather than increasing the embedding dimension or number of attention heads. The paper does not provide detailed reasoning for why depth is preferred, but this is consistent with the intuition that MoD creates a form of variable-depth processing (different tokens go through different numbers of effective layers), and deeper models can better exploit this flexibility.

Memory savings: The paper briefly notes that MoD transformers showed "memory savings relative to equivalently sized baseline models at larger sizes, with some variants requiring fewer total devices (i.e., a smaller TPU topology)." This is attributed to the reduced number of tokens being processed in self-attention and MLP layers, which reduces the size of intermediate activations and therefore the memory footprint. The paper speculates that these savings "could have significant positive effects in regards to the KV cache size during autoregressive sampling" — the key-value cache for self-attention in the routing blocks would only need to store $C$ entries rather than $S$ entries, a substantial reduction for inference.

Mixture-of-Depths-and-Experts (MoDE) integration (Section 4.3): The MoD routing mechanism can be combined with MoE in two ways:

Staged MoDE: MoD routing is applied first (deciding which tokens participate in self-attention), followed by MoE routing (deciding which expert MLP each selected token goes to). This allows tokens to skip self-attention entirely.
Integrated MoDE: The MoD and MoE routing are unified — tokens are routed to either one of the expert MLPs or to a "no-op" expert (equivalent to the residual path in MoD). This simplifies the routing machinery but means that even tokens routed to the residual path still participate in self-attention (since routing is applied after attention).

The paper reports that integrated MoDE "was distinctly better than simply reducing the capacity of experts in conventional MoE models, and relying on token dropping to implement residual routing." The interpretation is that "tokens explicitly learn to choose the residual path around the experts, as opposed to preferring an expert but being dropped when implemented as a capacity reduction." This suggests that learned routing to a no-op is qualitatively different from capacity-based dropping — in the latter case, tokens are forced to route to an expert they didn't prefer and are then dropped, potentially losing information, whereas in MoDE, tokens can learn that the residual path is the best choice for them and explicitly select it.

Design Choices and Their Justifications (Summary)

Expert-choice over token-choice routing: avoids auxiliary balancing loss; enables relative competition among tokens for limited computation slots; naturally partitions tokens into two mutually exclusive sets for the two-path case.
Router weight multiplication on block output: places router weights on the gradient path, enabling the router to learn from the downstream language modeling objective which tokens benefit from processing. This is an end-to-end learned signal rather than a separately defined auxiliary objective.
Interleaved routing blocks (every other block) rather than routing every block: ensures the model never loses the ability to do full-sequence computation, providing periodic "full information" layers. The paper does not deeply analyze why this is necessary but the empirical result is decisive.
Extreme capacity reduction (12.5%) rather than modest reduction: the paper found gradual improvements when reducing capacity all the way down to 12.5%, with degradation only beyond that point. This is a non-obvious result — one might have expected optimal capacity to be closer to 50% or 75% — and it implies that the overwhelming majority of token processing in vanilla transformers is redundant (at least for alternating layers).
Linear projection router rather than MLP router: computational efficiency — the router must not add significant overhead, and a single dot product adds minimal FLOPs.
Auxiliary predictor or BCE loss for autoregressive sampling rather than redesigning the routing mechanism to be inherently causal: this cleanly separates the training concern (non-causal top- $k$ for optimal routing) from the inference concern (causal per-token decisions), and the 99% predictor accuracy suggests the non-causal information is not essential for making good routing decisions.
Adding depth rather than width when scaling MoD models: this leverages the variable-depth property of MoD — since different tokens go through different numbers of effective layers, deeper models allow more gradations of processing depth and more flexibility in routing.

4. Key Insights and Innovations

Innovation 1: A New Axis for the Compute-Performance Tradeoff — Routing Between Radically Different Computational Paths, Not Just Between Equivalent Experts

The dominant paradigm for conditional computation in transformers, established by Mixture-of-Experts (Shazeer et al., 2017; Fedus et al., 2022; Lepikhin et al., 2020), routes tokens between computationally equivalent alternatives — multiple MLP experts that all perform the same type of operation (a feedforward transformation) but with different learned parameters. The total compute expenditure remains roughly constant because every token goes through some expert. The intellectual contribution of MoD is to break this symmetry: route between qualitatively different computational paths, one of which costs essentially nothing. Specifically, MoD routes between a full transformer block (self-attention + MLP) and a pure residual connection — a null operation that leaves the token's representation unchanged.

This is not an incremental tweak to MoE routing. It fundamentally changes what the routing decision means. In MoE, routing answers the question: "Which specialized MLP is best for this token?" In MoD, routing answers: "Does this token need any computation here at all?" The second question is both more fundamental — it targets the baseline assumption that all tokens require all layers — and more impactful, because answering "no" genuinely reduces total FLOPs rather than reallocating fixed FLOPs across experts. The paper draws this contrast explicitly (Section 2):

"MoE transformers use conditional logic to route tokens to one of many expert MLPs while keeping total compute expenditure constant. Our mixture-of-depths method can be thought of as using the routing logic from MoE transformers, but rather than having multiple experts, MoD deploys a single expert which can be dynamically skipped."

The significance goes beyond the specific implementation. This reframes conditional computation from a load-balancing problem (how to distribute tokens evenly across experts) to a resource allocation problem (how to decide which tokens get processing at all). The fact that a simple linear projection router can learn to make this binary choice — process or skip — for each token at each routing block, and that doing so at an aggressive 12.5% capacity on interleaved layers improves performance relative to an isoFLOP baseline (Figure 3), implies something non-obvious about the standard transformer architecture: a large fraction of the computation in vanilla transformers is not merely redundant, but actively counterproductive in a resource-constrained setting. If random routing (the stochastic control) destroys performance while learned routing improves it, the learned routing isn't just avoiding waste — it's concentrating compute where it has the highest marginal return on the language modeling objective.

This insight also connects MoD to a broader conceptual space that the paper sketches in the Discussion (Section 5): if routing can effectively choose between "full block" and "null," it can plausibly choose between arbitrary computation types — memory lookups, tool-use functions, specialized reasoning modules — with costs balanced by capacity. MoD thus serves as a proof of concept for heterogeneous computation routing in transformers, where different tokens undergo qualitatively different types of processing at different depths, all within a static compute budget.

Innovation 2: Recovering Causal Routing from Non-Causal Top-k via a Distillation Approach — and the Implicit Finding That Future Context Is Largely Unnecessary for Token Importance Decisions

The paper's solution to the autoregressive sampling problem is, at first glance, a straightforward engineering fix: since top-k over the full sequence is non-causal, train an auxiliary mechanism (a BCE loss or a small predictor MLP) that approximates the top-k decision using only per-token information. But the intellectual significance runs deeper than the mechanism itself. The auxiliary predictor achieves ~99% accuracy (Section 3.5, Figure 6), which carries a strong implicit finding: whether a token should be processed at a given layer is almost entirely decidable from its own representation and past context, with negligible dependence on future tokens.

This is not obvious a priori. One might expect that a token's importance in a sequence depends on what comes after it — a seemingly innocuous token might gain retrospective importance when it turns out to introduce a key concept, or a token at the start of a sentence might only be recognizable as "easy" or "hard" once the full sentence is seen. The high predictor accuracy suggests that the router weights $r_i^l = w_\theta^T x_i^l$ — which are computed from the token's representation at layer $l$ , itself a function of all preceding tokens through causal self-attention — already contain sufficient signal to determine whether that token will be among the top-k in the full sequence. In other words, past context is sufficient to predict future routing importance, even though the top-k operation itself uses future router weights to set the percentile threshold.

This finding matters because it cleanly validates the conceptual separation that MoD relies on: the routing decision is not a global optimization over the whole sequence that requires mutual information between distant tokens, but rather a local property of each token conditioned on its history. A token "knows" from its accumulated representation whether it's the kind of token that needs deep processing, independent of what specific tokens will appear later. This is what makes the distillation approach feasible — if the top-k selection were genuinely dependent on complex interactions between widely separated tokens (e.g., "process token 5 only if token 500 is a question mark"), the predictor would fail.

The paper also contributes a practical design pattern that generalizes beyond MoD: when a training-time operation is non-causal but empirically effective, you can distill it into a causal predictor for inference without redesigning the core mechanism. This is a form of amortization — the non-causal top-k provides a target that is computationally cheap to approximate at inference time because the underlying signal is mostly local. The two solutions (BCE loss vs. predictor MLP) offer a tradeoff between simplicity (BCE affects the primary loss by 0.2-0.3%) and modularity (the predictor MLP leaves the primary objective untouched), giving practitioners options depending on whether they prioritize performance purity or architectural cleanliness.

Innovation 3: The "Down and to the Right" IsoFLOP Shift as Evidence That Learned Sparsity Improves Scaling Efficiency — MoD Models Are Both Better and Bigger at the Same Training Budget

The paper's central empirical result — that MoD shifts the isoFLOP curve "down and to the right" (Section 4.1, Figure 4) — is more than a performance improvement. It demonstrates that per-token dynamic compute allocation improves the efficiency of the scaling process itself, not just the efficiency of individual forward passes. At a fixed training FLOP budget, the optimal MoD model achieves lower loss and has more parameters than the optimal vanilla baseline. This means the FLOP savings from routing are sufficient to fund a larger model (more parameters consume more FLOPs per token, but fewer tokens are processed per forward pass in routing blocks), and the net effect is positive — the larger, sparsely-computed model generalizes better than the smaller, densely-computed baseline.

This finding challenges an implicit assumption in the scaling laws literature (Hoffmann et al., 2022; Kaplan et al., 2020): that the primary axis for trading off compute and model size is the ratio of parameters to training tokens, with the per-token compute cost being a fixed function of model size. MoD introduces a new degree of freedom — forward-pass sparsity — that decouples model size from per-token FLOPs. You can have a larger model that costs less per forward pass than a smaller baseline because a large fraction of its parameters are selectively applied. The isoFLOP analysis (Figure 4) shows that this decoupling is beneficial across three orders of magnitude in training budget (6e18, 2e19, 1e20 FLOPs), suggesting it is a robust property of the scaling landscape rather than a small-budget artifact.

The practical consequence — that there exist MoD variants that are both faster to step and better-performing than the isoFLOP-optimal baseline (e.g., model #3 in Figure 3, which is ~60% faster while achieving slightly lower loss) — is significant for deployment economics. Prior to this work, getting a faster model typically meant either training a smaller model (sacrificing performance) or overtraining a small model on more data (which the paper acknowledges is "still possible with MoD transformers, and speed gains should compound"). MoD offers a third path: maintain or improve performance while reducing per-step FLOPs through learned compute allocation, without requiring additional training data.

The paper also contributes a useful heuristic for practitioners: the optimal MoD model is the one that uses as many FLOPs per forward pass as the isoFLOP-optimal baseline (Section 4.1). This provides a concrete recipe — given a target training budget, find the baseline-optimal model size, then scale the MoD variant (preferentially by adding depth) until its per-forward-pass FLOPs match the baseline, and you have the isoFLOP-optimal MoD configuration for that capacity setting. This reduces the hyperparameter search from a 2D grid (model size and capacity) to essentially a 1D search (find the matching FLOPs-per-forward-pass point), which is a substantial practical simplification.

Innovation 4: Verifying That Extreme Capacity Reduction (12.5%) Is Optimal — and What This Implies About Redundancy in Standard Transformers

The hyperparameter sweep in Figure 3 reveals an empirical result that is both surprising and theoretically suggestive: performance improves monotonically as capacity is reduced from 95% down to 12.5%, and only degrades below that point. This means that at a routing block, the model performs better when it processes only 12.5% of tokens and routes the other 87.5% around the block, compared to processing 50% or 75% or 95% of tokens.

This is not merely a "more sparsity is better up to a point" finding — the optimal point is extraordinarily aggressive, with the vast majority of tokens receiving no computation at alternating layers. The implication is that standard transformers are dramatically over-computing at the token level. If a model can achieve better performance by skipping 87.5% of tokens at every other layer, then for those layers in a vanilla transformer, the tokens that were "unnecessarily" processed were not just neutral — their processing was actively interfering with the model's ability to form good representations, perhaps by introducing noise or by consuming capacity in the residual stream that could have been better used for the truly important tokens.

This finding also provides an empirical anchor for a conceptual claim the paper makes in the Discussion: that "FLOPs may be inefficiently used in vanilla transformer models, and that there may be more efficient ways for them to be expended." The 12.5% capacity number quantifies the scale of the inefficiency — at least for interleaved routing blocks, roughly 7/8 of the per-token computation is unnecessary (and possibly harmful) when the network is allowed to choose which tokens to process. This is a much stronger claim than the paper's modest framing of "not all tokens require the same amount of time or effort" — it suggests that at alternating layers, almost no tokens require computation, and the network's performance is limited not by insufficient processing but by the inability of vanilla transformers to withhold processing from tokens that don't need it.

The corollary that "routing every other block was crucial for strong performance" (Section 4.1) adds nuance: full-capacity blocks remain essential, presumably because they provide the global information exchange that allows the routing blocks to be so aggressively sparse. The architecture that works is not uniform sparsity but alternating dense and sparse computation, where the dense blocks ensure that every token periodically gets updated with full-context information, and the sparse blocks selectively deepen processing for the subset of tokens that benefit from it. This alternating pattern is reminiscent of how convolutional networks alternate between spatial pooling (reducing resolution) and convolutional layers (deepening features), but adapted to the token dimension in transformers.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper uses a standard language modeling corpus for pretraining—the exact dataset is not named in the main text, but all experiments train decoder-only transformers on next-token prediction. Training data consists of sequences of length 2048 tokens; autoregressive evaluation is performed on a held-out set of 256,000 sequences (500M tokens). The lack of a named dataset (e.g., C4, The Pile, etc.) is a notable omission, as it makes exact replication difficult without additional information.
Base model(s). The paper trains decoder-only transformers from scratch, varying model size from 60M to 3B parameters across experiments. No pretrained model is used; all comparisons are between MoD variants and vanilla transformer baselines trained under identical conditions. The models use standard transformer architecture with multi-head self-attention and MLP blocks. The paper does not specify the exact number of heads or embedding dimensions for each model size, though these would be determined by the standard scaling configurations used.
Metrics. The primary metric is training loss on the language modeling objective, measured as "final log probability" (Section 4.1)—this is essentially the negative log-likelihood (or log perplexity, up to a constant factor). For autoregressive evaluation (Section 4.2), the same metric is computed on held-out data. When comparing models, the paper focuses on isoFLOP comparisons: models are trained with the same total pretraining FLOP budget, and the model achieving the lowest training loss at that budget is considered optimal.
Baselines. The paper uses three types of baselines: (1) Vanilla transformer — standard decoder-only transformers with the same architecture but no routing (every token participates in every block at every layer). These are trained at multiple sizes to trace out baseline isoFLOP curves. (2) Stochastic routing MoD — same architecture as MoD but with router weights sampled from a Gaussian distribution rather than learned, applying the same top- $k$ capacity reduction. This is explicitly described as a "control" to isolate the effect of learned routing. (3) Baselines at multiple FLOP budgets — for each of the three training budgets (6e18, 2e19, 1e20 FLOPs), multiple baseline model sizes are trained to establish the isoFLOP-optimal baseline configuration.
Generation budget / compute accounting. Compute is measured in total pretraining FLOPs (forward + backward passes). The paper uses three budgets: 6e18, 2e19, and 1e20 FLOPs. Within each budget, models of different sizes are trained to convergence, and the model minimizing training loss is the isoFLOP-optimal configuration. Per-forward-pass FLOPs are also reported (normalized relative to the isoFLOP-optimal baseline) as a proxy for step speed. The paper notes that per-forward-pass FLOPs and wall-clock step time are "tightly correlated" in their experiments. Importantly, the comparison is genuinely isoFLOP: MoD models with reduced per-forward-pass FLOPs do not get to train for more steps within the same total budget—the budget is total FLOPs, so a model with 50% fewer FLOPs per forward pass would need to train for 2× more steps to consume the same total budget. The paper accounts for this, and the headline result is that MoD models match or exceed baseline performance within the fixed total FLOP budget despite (or because of) using fewer FLOPs per step.
Cross-validation / statistical protocol. The paper does not report cross-validation or statistical significance testing. Results are reported as point estimates from single training runs at each configuration. The hyperparameter sweep (Figure 3) trains multiple MoD variants at 6e18 FLOPs to select the optimal configuration (12.5% capacity, every-other-block routing), and this configuration is then used at larger FLOP budgets without further tuning. The autoregressive evaluation (Figure 6) uses 256,000 held-out sequences for evaluation.

Main Quantitative Results

Hyperparameter Tuning: Capacity and Routing Frequency Sweep (Figure 3)

The initial sweep at 6e18 FLOPs establishes the optimal MoD configuration. The key findings are:

Capacity: Performance improves monotonically as capacity is reduced from 95% down to 12.5% of the sequence. The paper reports "gradual improvements were observed when reducing the capacity down to 12.5% of the total sequence... with performance degrading beyond this point" (Section 4.1). At 12.5% capacity with sequence length 2048, this means only 256 tokens are processed by each routing block, while 1792 tokens (87.5%) bypass it. The performance improvement from reducing capacity is a central surprise — the model gets better while doing less computation per block.

Routing frequency: The paper states that "routing every other block was crucial for strong performance" (Section 4.1). The optimal configuration interleaves MoD routing blocks (with 12.5% capacity) with standard full-capacity blocks. The paper does not report results for routing every block or routing at other frequencies, but the strong language ("crucial") indicates that the performance gap was substantial.

Learned vs. stochastic routing: MoD transformers with stochastic routing (Gaussian-sampled router weights) perform "drastically worse than both the baseline and normal MoD transformer" (Section 4.1, Figure 3). The stochastic model is shown as a single data point in Figure 3 (left panel, positioned significantly above both the baseline isoFLOP curve and the MoD variants), establishing that capacity reduction alone is harmful without learned token selection.

The "down and to the right" shift: Figure 3 (left panel) shows the isoFLOP curve for MoD variants (grey shaded region) positioned below and to the right of the baseline isoFLOP curve. The optimal MoD model achieves lower training loss than the optimal baseline at the same total FLOP budget, and does so with more parameters. This is the core empirical finding that the paper scales up to larger budgets.

Speed-performance tradeoff discovery: Figure 3 identifies a specific MoD variant — model #3, a 220M parameter configuration — that "slightly outperforms the isoFLOP optimal baseline (also 220M, model #1), but is upwards of 60% faster to step during training" (Section 4.1). The learning curves (Figure 3, right panel) show model #3 and model #1 converging to nearly identical final loss, with model #3 marginally lower. This demonstrates that MoD can achieve performance parity with the baseline while offering substantial speed advantages.

IsoFLOP Scaling Analysis (Figure 4)

The paper extends the analysis to 2e19 and 1e20 FLOPs using the fixed 12.5% capacity, every-other-block configuration. Figure 4 shows isoFLOP curves at three budgets, and the findings are:

The "down and to the right" shift persists across scales: At all three FLOP budgets (6e18, 2e19, 1e20), the optimal MoD transformer achieves lower training loss than the optimal baseline transformer, and has more parameters. This is shown in the left panel of Figure 4, where the MoD curves sit below the baseline curves and extend further to the right (higher parameter counts). The exact loss values are not reported in the text, but the visual separation between the curves is consistent and clear.

Speed advantages are preserved: The right panel of Figure 4 shows normalized FLOPs per forward pass (relative to the isoFLOP-optimal baseline at each budget). At each of the three budgets, there exist MoD variants that achieve lower loss than the isoFLOP-optimal baseline while also requiring fewer FLOPs per forward pass. Specifically:

At 6e18 FLOPs: MoD variants at approximately 0.6-0.8× the baseline's forward-pass FLOPs achieve lower loss.
At 2e19 FLOPs: The spread is similar, with the optimal MoD requiring roughly 0.7× the baseline's forward-pass FLOPs.
At 1e20 FLOPs: The pattern holds, though the paper notes the advantage narrows slightly at larger scales.

Optimal MoD uses matched forward-pass FLOPs: The paper reports a key regularity: "the optimal MoD transformer is that which uses as many FLOPs per forward pass as the isoFLOP optimal baseline" (Section 4.1). This is visible in Figure 4 (right panel): the minimum of each MoD curve (lowest loss) occurs at the point where normalized FLOPs per forward pass equals 1.0 — i.e., matching the baseline's per-step compute. This provides a practical heuristic for practitioners: to find the isoFLOP-optimal MoD configuration, scale the model size (preferentially by adding depth) until its per-forward-pass FLOPs match the optimal baseline's, and you have the best configuration for that capacity setting.

Scaling strategy (depth vs. width): The paper states that "it is better to add depth than to add width when adding FLOPs to the model" (Section 4.1). No ablation is shown comparing depth-scaling vs. width-scaling, but the statement is presented as an empirical finding from the isoFLOP analysis. This is consistent with the intuition that MoD creates variable-depth processing and deeper models can better exploit this flexibility.

Autoregressive Evaluation (Figure 6)

Figure 6 addresses the critical question: does the performance demonstrated during training (with non-causal top-k routing) carry over to autoregressive sampling (with causal routing)?

Predictor accuracy: The auxiliary classifier ("predictive router" or BCE-calibrated router) that enables causal routing achieves "upwards of 97% accurate soon into training" for the 6e18 FLOP model configuration (Section 3.5, Figure 6). Exact accuracy curves are not shown, but the paper reports the 97% figure and 99% for the final accuracy (Section 3.5, text).

Performance retention: When switching from non-causal top-k routing to predictor-based causal routing during autoregressive evaluation, the paper reports "little performance degradation" (Section 4.2). Figure 6 shows the training loss and autoregressive evaluation loss for multiple MoD variants compared to baselines. The key result: "As in the training setting, there exist MoD variants that are better performing than the isoFLOP-optimal baseline, while requiring fewer FLOPs per forward pass" during autoregressive evaluation. This confirms that the FLOP savings and performance advantages are not artifacts of teacher-forcing and translate to the generation setting.

Comparison to baselines: Figure 6 plots both training and evaluation metrics against normalized FLOPs per forward pass. The MoD evaluation points (triangles in the right panel) sit below the baseline evaluation points at several speed points, meaning MoD achieves lower evaluation loss while being faster to step. The exact loss values are not reported numerically.

Mixture-of-Depths-and-Experts (MoDE) Integration (Figure 7)

Figure 7 shows results for combining MoD routing with MoE (Mixture-of-Experts), creating MoDE models. Two integration strategies are tested:

Staged MoDE: MoD routing is applied first (selecting which tokens participate in self-attention), followed by standard MoE routing (selecting which expert MLP each selected token goes to). This allows tokens to skip self-attention entirely.

Integrated MoDE: A single routing operation funnels tokens to either one of the expert MLPs or to a "no-op" expert (equivalent to the residual path). This means all tokens still participate in self-attention (since routing is applied after attention), but some tokens skip the expert MLP computation.

Key finding: The paper reports that "implementing MoDE in the integrated manner was distinctly better than simply reducing the capacity of experts in conventional MoE models, and relying on token dropping to implement residual routing" (Section 4.3). The interpretation is that with integrated MoDE routing, "tokens explicitly learn to choose the residual path around the experts, as opposed to preferring an expert but being dropped when implemented as a capacity reduction." This is a qualitative finding about the learning dynamics: explicit routing to a no-op produces better results than capacity-based dropping because the router can genuinely learn that the residual path is optimal for some tokens, rather than having tokens that wanted computation get arbitrarily dropped when capacity is exceeded.

Figure 7 shows performance curves for both staged and integrated MoDE compared to pure MoD and pure MoE baselines. The integrated MoDE achieves the best performance among the variants shown, suggesting that the routing benefits of MoD and the expert-specialization benefits of MoE can compound.

Memory and Hardware Efficiency (Not Quantified in Detail)

The paper briefly notes memory savings from MoD but does not provide quantitative analysis. The observation in Section 4.1 is: "We noticed that MoD transformers had memory savings relative to equivalently sized baseline models at larger sizes, with some variants requiring fewer total devices (i.e., a smaller TPU topology)." The paper speculates that this "could have significant positive effects in regards to the KV cache size during autoregressive sampling" because the key-value cache for routing blocks would only store entries for $C$ tokens rather than $S$ tokens. However, no measurements of actual KV cache size, memory footprint, or device count are reported.

Ablation Studies and Robustness Checks

Capacity sweep (12.5% to 95%): The paper reports a range of capacities tested during the 6e18 FLOP sweep (Figure 3). Capacity was varied from 95% down to 12.5% of the sequence length. The finding is that performance improves as capacity decreases, with the optimum at 12.5% and degradation below that point. The sweep is shown in Figure 3 (left), where models with lower capacity (higher sparsity) generally achieve lower loss. The paper does not report intermediate capacity points or the exact loss values, making it difficult to assess how sharp the optimum is — whether 15% or 10% would perform similarly, or whether 12.5% is a precise sweet spot.
Routing frequency (every block vs. every other block): The paper states that routing every other block was "crucial" but does not show the comparison data. Figure 3 presumably includes models with different routing frequencies, but the specific results for routing-every-block are not discussed. This is a significant omission, as the performance penalty of routing every block would help characterize how much full-capacity computation is necessary.
Stochastic routing (learned vs. random): The stochastic routing baseline — using top-k on Gaussian-sampled router weights — performs "drastically worse" than learned routing (Figure 3). This establishes that the performance of MoD is not attributable to structured dropout or regularization effects from token dropping. However, the paper does not explore why stochastic routing fails — whether it's because important tokens get dropped, because unimportant tokens get processed (wasting capacity), or both. A finer-grained analysis (e.g., measuring the correlation between stochastic selection and token "importance" as determined post-hoc) is not provided.
Router weight multiplication (on vs. off the gradient path): The paper mentions that router weights are multiplied onto the block output for selected tokens, placing them "along the gradient path" (Section 3.4). The authors note they "experimented with versions where the router weights are also included along the computational path for those tokens that bypass the block's computations, but it seems to be sufficient — and implementationally simpler — to only include the router weights along the computational path for those tokens that do not bypass." No quantitative comparison is shown. This is a minor ablation that the paper treats as an implementation detail rather than a central finding.
Autoregressive routing: BCE loss vs. predictor MLP: Two solutions are proposed for causal routing at inference (Section 3.5): an auxiliary BCE loss that calibrates router outputs, and a separate predictor MLP trained with stop-gradient. The paper states that the BCE loss "empirically affects the primary language modeling objective by approximately 0.2-0.3%" while the predictor MLP "does not affect the language modeling objective." However, Figure 6 does not distinguish between these two methods — the autoregressive evaluation results are presumably from one of them (likely the predictor, since it doesn't affect the primary loss), but this is not explicitly stated. A direct comparison of the two approaches would clarify whether the 0.2-0.3% degradation from the BCE approach is visible in downstream performance or is negligible in practice.
Token-level routing analysis (Figure 5): The paper provides a qualitative analysis of routing decisions in a trained MoD model. Figure 5 (left) visualizes which tokens are selected (dark) vs. bypassed (light) across layers for a truncated 64-token sequence. The paper notes a "vertical band of dark blue towards the end of the sequence," suggesting the model preferentially routes late-sequence tokens to computation. Preliminary analysis "suggest that the tokens that engage with blocks more frequently are correlated with output predictions that have higher entropy, which possibly corresponds to predictions that are more difficult to make" (Section 4.1). This is a correlational observation rather than a causal ablation, but it provides face validity that the router is learning something sensible about token difficulty.
Scaling to larger FLOP budgets (Figure 4): The isoFLOP analysis at 6e18, 2e19, and 1e20 FLOPs can be viewed as a robustness check: does the MoD advantage persist at larger scales? Figure 4 shows that the "down and to the right" shift is present at all three budgets, though the gap between the MoD isoFLOP curve and the baseline curve appears to narrow slightly at 1e20 FLOPs. The paper does not comment on this narrowing or whether it might continue at even larger scales (e.g., 1e21 FLOPs). This is a potentially important limitation — if the advantage continues to diminish with scale, MoD might provide diminishing returns in the regime of frontier models trained on 1e22+ FLOPs.
Depth vs. width scaling: The paper's claim that "it is better to add depth than to add width" is stated as an empirical finding without supporting ablation data. An experiment varying model shape (deep-narrow vs. shallow-wide) while holding total parameters and capacity constant would strengthen this claim. As presented, it may be a configuration choice rather than a robustly tested finding.

Critical Assessment

Claim 1: "MoD transformers learn to dynamically allocate compute and match baseline performance for equivalent FLOPS and wall-clock times to train, but require a fraction of the FLOPs per forward pass"

What the evidence shows: This claim is supported with qualifications. Figure 3 shows a specific MoD variant (model #3, 220M parameters) that matches the isoFLOP-optimal baseline (also 220M) in training loss while being ~60% faster per step. Figure 4 shows that across three FLOP budgets, there exist MoD variants that achieve lower training loss than the isoFLOP-optimal baseline while requiring fewer FLOPs per forward pass. The autoregressive evaluation in Figure 6 confirms that these advantages persist at inference time.