A normal transformer forces every token through every layer, so the word "the" costs exactly as much compute as "photosynthesis." That's obviously wasteful, and Mixture-of-Depths (DeepMind, 2024) is the fix that finally made the waste addressable to me. It adds a tiny router to each block that scores every token and keeps only the top-k to run the full attention+MLP; the rest skip the block entirely, riding the residual connection for ~0 FLOPs. Depth gets allocated dynamically, per token, per layer — at the same total budget. I built an 8-layer, 12-token demo where every router score and FLOP count is computed live. Here's how it works.
The waste we're removing
A plain transformer runs L blocks over all T tokens — L·T block-passes, with filler words getting the same budget as information-dense ones.
const L = 8;
const tokens = embed(sentence); // T token vectors
// dense forward = L * T block-passes.
// "the" costs as much as "photosynthesis" <- the waste we remove.
The router is one tiny linear layer
Give each block a router: a single weight vector that turns each token's vector into one scalar — "how much does this token want this block's compute?" It runs on all tokens, but it's one dot product each, so it's well under 0.1% of a block.
function router(x, Wr){
return x.map(tok => dot(tok, Wr)); // one "compute demand" number per token
}
Top-k is a fixed budget — and that's the whole trick
Pick a capacity (say 50%) and keep only the top k = capacity·T tokens by score. Crucially k is fixed — the same count every pass — not a threshold. A threshold would process a different number of tokens each input, giving a dynamic compute graph that GPUs hate. Top-k keeps the graph static: known FLOPs, clean batching.
function topK(weights, capacity){
const k = Math.round(capacity * weights.length); // FIXED k = the budget
const order = weights.map((w,i)=>[w,i]).sort((a,b)=>b[0]-a[0]);
return new Set(order.slice(0,k).map(p=>p[1]));
}
Process the chosen; the rest ride the residual
This is the heart of it. Chosen tokens get the full attention+MLP; every other token is returned unchanged, riding the residual straight past the block for ~0 FLOPs. No token is dropped — it just isn't transformed here.
return x.map((tok,i) =>
keep.has(i)
? tok + sigmoid(w[i]) * attnMlp(tok, x) // CHOSEN: full block, scaled by router score
: tok // SKIPPED: the residual shortcut, ~0 FLOPs
);
Scaling the output by the router score (through a sigmoid) is what makes it trainable: a hard top-k has no gradient, but putting the score on the compute path lets the loss flow back into the router, so it learns on its own that content words deserve depth and "the" doesn't — no labels required.
Counting the savings
Each layer routes independently, so a token can be processed at layer 2, skip 3, run again at 4. The dense model does L·T block-passes; MoD does L·k plus the tiny router, so savings are almost exactly 1 − capacity.
const modCost = L * Math.round(CAPACITY * T) + L * T * routerFLOPs;
const saved = 1 - modCost / (L * T); // ~= 1 - CAPACITY -> ~50% fewer FLOPs
The paper reports matching a dense model with ~50% of the per-forward FLOPs. One catch at inference: top-k needs the whole sequence, which isn't causal — so MoD trains a small predictor that decides "process this token?" from the token alone, preserving the fixed budget while staying autoregressive and KV-cache-friendly.
MoD vs Mixture-of-Experts
Both "route" tokens, but they save different things. MoE routes every token to a different expert — it makes the model wider, adding parameters but spending the same compute per token. MoD routes only the top-k through the block — same parameters, fewer FLOPs. MoE trades width; MoD trades depth. They're orthogonal, so you can even do both — "MoDE" blocks route to experts and skip tokens.
Set the capacity and watch the stack route tokens, layer by layer:
https://dev48v.infy.uk/ai/days/day47-mixture-of-depths.html
Top comments (0)