Mixture-of-Depths

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.

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.

<!-- {{c_demo_title}} -->
<div class="controls">
  <label for="budget">{{lbl_budget}} <span id="budgetVal">50</span>%</label>
  <input type="range" id="budget" min="10" max="90" step="5" value="50" title="{{lbl_budget_tip}}">
  <button id="randomize" type="button" title="{{lbl_randomize_tip}}">{{btn_randomize}}</button>
  <button id="route" type="button">{{btn_route}}</button>
</div>
<div class="legend">
  <span class="leg-box enter"></span>{{lbl_enters_layer}}
  <span class="leg-box skip"></span>{{lbl_skips_layer}}
</div>
<div id="tokens" class="tokens" aria-label="{{aria_token_grid}}"></div>
<div id="status" class="status" role="status"></div>
<div id="bar-wrap" class="bar-wrap" aria-label="{{aria_score_chart}}">
  <div id="bar-chart" class="bar-chart"></div>
  <div id="threshold-line" class="threshold-line"></div>
</div>

/* {{c_reset}} */
* { box-sizing: border-box; margin: 0; }
body { font-family: system-ui, sans-serif; color: #222; padding: .5rem; }
/* {{c_controls}} */
.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; }
/* {{c_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; }
/* {{c_token_grid}} */
.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; }
/* {{c_status}} */
.status { font-size: .9rem; font-weight: 600; min-height: 1.3em; margin-bottom: .4rem; }
.status.ok  { color: #0a7d33; }
.status.info { color: #1d3557; }
/* {{c_barchart}} */
.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; }

// Code not found

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-$k$ 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.

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.

"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.

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.