mixture-of-depthstransformer-architecturellm-inferenceefficiencytutorial

Mixture-of-Depths: Dynamic Token Routing in LLMs Guide

September 20, 2026

Mixture-of-Depths Architecture: Per-Block Router Scoring, Top-k Token Capacity Allocation, and Residual Skip Connections

The Problem: The Uniform Compute Waste in Standard Transformers

In conventional transformer models (such as standard Llama or GPT architectures), computation is completely static: every token passes through every single layer.

Whether an input sequence contains a trivial punctuation mark (e.g., a comma or period), a predictable conjunction (e.g., “and”, “the”), or a crucial mathematical operator in a complex proof, the model expends the exact same number of Floating Point Operations (FLOPs) across all $L$ transformer blocks:

$$\text{FLOPs}{\text{Token}} = 2 \times L \times d{\text{model}} \times \left(4 d_{\text{model}} + 8 d_{\text{ffn}}\right)$$

This uniform compute allocation suffers from severe structural inefficiencies:

Compute Misallocation: Simple, predictable tokens require very little semantic processing, yet they burn expensive Tensor Core cycles across 32 or 64 deep attention layers.
The Early-Exit Pipeline Stall: Previous attempts at conditional compute utilized Early Exiting—halting token processing once a classification confidence threshold was met. However, early exiting creates dynamic, jagged batch sizes that shatter GPU SIMD hardware pipelines, causing memory stalls and synchronization deadlocks.
Memory Bandwidth Throttling: Moving all intermediate hidden states through non-essential Multi-Head Attention blocks wastes scarce GPU high-bandwidth memory (HBM), exacerbating generation latency.

The Solution: DeepMind’s Mixture-of-Depths (MoD) Architecture

To eliminate uniform compute waste while preserving predictable hardware parallelization, research from Google DeepMind introduced Mixture-of-Depths (MoD).

Instead of deciding whether an entire sequence should exit early, Mixture-of-Depths dynamically routes individual tokens around transformer blocks on a layer-by-layer basis using a strictly fixed token capacity.

Mixture-of-Depths Execution Flow: Router Linear Projection, Top-k Token Gating, Selective Block Computation, and Identity Residual Recombination

The Three Mathematical Mechanics of MoD

Learned Scalar Routing: At each transformer layer, a lightweight linear router projects each token’s hidden state $\mathbf{x}_i \in \mathbb{R}^d$ into a scalar routing weight $r_i \in \mathbb{R}$:

$$r_i = \mathbf{w}_{\text{router}}^\top \mathbf{x}_i$$

Fixed-Capacity Top-$k$ Gating: Unlike threshold-based routing (which produces unpredictable batch shapes), MoD enforces a fixed compute budget: only the top-$k$ highest-scoring tokens in a sequence are permitted to enter the compute block. For a sequence length $T$ and a capacity factor $C \in (0, 1]$ (typically $C = 0.5$):

$$k = \lfloor C \times T \rfloor$$

Only the $k$ tokens with the largest router scores are selected. Because $k$ is constant across every batch, GPU tensor operations remain strictly deterministic with zero pipeline bubbles.

Residual Pass-Through (Zero Compute): The selected top-$k$ tokens pass through the standard Multi-Head Self-Attention and MLP blocks. The unselected $(T - k)$ tokens completely bypass the block via the identity residual connection:

$$\mathbf{y}_i = \begin{cases} \mathbf{x}_i + r_i \cdot f(\mathbf{x}_i) & \text{if } i \in \operatorname{top-}k({r_1, \dots, r_T}) \ \mathbf{x}_i & \text{otherwise} \end{cases}$$

Architectural Comparison: Dense vs. MoE vs. Mixture-of-Depths

Architecture	Routing Dimension	Active Parameters / Token	Static Compute FLOPs	Primary Hardware Benefit
Standard Dense Transformer	None (All tokens visit all layers)	100%	Fixed (High)	Baseline memory access
Mixture-of-Experts (MoE)	Width (Routes across parallel MLPs)	~20% – 30%	Fixed (High per active expert)	Enormous parameter scaling
Mixture-of-Depths (MoD)	Depth (Routes across sequence layers)	50% (Layer by layer)	Strictly bounded (50% reduction)	50% FLOP reduction & faster generation
MoDE (MoD + MoE)	Depth AND Width	~10% – 15%	Minimized	Maximum parameter efficiency & inference speed

Implementation: Building a Mixture-of-Depths Block in PyTorch

Let us implement a complete Mixture-of-Depths transformer layer in PyTorch featuring top-$k$ capacity routing, selective tensor indexing, and residual pass-through.

Step 1: Implement the Top-k Capacity Router

Create the linear router module that scores tokens and computes the selection mask:

import torch
import torch.nn as nn
from typing import Tuple

class MoDRouter(nn.Module):
    def __init__(self, hidden_dim: int, capacity_factor: float = 0.5):
        super().__init__()
        self.capacity_factor = capacity_factor
        self.router_linear = nn.Linear(hidden_dim, 1, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: Input tensor of shape [batch_size, seq_len, hidden_dim]
        Returns:
            topk_indices: Indices of tokens selected for computation [batch_size, k]
            routing_weights: Normalized scalar weights for selected tokens [batch_size, k, 1]
            topk_mask: Boolean mask of shape [batch_size, seq_len]
        """
        batch_size, seq_len, _ = x.shape
        k = max(1, int(self.capacity_factor * seq_len))

        # Compute router scalar score for each token: [batch_size, seq_len]
        router_logits = self.router_linear(x).squeeze(-1)

        # Select top-k tokens per sequence
        topk_scores, topk_indices = torch.topk(router_logits, k, dim=-1)

        # Softmax normalize across top-k tokens
        routing_weights = torch.softmax(topk_scores, dim=-1).unsqueeze(-1)

        # Create boolean mask
        topk_mask = torch.zeros(batch_size, seq_len, dtype=torch.bool, device=x.device)
        topk_mask.scatter_(1, topk_indices, True)

        return topk_indices, routing_weights, topk_mask

Step 2: Implement the MoD Layer with Selective Tensor Slicing

Route selected tokens through Multi-Head Attention and an MLP block while passing unselected tokens via residual:

class MoDTransformerLayer(nn.Module):
    def __init__(self, hidden_dim: int, num_heads: int, capacity_factor: float = 0.5):
        super().__init__()
        self.router = MoDRouter(hidden_dim, capacity_factor)

        # Standard compute blocks (LayerNorm + MHA + MLP)
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.attn = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(hidden_dim)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_dim, 4 * hidden_dim),
            nn.GELU(),
            nn.Linear(4 * hidden_dim, hidden_dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, hidden_dim = x.shape
        topk_indices, routing_weights, topk_mask = self.router(x)

        # Initialize output with residual connection
        output = x.clone()

        # Gather selected tokens: shape [batch_size, k, hidden_dim]
        batch_indices = torch.arange(batch_size, device=x.device).unsqueeze(1).expand(-1, topk_indices.size(1))
        selected_tokens = x[batch_indices, topk_indices]

        # 1. Attention Block over selected tokens
        normed_tokens = self.norm1(selected_tokens)
        attn_out, _ = self.attn(normed_tokens, normed_tokens, normed_tokens)
        selected_tokens = selected_tokens + attn_out

        # 2. MLP Block over selected tokens
        normed_tokens = self.norm2(selected_tokens)
        mlp_out = self.mlp(normed_tokens)
        selected_tokens = selected_tokens + mlp_out

        # 3. Scale by router weights
        processed_tokens = selected_tokens * routing_weights

        # 4. Scatter back into the output tensor
        output[batch_indices, topk_indices] = output[batch_indices, topk_indices] + processed_tokens

        return output

Step 3: Forward Pass Demonstration

Run a sequence through the MoD layer and observe token capacity execution:

# Initialize a 512-dim MoD layer with 50% compute capacity
layer = MoDTransformerLayer(hidden_dim=512, num_heads=8, capacity_factor=0.5)

# Simulate batch of 2 sequences, 16 tokens each
dummy_input = torch.randn(2, 16, 512)
output = layer(dummy_input)

print(f"[MoD Layer] Input shape:  {dummy_input.shape}")
print(f"[MoD Layer] Output shape: {output.shape}")
print(f"[Success] Processed exactly 8 of 16 tokens through attention; 8 tokens skipped computation.")

Synergies with MoE (MoDE) and Hardware Production Sizing in 2026

When deploying Mixture-of-Depths across large inference fleets:

The MoDE Architecture (Mixture-of-Depths & Experts): The ultimate efficiency architecture combines MoD with MoE. A layer router first decides which tokens receive computation (MoD), and an expert router then assigns those tokens to specialized expert networks (MoE). This achieves extreme sparsity, cutting FLOPs by up to 75%.
Prefill vs. Decode Phase Acceleration: During long-context prefill (e.g., ingesting a 10,000-token prompt), MoD cuts prompt ingestion times in half. During autoregressive decoding, tokens with high router scores (e.g., rare vocabulary or code symbols) receive deep attention, while common syntax passes through with minimal latency.
Training Stability via Auxiliary Load Balancing: Similar to MoE routing, unconstrained routers can collapse into routing only specific token positions. MoD training utilizes auxiliary loss functions to ensure router capacity is distributed evenly across semantic entities throughout the sequence.

Conclusion

Mixture-of-Depths breaks the dogma of uniform layer execution in transformer neural networks. By enforcing fixed-capacity per-block routing where only the most salient tokens enter attention blocks, MoD reduces computational FLOPs and inference latency by 50% without compromising foundational model intelligence.