Mixture of Depths (MoD)

Mixture of Depths (MoD) is a transformer technique in which a learned router decides, block by block, which tokens are processed and which skip the block through its residual connection, so different tokens receive different amounts of computation. David Raposo, Adam Santoro, and colleagues at Google DeepMind introduced it in April 2024 in Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models. A fixed capacity k sets how many tokens each routed block processes, so the compute budget and tensor shapes are known in advance even though the choice of tokens depends on the input. MoD large language models (LLMs) matched baseline quality at equal training FLOPs and wall-clock time while using a fraction of the FLOPs per forward pass, and stepped upwards of 50% faster during sampling. At a fixed training budget, the best MoD model was larger than the best baseline and reached a lower loss.

How MoD Routing Works#

Before a routed block, a linear router gives every token a scalar score. The k highest-scoring tokens go through the block's self-attention and MLP, and their outputs are multiplied by their router scores and added back to the residual stream. Multiplying by the score puts the router on the gradient path, so backpropagation can teach it which tokens benefit. The remaining tokens pass through unchanged and are left out of that block's attention, so later tokens cannot attend to them there and the attention cost shrinks with the square of the capacity: half the tokens cost a quarter of the query-key products. The best setting in the paper routed every other block at a capacity of 12.5% of the sequence. A minimal version of one routed block, here with only the MLP, in PyTorch:

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

It prints Block ran on 16 of 128 tokens; 112 skipped it unchanged: the block's MLP did one-eighth of its usual work, and the output tensor keeps the full sequence shape for the next block.

Routing During Generation#

Top-k selection is non-causal: whether a token ranks in the top k of its sequence depends on the scores of tokens that come after it, which do not exist yet when a model generates text one token at a time. The authors made routing causal in two ways: an auxiliary binary cross-entropy loss that teaches the router's output to predict top-k membership, which shifted the language modeling loss by only 0.2-0.3%, or a small separate MLP predictor on stop-gradient inputs, which quickly reached 97-99% accuracy. Either way, autoregressive sampling lost little quality. Trained routers also show a pattern: some tokens engage every routed block along the model's depth, while others route around blocks whenever they can.

MoD vs. Related Techniques#

• Mixture of experts (MoE): Routes every token to one or a few expert MLPs within a layer, adding parameters while keeping per-token compute fixed. The paper also combined the two as MoDE, in which a token can be routed to an expert or to a no-op, and the gains compounded.
• Early exit: A token stops after some layer and skips all the rest. MoD can skip one block and still use later ones.
• Token merging (ToMe): Merges similar tokens in a vision transformer so later layers process fewer of them, with no learned router.
• Model pruning: Removes weights permanently, so every input gets the same smaller network.
• Model routing: Sends a whole request to one of several models, the same idea at the level of complete models.

Most networks fix their depth when they are built. Ultralytics YOLO26 sets it per model scale: the depth multiplier in its model YAML is 0.50 for the n, s, and m scales and 1.00 for l and x, so a C3k2 stage listed with 2 repeats runs 1 block in YOLO26n and 2 in YOLO26x, and every image passes through all of them.