Mixture of Depths (MoD)

Mixture of Depths (MoD) is a transformer architecture technique that lets a model decide, layer by layer, which tokens are worth spending full compute on and which can skip a layer’s attention and MLP block altogether. Introduced by Google DeepMind researchers David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro in the 2024 paper “Mixture-of-Depths: Dynamically allocating compute in transformer-based language models” (arXiv:2404.02258), it addresses a specific inefficiency in standard transformers: every token gets identical compute at every layer, whether that token is a load-bearing content word or a trivial punctuation mark or repeated filler. MoD gives the model a learned router that allocates FLOPs unevenly across the sequence at each layer, while keeping the total compute budget fixed and known ahead of time, which is the detail that separates it from other conditional-computation ideas including mixture-of-experts.

The Problem: Uniform Compute, Non-Uniform Difficulty

A standard transformer runs every token in a sequence through the same stack of attention and MLP blocks, layer after layer, at identical cost per token. But not every token needs that much processing. Predicting the token after “the” in “the cat sat on the ___” is a much easier problem than predicting the next token after a rare technical term in a dense paragraph, yet a vanilla transformer spends the same FLOPs on both. Prior attempts to exploit this unevenness, generally grouped under “early exiting” or conditional computation, let individual tokens exit the network early once a confidence threshold is hit. MoD takes a different, more structured approach: instead of a dynamic per-token decision that can vary unpredictably at runtime, it fixes in advance exactly how many tokens per layer will get full processing, and lets the router decide only which tokens fill those fixed slots.

How MoD Routing Works

Each MoD-enabled transformer layer wraps its usual attention-plus-MLP block with a router and a capacity limit:

1. A capacity is set per layer, as a fraction of the sequence length. The paper’s strongest results use small capacities, as low as 12.5% of tokens per layer, meaning seven out of every eight tokens skip that layer’s full block.
2. A learned router scores every token. A simple linear layer produces one scalar routing weight per token position, indicating how much that token “wants” to go through the full block at this layer.
3. Top-k selection picks the winners. The router sorts tokens by that scalar weight and keeps only the top capacity of them; because capacity is a fixed number decided in advance rather than a threshold that could pass a variable number of tokens, the computation graph has static, predictable tensor shapes.
4. Selected tokens go through attention and the MLP; the rest skip via the residual stream. Tokens that don’t make the cut bypass the block’s computation entirely and pass straight through to the next layer unchanged, exactly the way a residual connection normally carries information around a sublayer.
5. The router’s output gates the result. The routing weight (passed through a sigmoid) multiplies the processed token’s output before it’s written back, so the router receives a gradient signal and learns to route well over training.

Because attention is computed only over the tokens the router selected for a given layer, MoD also shrinks the attention computation itself at that layer: fewer tokens in the routed set means a smaller query/key/value matrix for that layer’s self-attention, not just a skipped MLP.

# Scenario: inspecting which tokens got routed through a given layer at inference time,
# useful for debugging whether the router is behaving sensibly on real input
def routed_positions(router, x, capacity_fraction):
    seq_len = x.shape[1]
    capacity = max(1, int(seq_len * capacity_fraction))
    logits = router(x).squeeze(-1)
    _, idx = logits.topk(capacity, dim=-1)
    return idx.sort(dim=-1).values  # positions that get full attention+MLP this layer

MoD vs. Mixture-of-Experts

MoD and Mixture-of-Experts are both “routing” techniques and both are frequently, and easily, confused for variants of the same idea. They route completely different things:

• MoE routes which expert processes a token. Every token still gets processed by some MLP block (typically one or two of many available experts), so total compute per token is roughly constant, but which parameters do the work varies token to token. MoE’s routing problem is about specialization: which expert is best suited to this token.
• MoD routes whether a token gets processed at all, at this layer. Every routed token goes through the exact same single block, but a fixed fraction of tokens skip the block entirely. MoD’s routing problem is about allocation: is this token worth spending compute on right now.

That difference has a direct consequence for system-level predictability. Token-choice MoE routing can be dynamic and imbalanced in practice: some experts get overloaded with more tokens than others in a given batch, which is exactly the load-imbalance problem that MoE architectures need auxiliary losses or capacity limits to manage. MoD sidesteps this by construction: since capacity is a fixed number of tokens per layer chosen ahead of time, the computation graph has static, known tensor shapes regardless of input, which is a large part of why the original paper describes MoD transformers as achieving comparable training time to vanilla transformers despite the added routing logic. The two techniques are also complementary rather than competing: a model can combine MoD’s “should this token compute at this layer” decision with MoE’s “which expert should process it” decision in the same architecture, and follow-up work (including “MoDification: Mixture of Depths Made Easy”) has explored exactly this kind of combination and retrofit onto existing pretrained models.

graph TB
    classDef default fill:#ffffff,stroke:#4338CA,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef data    fill:#EEF0F7,stroke:#0D9488,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef process fill:#F7F8FC,stroke:#6366F1,stroke-width:2px,color:#0F172A,rx:8px,ry:8px;
    classDef output  fill:#4338CA,stroke:#4338CA,stroke-width:2px,color:#ffffff,rx:8px,ry:8px;

    subgraph MOD["Mixture of Depths: routes WHETHER"]
        T1([Token sequence]):::data --> R1[Router scores<br/>every token]:::process --> K1[Top-k tokens selected<br/>by fixed capacity]:::process
        K1 --> B1[Selected: full attention + MLP]:::output
        K1 --> S1[Unselected: skip via residual]:::output
    end

    subgraph MOE["Mixture of Experts: routes WHICH"]
        T2([Token sequence]):::data --> R2[Router scores<br/>every expert, per token]:::process --> K2[Top-k experts selected<br/>per token]:::process
        K2 --> B2[Every token processed by<br/>1-2 of N expert MLPs]:::output
    end

The Capacity Fraction: The Real Trade-off

The one number that determines MoD’s whole compute-versus-quality trade-off is the capacity fraction: what proportion of tokens per layer get the full block. Set it to 100% and MoD degenerates into a normal transformer, every token processed at every layer, no compute saved, no quality lost from skipping. Push it down toward the paper’s more aggressive settings (12.5%) and the model skips the vast majority of token-layer computations, which is where the FLOPs savings come from, but pushes more of the sequence’s information-carrying burden onto the residual stream and the tokens that do get selected. The original paper reports that a well-trained MoD model can match an isoFLOP-optimal vanilla transformer’s training loss while using substantially fewer FLOPs per forward pass (upward of 50% fewer in some configurations), or alternatively outperform a vanilla transformer by a small margin (up to roughly 1.5% on the training objective) under an equal compute budget, because the saved FLOPs at skipped layers can be reinvested elsewhere (larger models, more training steps) rather than simply discarded. The catch is that capacity is usually fixed per layer at design time as a hyperparameter, not adjusted per input at inference, so choosing it is a training-time architectural decision rather than something adjusted per query.

Where MoD Sits Among Compute-Scaling Ideas

MoD is one of several 2024-2026 ideas about spending compute more deliberately rather than uniformly, and it’s useful to place it against the others rather than treat it in isolation. Test-time compute scaling spends more compute at inference by having a model think longer or sample more before answering; MoD instead spends less compute per forward pass by skipping unnecessary work inside a single pass. They are not competing techniques: a model could use MoD to make each individual forward pass cheaper and separately use test-time compute scaling (extra reasoning steps, best-of-n sampling) on top of those cheaper passes. Within architecture-level efficiency techniques specifically, MoD is closer in spirit to attention-efficiency methods like grouped-query attention and sliding window attention, all of which reduce compute or memory without changing what the model is fundamentally capable of representing, than it is to capacity-expanding techniques like MoE, which add parameters and specialization rather than remove computation.

What’s New (2025-2026)

• Retrofit and hybridization work. Follow-up research such as “MoDification: Mixture of Depths Made Easy” has focused on applying MoD-style routing to already-pretrained models rather than requiring training a new model from scratch with MoD baked in from the start, lowering the barrier to adopting the technique on existing model families.
• Combination with MoE architectures. Because MoD (whether a token computes at this layer) and MoE (which expert handles it) answer different questions, 2025-2026 research has explored layers that do both: route a subset of tokens for full processing, then further route those selected tokens to a specific expert, stacking the compute savings of both techniques rather than picking one.
• Continued interest as inference costs dominate deployment economics. As serving costs, not just training costs, have become the binding constraint for many teams running LLMs at scale, MoD’s promise of a smaller, predictable per-token compute footprint at inference (not just training) has kept it an active area of applied research alongside other inference-efficiency techniques like flash attention.

MoD is not a replacement for MoE, attention-efficiency tricks, or test-time compute scaling; it is a distinct lever, deciding how many tokens at each layer get full computation at all, that can be combined with most of the others. Its main practical appeal is the predictability that top-k routing buys: unlike token-choice MoE’s potential for runtime load imbalance, an MoD model’s compute cost per forward pass is fixed and known the moment the capacity fraction is chosen.

import torch
import torch.nn as nn

class MoDBlock(nn.Module):
    # A transformer block where only `capacity` tokens per sequence
    # get the full attention+MLP treatment at this layer
    def __init__(self, dim, attn_mlp_block, capacity_fraction=0.125):
        super().__init__()
        self.router = nn.Linear(dim, 1)          # scalar routing weight per token
        self.block = attn_mlp_block               # the normal attention + MLP sublayer
        self.capacity_fraction = capacity_fraction

    def forward(self, x):
        b, seq_len, dim = x.shape
        capacity = max(1, int(seq_len * self.capacity_fraction))

        router_logits = self.router(x).squeeze(-1)      # (b, seq_len)
        weights, idx = torch.topk(router_logits, capacity, dim=-1)

        selected = torch.gather(x, 1, idx.unsqueeze(-1).expand(-1, -1, dim))
        processed = self.block(selected) * torch.sigmoid(weights).unsqueeze(-1)

        out = x.clone()
        out.scatter_(1, idx.unsqueeze(-1).expand(-1, -1, dim), processed)
        # tokens NOT in idx bypass the block entirely via this residual copy
        return out