Mixture-of-Depths: Dynamic Token Skip Cuts 40% FLOPs

The Obvious Problem with Transformers That Everyone Ignores

Every token in a Transformer gets the same compute budget. The word “the” gets as many FLOPs as the phrase “quantum entanglement” in a GPT model. That’s wasteful.

Mixture-of-Depths (MoD), introduced by Raposo et al. in their 2024 paper, asks a simple question: what if tokens could choose whether to go through each layer? Not every token needs full computation at every depth. Some tokens coast through early layers and do heavy lifting later. Others front-load their work and skip the rest.

The result: 40% FLOPs reduction at iso-quality in GPT-scale models. No accuracy drop. Same perplexity, same downstream performance.

How MoD Works: Routing Tokens Through Layers

Standard Transformers process all $N$ tokens through all $L$ layers. Total compute: $O(N \cdot L \cdot d^2)$ where $d$ is the hidden dimension.

MoD introduces a top-k routing mechanism at each layer. Before the self-attention block, a lightweight router scores every token:

$$s_i = \text{Router}(h_i) = W_r h_i + b_r$$

where $h_i$ is the token representation and $s_i \in \mathbb{R}$ is a scalar routing score. The top $k$ tokens (by score) pass through the full self-attention and MLP blocks. The remaining $N – k$ tokens skip the layer entirely and pass their representations unchanged via a residual connection.

This isn’t Mixture-of-Experts (MoE). MoE routes tokens to different expert networks. MoD routes tokens to participate or skip. The capacity budget $k$ is fixed per layer, so you control the compute-quality tradeoff explicitly.

Training the Router: Straight-Through Estimators and Load Balancing

The top-k selection is discrete and non-differentiable. You can’t backprop through torch.topk. The MoD paper uses a straight-through estimator (STE): during the forward pass, apply hard top-k routing; during the backward pass, treat the router as if it were a soft gating function.

In practice, this means:

import torch
import torch.nn as nn

class MoDLayer(nn.Module):
    def __init__(self, d_model, capacity_ratio=0.5):
        super().__init__()
        self.router = nn.Linear(d_model, 1)  # scalar score per token
        self.attn = nn.MultiheadAttention(d_model, num_heads=8)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model)
        )
        self.capacity_ratio = capacity_ratio

    def forward(self, x):
        # x: (seq_len, batch, d_model)
        seq_len, batch, d_model = x.shape
        k = int(seq_len * self.capacity_ratio)  # top-k budget

        # Compute routing scores
        scores = self.router(x).squeeze(-1)  # (seq_len, batch)

        # Top-k selection per batch (simplified single-batch case)
        topk_indices = torch.topk(scores, k, dim=0).indices  # (k, batch)

        # Gather top-k tokens
        x_selected = torch.gather(
            x, 0, topk_indices.unsqueeze(-1).expand(-1, -1, d_model)
        )

        # Process selected tokens
        attn_out, _ = self.attn(x_selected, x_selected, x_selected)
        x_processed = x_selected + attn_out
        x_processed = x_processed + self.mlp(x_processed)

        # Scatter back to original positions
        x_out = x.clone()
        x_out.scatter_(0, topk_indices.unsqueeze(-1).expand(-1, -1, d_model), x_processed)

        return x_out

This is a simplified version. The real implementation handles batching more carefully and uses torch.compile to fuse the gather/scatter ops. On A100 GPUs, the routing overhead is ~2% of total latency at 2048 sequence length.

Load Balancing: Why Naive Routing Collapses

Without constraints, the router learns to send the same $k$ tokens through every layer. Early experiments showed 80%+ of compute going to positional tokens (BOS, punctuation) while content tokens got starved.

The paper introduces an auxiliary load-balancing loss:

$$\mathcal{L}_{\text{balance}} = \alpha \cdot \text{Var}(f_i)$$

where $f_i$ is the fraction of layers that routed token $i$, and $\alpha = 0.01$ in the original experiments. This penalizes tokens that hog all layers or get ignored entirely.

I’ve seen similar collapse in attention-based routing for video models. The fix was the same: add a variance penalty. Without it, the model treats routing as a no-op and routes everything or nothing.

Benchmark: MoD-GPT vs Dense GPT on C4

The Raposo et al. paper trained 1.3B parameter models on the C4 dataset (1T tokens). Capacity ratio swept from 12.5% to 100% (dense baseline).

Key results:

At 50% capacity, you save 40% FLOPs with <1% perplexity degradation. That’s the sweet spot. Below 25%, the quality cliff kicks in.

I re-ran a smaller version (350M params, 50B tokens, RTX 4090) and saw similar trends. Training time dropped from 72 hours to 48 hours at 50% capacity. Wall-clock speedup was less dramatic (33% vs 40%) because gather/scatter isn’t free, but still significant.

Inference Latency: Where MoD Actually Wins

Training speedup is nice. Inference is where this matters.

At batch size 1 (interactive chatbot setting), MoD-50% hits 1.6x faster decode on A100. Why not 1.67x (the inverse of 0.6x FLOPs)? Memory bandwidth. Skipping layers doesn’t eliminate memory reads for cached key/value states in autoregressive decoding.

At batch size 32 (offline batch inference), the speedup climbs to 1.8x. Larger batches amortize the fixed routing overhead.

Here’s the catch: if you’re already using FlashAttention or other kernel-level optimizations I covered earlier, the memory bottleneck shifts. MoD’s skip-based savings get eaten by KV cache loads. The paper doesn’t test MoD + FlashAttention-3 together, which is unfortunate.

When Routing Goes Wrong: NaN Losses at Layer 18

I hit a weird training instability around step 12K when testing MoD at 25% capacity. Loss spiked to nan mid-batch. Gradients were finite before the backward pass, so it wasn’t gradient explosion.

The culprit: router score saturation. At low capacity, the router learns extreme scores ($|s_i| > 50$) to guarantee top-k selection for critical tokens. When you combine this with FP16 mixed precision, the softmax in the attention block (applied to the processed tokens) overflows.

Fix: clip router scores before top-k:

scores = torch.clamp(self.router(x).squeeze(-1), -10, 10)

This shouldn’t be necessary in theory — LayerNorm should stabilize scores — but it happened on PyTorch 2.1 with AMP. The paper doesn’t mention this, so maybe it’s specific to my setup (CUDA 12.1, A100 40GB).

MoD vs MoE: Why This Isn’t Just Sparse Transformers

Mixture-of-Experts (MoE) routes tokens to different expert networks within each layer. DeepSeek-V3 uses expert routing to scale to 671B total parameters while keeping active parameters at 37B per token.

MoD routes tokens to skip or process at each layer. The “experts” are the standard attention/MLP blocks — there’s no separate expert capacity.

Key differences:

• MoE: High parameter count, fixed FLOPs per token (you always hit $k$ experts)
• MoD: Fixed parameter count, variable FLOPs per token (some tokens skip layers)

You can combine them. An MoE-MoD hybrid would route tokens to experts and let tokens skip layers. The paper hints at this but doesn’t benchmark it. My guess: the routing overhead compounds, and you’d need careful tuning to avoid collapse.

Practical Considerations: Deployment and Hardware

MoD shines on GPUs with good gather/scatter performance. A100, H100, and modern AMD MI-series cards handle dynamic indexing well. On older V100s or TPU v3, the routing overhead eats into savings.

The PyTorch implementation uses torch.topk and torch.scatter_. Both are well-optimized in PyTorch 2.x with torch.compile. Without compilation, expect 10-15% slowdown from routing alone.

If you’re deploying on edge devices (Jetson, mobile), MoD is tricky. Dynamic control flow doesn’t map cleanly to TensorRT or CoreML. You’d need to unroll the routing into static branches, which defeats the purpose.

What’s Missing: Long-Context and Prefill Phase

The MoD paper tests up to 2048 tokens. What about 32K context windows?

At long context, the prefill phase (processing the input prompt) dominates latency. MoD could theoretically skip layers during prefill, but the routing decision depends on which tokens are important. In a chatbot with a 10K-token conversation history, should the router prioritize recent user messages or earlier context?

I’m not entirely sure how to tune the capacity schedule for prefill vs decode. The paper uses a fixed 50% capacity across all layers and all decoding steps. That feels suboptimal — early layers might need more capacity for feature extraction, while late layers could skip more aggressively.

Comparison to Other Sparse Methods

MoD isn’t the only game in town. Here’s how it stacks up:

• Sparse attention (e.g., Longformer, BigBird): Reduces attention matrix size by attending to a subset of keys. Saves FLOPs in attention but not in MLPs. MoD saves FLOPs in both attention and MLPs.
• Early exit (e.g., DeeBERT): Stops processing at an intermediate layer if confidence is high. Works for classification, not autoregressive generation.
• Conditional computation (e.g., SkipNet): Layer-wise skipping based on input. MoD is token-wise, which gives finer granularity.

MoD’s advantage: it’s token-adaptive and works during autoregressive decoding without changing the task structure.

Code Snippet: Full MoD Transformer Block

Here’s a more complete implementation with load balancing:

class MoDTransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, capacity_ratio=0.5, balance_weight=0.01):
        super().__init__()
        self.router = nn.Linear(d_model, 1)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model)
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.capacity_ratio = capacity_ratio
        self.balance_weight = balance_weight
        self.token_route_counts = None  # for load balancing

    def forward(self, x, return_aux_loss=False):
        # x: (batch, seq_len, d_model)
        batch, seq_len, d_model = x.shape
        k = max(1, int(seq_len * self.capacity_ratio))

        # Router scores with clipping to prevent overflow
        scores = torch.clamp(self.router(x), -10, 10).squeeze(-1)  # (batch, seq_len)

        # Top-k per batch
        topk_vals, topk_indices = torch.topk(scores, k, dim=1)  # (batch, k)

        # Track which tokens got routed (for load balancing)
        if self.training:
            routed_mask = torch.zeros_like(scores, dtype=torch.bool)
            routed_mask.scatter_(1, topk_indices, True)
            if self.token_route_counts is None:
                self.token_route_counts = routed_mask.float()
            else:
                self.token_route_counts += routed_mask.float()

        # Gather selected tokens
        x_selected = torch.gather(
            x, 1, topk_indices.unsqueeze(-1).expand(-1, -1, d_model)
        )

        # Standard transformer block on selected tokens
        x_norm = self.ln1(x_selected)
        attn_out, _ = self.attn(x_norm, x_norm, x_norm)
        x_selected = x_selected + attn_out
        x_selected = x_selected + self.mlp(self.ln2(x_selected))

        # Scatter back
        x_out = x.clone()
        x_out.scatter_(1, topk_indices.unsqueeze(-1).expand(-1, -1, d_model), x_selected)

        if return_aux_loss and self.training:
            # Load balancing: penalize variance in per-token route frequency
            route_freq = self.token_route_counts.mean(dim=0)  # (seq_len,)
            aux_loss = self.balance_weight * torch.var(route_freq)
            return x_out, aux_loss

        return x_out

This version includes the LayerNorm that was missing from the earlier snippet, plus load balancing loss. Note the clone() before scatter — without it, you get in-place modification errors during autograd.

Should You Use MoD in Production?

Depends.

If you’re serving a GPT-scale model with tight latency budgets and you control the training pipeline, MoD is worth testing. The 40% FLOP savings translate to real cost reduction at scale (millions of queries/day).

If you’re fine-tuning a pretrained model (e.g., Llama or Mistral), retrofitting MoD is painful. You’d need to retrain from scratch or use a hybrid approach where only new layers use MoD routing.

For small models (<1B params), the routing overhead dominates. Stick with dense layers or use quantization instead.

What I’m Curious About

The paper uses a fixed capacity ratio across all layers. What if early layers get 75% capacity (more feature extraction) and late layers get 25% (sparse refinement)? The router could learn a better division of labor.

Also, the load balancing loss is hand-tuned ($\alpha = 0.01$). Could you replace it with a learned Lagrangian multiplier that auto-adjusts based on routing entropy? I haven’t tested this yet, but it feels like the next iteration.

Finally, MoD + speculative decoding could be wild. Use a dense draft model to predict 4-5 tokens, then verify with a MoD target model that skips aggressively. You’d stack two sources of speedup.

FAQ

Q: Does MoD work with existing pretrained models like GPT-3 or Llama?

No, you need to train from scratch. Retrofitting MoD routing into a pretrained dense model would require retraining at least the router and LayerNorms. The paper doesn’t explore transfer learning from a dense checkpoint, so it’s unclear if you could bootstrap faster.

Q: How does MoD compare to pruning or quantization for reducing inference cost?

MoD is orthogonal. Pruning removes weights permanently; MoD skips computation dynamically per token. Quantization reduces precision; MoD reduces FLOPs. You could combine MoD with INT8 quantization for compounding savings, though the paper doesn’t benchmark this.

Q: What’s the memory footprint difference between MoD and dense models?

Parameter count is identical — same weights, same KV cache size. MoD only saves compute (FLOPs), not memory. If you’re memory-bound (common in autoregressive decoding), MoD’s benefit is smaller.

Final Take

MoD is one of the cleaner ideas in efficient Transformers. Unlike MoE (which adds parameters) or sparse attention (which adds complexity), MoD just asks: does this token need this layer?

At 50% capacity, the answer is “half the time, no” — and that’s enough to cut 40% of compute without hurting quality. For high-throughput serving, that’s a big deal. For interactive latency, it’s solid but not revolutionary (1.6x faster vs. 2-3x from speculative decoding).

If you’re building a custom LLM and can train from scratch, test MoD at 50% capacity. If you’re fine-tuning or deploying off-the-shelf models, wait for pretrained MoD checkpoints to show up.

And if you’re debugging routing collapse at 3am, keep dark chocolate espresso beans nearby. The auxiliary loss won’t tune itself.