Introduction

Every word in a sentence is different. "The" carries almost no surprise; "antidisestablishmentarianism" is packed with information. Yet inside a standard transformer, every token travels through exactly the same stack of layers, consuming the same amount of compute regardless of how much thinking it actually needs.

Mixture-of-Depths (MoD) breaks that rule. Published by Raposo et al. at Google DeepMind in 2024, MoD adds a lightweight router at each transformer layer. The router scores each token and decides: process it fully here, or skip this layer entirely and pass through unchanged. A budget controls how many tokens get processed — for example, only 12.5 % of tokens might enter the most expensive layers.

The result is a model that uses the same number of parameters as a standard transformer but can run significantly faster at inference, because whole slices of the compute graph simply never execute for easy tokens. The router learns its routing policy jointly with the rest of the model — no hand-engineering required.

This connects to a broader family of ideas: see Mixture-of-Experts for the related technique that routes tokens to different experts rather than skipping layers.

Try It

The demo below simulates a MoD router deciding which tokens enter a transformer layer and which skip it. Each token has a router score (higher = more informative). Set the budget (what fraction of tokens may enter the layer) and press Route tokens to see which ones pass through.

<!-- Mixture-of-Depths token router demo -->
<div class="controls">
  <label for="budget">Budget: <span id="budgetVal">50</span>%</label>
  <input type="range" id="budget" min="10" max="90" step="5" value="50" title="Fraction of tokens allowed to enter this layer">
  <button id="randomize" type="button" title="Pick a new sentence">New tokens</button>
  <button id="route" type="button">Route tokens</button>
</div>
<div class="legend">
  <span class="leg-box enter"></span>Enters layer
  <span class="leg-box skip"></span>Skips layer
</div>
<div id="tokens" class="tokens" aria-label="Token routing result grid"></div>
<div id="status" class="status" role="status"></div>
<div id="bar-wrap" class="bar-wrap" aria-label="Router score bar chart">
  <div id="bar-chart" class="bar-chart"></div>
  <div id="threshold-line" class="threshold-line"></div>
</div>

/* base reset */
* { box-sizing: border-box; margin: 0; }
body { font-family: system-ui, sans-serif; color: #222; padding: .5rem; }
/* controls bar */
.controls { display: flex; flex-wrap: wrap; align-items: center; gap: .5rem .8rem; margin-bottom: .5rem; }
label { font-size: .85rem; color: #444; }
input[type=range] { width: 140px; accent-color: #1d3557; }
button { font: 600 13px system-ui; padding: .35rem .8rem; border: 1px solid #1d3557;
         border-radius: 7px; cursor: pointer; }
#route { background: #1d3557; color: #fff; }
#randomize { background: #fff; color: #1d3557; }
/* color legend */
.legend { display: flex; gap: 1rem; font-size: .8rem; color: #555; margin-bottom: .4rem; align-items: center; }
.leg-box { display: inline-block; width: 14px; height: 14px; border-radius: 4px; }
.leg-box.enter { background: #2a9d8f; }
.leg-box.skip  { background: #adb5bd; }
/* token chips */
.tokens { display: flex; flex-wrap: wrap; gap: 5px; margin-bottom: .6rem; }
.token { padding: .25rem .5rem; border-radius: 6px; font-size: .82rem; font-weight: 600;
         transition: background .25s, color .25s; border: 1px solid transparent; }
.token.enter { background: #2a9d8f; color: #fff; border-color: #21867a; }
.token.skip  { background: #dee2e6; color: #555; border-color: #ced4da; }
.token.neutral { background: #e9ecef; color: #333; border-color: #dee2e6; }
/* status line */
.status { font-size: .9rem; font-weight: 600; min-height: 1.3em; margin-bottom: .4rem; }
.status.ok  { color: #0a7d33; }
.status.info { color: #1d3557; }
/* score bar chart */
.bar-wrap { position: relative; margin-top: .3rem; }
.bar-chart { display: flex; align-items: flex-end; gap: 3px; height: 80px; }
.bar { flex: 1; border-radius: 3px 3px 0 0; transition: height .3s, background .3s; min-height: 4px; }
.bar.enter { background: #2a9d8f; }
.bar.skip  { background: #adb5bd; }
.threshold-line { position: absolute; left: 0; right: 0; height: 2px;
                  background: #e63946; pointer-events: none; transition: top .3s; }

const fmt = (s, v) => s.replace(/%(\w+)%/g, (_, k) => String(v[k] ?? ''));
// Mixture-of-Depths: a router at each transformer layer scores each token.
// Only the top-k scoring tokens (by budget) enter the layer; the rest skip via residual.

const SENTENCES = [
  ["The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"],
  ["Attention", "is", "all", "you", "need", "for", "sequence", "modeling"],
  ["Not", "every", "token", "requires", "deep", "computation"],
  ["Mixture", "of", "depths", "routes", "tokens", "to", "skip", "layers"]
];

// Active token list from a sample sentence
let tokens = [...SENTENCES[0]];
// Router scores (0-100) assigned to each token
let scores = [];

function randomScores(n) {
  // Assign a random router score to each token (simulates a learned router)
  return Array.from({length: n}, () => Math.round(Math.random() * 100));
}

function init() {
  scores = randomScores(tokens.length);
  render(null);
}

function render(routed) {
  // Draw token chips colored by routing decision
  const container = document.getElementById('tokens');
  container.innerHTML = '';
  tokens.forEach((word, i) => {
    const el = document.createElement('span');
    el.className = 'token ' + (routed === null ? 'neutral' : routed[i] ? 'enter' : 'skip');
    el.textContent = word;
    el.title = fmt('score: %v%', { v: scores[i] });
    container.appendChild(el);
  });
  renderBars(routed);
}

function renderBars(routed) {
  const chart = document.getElementById('bar-chart');
  chart.innerHTML = '';
  const maxScore = Math.max(...scores, 1);
  scores.forEach((s, i) => {
    const bar = document.createElement('div');
    bar.className = 'bar ' + (routed === null ? 'skip' : routed[i] ? 'enter' : 'skip');
    bar.style.height = (s / maxScore * 100) + '%';
    bar.title = fmt('%v% — score: %s%', { v: tokens[i], s: s });
    chart.appendChild(bar);
  });
  // Draw the score threshold line on the bar chart
  const budget = parseInt(document.getElementById('budget').value);
  const k = Math.max(1, Math.round(tokens.length * budget / 100));
  const sorted = [...scores].sort((a, b) => b - a);
  const threshold = sorted[k - 1] ?? 0;
  const thresholdFrac = threshold / Math.max(...scores, 1);
  const wrap = document.getElementById('bar-wrap');
  document.getElementById('threshold-line').style.top =
    (wrap.clientHeight - thresholdFrac * 80) + 'px';
}

function doRoute() {
  // Top-k selection: route the highest-scoring tokens into the layer
  const budget = parseInt(document.getElementById('budget').value);
  const k = Math.max(1, Math.round(tokens.length * budget / 100));
  // Threshold is the k-th highest score
  const sorted = [...scores].sort((a, b) => b - a);
  const threshold = sorted[k - 1] ?? 0;
  const routed = scores.map(s => s >= threshold);
  const entered = routed.filter(Boolean).length;
  render(routed);
  document.getElementById('status').textContent =
    fmt('%entered% / %length% tokens enter the layer (%budget%% budget)', { entered: entered, length: tokens.length, budget: budget });
  document.getElementById('status').className = 'status ok';
  console.log('MoD routing:', entered, '/', tokens.length, '| threshold score', threshold);
}

document.getElementById('route').addEventListener('click', doRoute);
document.getElementById('randomize').addEventListener('click', () => {
  const idx = Math.floor(Math.random() * SENTENCES.length);
  tokens = [...SENTENCES[idx]];
  scores = randomScores(tokens.length);
  document.getElementById('status').textContent = '';
  document.getElementById('status').className = 'status';
  render(null);
});
document.getElementById('budget').addEventListener('input', e => {
  document.getElementById('budgetVal').textContent = e.target.value;
  render(null);
});

init();
Result
Console

Tokens above the score threshold enter the layer and get updated representations. Tokens below the threshold skip it — their embedding passes through unchanged via a residual connection. The key insight is that the threshold is not hand-picked: a top-kk selection of the highest-scoring tokens is used, and the router learns to give high scores to tokens that genuinely need this layer's computation.

The Real Complexity

Skipping layers sounds simple, but several subtle issues make it non-trivial:

  • Causal masking at inference: autoregressive generation visits tokens one by one. At each step the model cannot know future tokens, so it cannot re-run routing decisions it already committed to. MoD addresses this by caching routing decisions alongside key–value pairs.
  • The routing must be learned, not heuristic: early attempts to use fixed rules (skip every other layer, or skip short words) hurt accuracy. The router is trained with the rest of the model using a simple auxiliary loss so that the routing decisions become meaningful.
  • isoFLOP comparisons: naively comparing a MoD model to a dense baseline is unfair. Raposo et al. carefully match total floating-point operations across configurations, showing MoD models reach the same loss as their dense equivalents while being faster.
  • Interaction with Mixture-of-Experts: MoD and MoE are orthogonal. A token that skips a layer avoids both the attention and the feed-forward expert computation in that layer. Combining them gives further savings.

The method is solved in the sense that the 2024 paper demonstrates clear wins. Open questions remain around optimal budget schedules across layers, learned budgets, and very long contexts.

Where It Matters

"Spend more compute where it's needed, less where it isn't" is an old idea in engineering. MoD brings it to the core of transformer inference:

  • Cheaper API calls: large language models are expensive to run. A MoD layer reduces the FLOPs for most tokens, directly cutting cost per token at scale.
  • On-device inference: mobile and edge devices have hard power budgets. Routing easy tokens around layers keeps quality high while fitting in tight constraints.
  • Long-context efficiency: in a 128 k-token document, the vast majority of tokens are routine. MoD lets the model focus depth on the rare, surprising tokens that actually move understanding forward.
  • Adaptive training: the same routing logic can be applied during training to reduce wall-clock time, not just inference cost.
  • Research into token importance: the router's learned scores provide an interpretability signal — high-scoring tokens are the ones the model "thinks" need the most processing.

The core tension MoD resolves — uniform compute vs. input-adaptive compute — will remain central as models grow and deployment costs become the bottleneck.

Conclusion

Mixture-of-Depths makes a deceptively simple bet: not every token needs every layer, and a lightweight router can learn to tell the difference. The result is a transformer that allocates its compute budget dynamically — spending deeply on hard tokens, sailing through easy ones.

The idea sits alongside Mixture-of-Experts as part of a broader shift toward conditional computation in neural networks: the model itself decides how much work to do, rather than the architecture imposing a fixed answer for every input. As models grow and inference costs bite harder, that flexibility is increasingly what separates efficient AI from expensive AI.

Share this article

Pick a channel — or use your device's native share sheet.

Comments

Loading comments...

https://www.kipuhub.com/en/article/mixture-of-depths/Content licensed under CC BY-NC 4.0.