Adaptive compute: spend effort only where it's needed

Not every input deserves the same compute

Everything so far made the model uniformly cheaper. But inputs are not uniform: classifying a blank wall is trivial, parsing a crowded street scene is hard, and yet a static network spends the same FLOPs on both. Adaptive (conditional) computation breaks that symmetry — it lets the network decide, per input or even per token, how much of itself to run. The promise is to cut average inference cost without lowering the ceiling for hard cases.

Early-exit networks: stop when you're sure

Early-exit networks attach lightweight classifier 'heads' to intermediate layers. As an input flows upward, each exit head produces a tentative prediction with a confidence score; if an early head is confident enough, inference stops there and skips all the remaining layers. Easy inputs leave early and cost a fraction of the full depth; hard inputs flow all the way to the top.

\hat{y}=g_\ell(h_\ell),\qquad \text{exit at layer }\ell \iff \max_{c}\,[\hat{y}]_c \ge \tau

Early exit: emit the prediction at the first head whose top-class confidence clears the threshold τ.

h = embed(x)
for layer, exit_head in zip(layers, heads):
    h = layer(h)
    p = exit_head(h)
    if confidence(p) > threshold:
        return p          # leave early, skip the rest
return p                  # used all layers

Early exit: return as soon as an intermediate head is confident enough.

The design tension is the confidence threshold: too low and you exit on wrong answers, too high and you never save anything. Calibrating the exit heads matters, and training them well — often by distilling the final layer's behavior into the early heads — is what makes the technique pay off in practice.

Mixture-of-depths: a compute budget per token

Early exits are coarse — the whole input leaves at once. Mixture-of-depths (MoD) is finer: at each layer, a small router selects only the top-k most useful tokens to be processed by that layer's attention and feed-forward block; the rest skip the layer entirely via the residual path. Over a sequence, the model spends its compute budget on the tokens that need it, layer by layer, while keeping a fixed and predictable total cost.

y_i=\begin{cases}x_i+f_\ell(x_i),& i\in\mathcal{S}_\ell\\ x_i,& \text{otherwise}\end{cases}\qquad \mathcal{S}_\ell=\text{top-}k\big(r_\ell(x_1),\dots,r_\ell(x_N)\big)

Mixture-of-depths: only the router's top-k tokens enter the block; the rest skip it via the residual path.

If this sounds like cousins of mixture-of-experts, it is — both are conditional routing. The difference: mixture-of-experts routes a token to one of several parallel experts (more parameters, same depth), while mixture-of-depths routes a token past a layer (same parameters, variable depth). They are complementary axes of dynamic compute, and modern designs sometimes use both.

Token merging: shrink the sequence, not the network

Transformer cost grows with sequence length — quadratically in attention. Token merging (ToMe) attacks that length directly: between layers, it finds pairs of tokens whose representations are highly similar and averages them into one, gradually shortening the sequence as it deepens. In a vision transformer, adjacent patches of sky or skin are nearly redundant, so merging them costs almost no accuracy while removing real work.

C_{\mathrm{attn}}=\mathcal{O}(N^2 d)\ \xrightarrow{\,N\to rN,\ r<1\,}\ \mathcal{O}(r^2 N^2 d)

Why token merging pays off: attention cost is quadratic in sequence length, so shrinking N to rN cuts it by r².

What makes token merging lovely is that it often needs no retraining — you can drop it into a pretrained model and trade a tunable amount of speed for accuracy at inference time. It is the sequence-axis sibling of pruning: instead of removing weights, you remove redundant positions.