Mixture of Depths

9 min read

Learn

Quick Reading

Estimated 9 mins

Prereq

Intermediate

Basic ML concepts helpful

Interactive

Static Playbook

Static guide & reference tables

The idea: not every token deserves full depth

A standard transformer runs every token through every layer, regardless of how 'hard' that token is to process. Predicting the next token after "The capital of France is" is nearly trivial by the time the model sees "Paris" context; predicting the next word in a dense technical proof may genuinely need every layer's worth of computation. Mixture of Depths (MoD) applies the same intuition that Mixture of Experts applies across width — route dynamically — but applies it across *depth*: at each layer, a router decides whether a given token needs to pass through that layer's computation at all, or can skip straight to a residual connection.

How the routing mechanism works

At each MoD-enabled layer, a lightweight router scores every token in the sequence. Only the top-k tokens (a fixed capacity, e.g. 12.5% or 25% of the sequence) are selected to go through the full block — attention and feed-forward computation. The remaining tokens skip the block entirely via the residual path, meaning their representation passes through unchanged at that layer. Because the capacity is fixed in advance (not data-dependent), MoD preserves predictable, static computation graphs — a major practical advantage over techniques with variable per-token compute, which are hard to batch efficiently on GPUs.

mixture_of_depths_sketch.py

python

import torch

def mod_layer(x, router, block, capacity_fraction=0.125):
    """
    x: (seq_len, d_model) token representations
    router: linear layer -> scalar score per token
    block: the full transformer block (attention + FFN)
    """
    seq_len = x.shape[0]
    scores = router(x).squeeze(-1)             # (seq_len,)
    k = max(1, int(seq_len * capacity_fraction))
    topk_vals, topk_idx = torch.topk(scores, k)

    out = x.clone()                              # default: skip (residual passthrough)
    selected = x[topk_idx]                       # tokens chosen for full compute
    processed = block(selected)                  # run attention + FFN only on these
    # weight by router score so gradients flow back into the router
    out[topk_idx] = processed * torch.sigmoid(topk_vals).unsqueeze(-1)
    return out

Compute savings without shrinking the model

Like MoE, MoD decouples total model capacity from per-token compute cost — but along the depth axis instead of the width axis. A model can keep its full number of layers (and therefore its full representational depth when needed) while spending far less FLOPs on average, because most tokens at most layers take the cheap residual skip. This is attractive combined with MoE: experts widen the model, MoD lets tokens skip whole blocks, and together they decouple both capacity dimensions from average inference cost.

Training considerations

The top-k selection is non-differentiable in its hard form, so MoD training typically uses the router's score to weight the output of processed tokens (as in the sketch above) so gradients can flow into the router even though the selection itself is discrete. Like MoE routing, this can be unstable early in training — the router has to learn a genuinely useful notion of 'this token needs this layer' rather than collapsing to an arbitrary fixed pattern. Curriculum tricks (e.g. starting with a higher capacity fraction and annealing it down) and auxiliary losses to encourage sensible usage are common mitigations, echoing MoE's load-balancing problem.

Mixture of Experts vs Mixture of Depths

Property	Mixture of Experts	Mixture of Depths
Axis of sparsity	Width (which expert FFN)	Depth (which layers run)
What's skipped	Non-selected experts at a layer	Entire block at a layer, via residual
Routing granularity	Per-token, per-layer, among experts	Per-token, per-layer, binary run/skip
Main failure mode	Expert collapse / load imbalance	Router picking a degenerate, unhelpful subset

Where this fits in the bigger picture

MoD is part of a broader family of conditional-computation techniques that ask: does every token really need every unit of compute the architecture provides? Early-exit networks, adaptive computation time, and MoD all target this question from slightly different angles — early exit stops a token's forward pass entirely once confident enough, while MoD lets a token selectively skip individual blocks and still potentially run through later ones. The common thread is treating 'depth' as a resource to be allocated per-token rather than spent uniformly.

What's Next?

Routing decisions — whether across experts or across depth — are only as good as the gating mechanism making them, so the natural next question is how these routing and gating networks themselves are designed and trained to make good decisions reliably.

I build these systems professionally.

Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.

Work With Me →