Mixture-of-Depths (MoD): Mathematical Foundations, Dynamic Compute Routing, Capacity-Constrained Tensors, and IsoFLOP Scaling

In standard autoregressive Transformer architectures, computational effort is distributed uniformly across all tokens in a sequence. Every token position $i \in \{1, \dots, S\}$ passes through every layer $l \in \{1, \dots, L\}$ , executing identical matrix multiplications across multi-head self-attention and feed-forward networks (FFN). This architectural constraint ignores the wide disparity in information density across natural language and formal reasoning: trivial grammatical connectors, punctuation, and predictable tokens require minimal processing, whereas complex multi-step reasoning tokens demand substantial non-linear transformation.

Prior attempts to introduce conditional computation into deep neural networks—such as Adaptive Computation Time (Graves, 2016), Universal Transformers (Dehghani et al., 2018), and early-exiting mechanisms like Confident Adaptive Language Modeling (Schuster et al., 2022)—often created dynamic computation graphs, variable sequence lengths per batch, and ragged tensor shapes. These non-uniform execution profiles conflict directly with modern hardware accelerators (TPUs and GPUs), which require static memory allocations and predetermined matrix dimensions to maximize systolic array utilization.

The Mixture-of-Depths (MoD) framework (Raposo et al., Google DeepMind, 2024) resolves this tension by enforcing a static computational budget per layer through capacity-constrained top- $k$ routing. By capping the number of tokens that participate in self-attention and MLP blocks at a predefined constant $C = k < S$ , MoD preserves static tensor shapes while allowing the network to dynamically assign depth across sequence positions.

1. The Mechanical Inefficiency of Uniform Compute

In a standard decoder Transformer with $L$ layers, hidden dimension $d$ , sequence length $S$ , and intermediate MLP dimension $d_{\text{ffn}} = 4d$ , the computational cost per token sequence per layer is dominated by matrix multiplications:

$\text{FLOPs}_{\text{attn}} = 4 S d^2 + 2 S^2 d$

$\text{FLOPs}_{\text{mlp}} = 2 \cdot (2 S d \cdot d_{\text{ffn}}) = 16 S d^2$

$\text{FLOPs}_{\text{layer}} = 20 S d^2 + 2 S^2 d$

Across all $L$ layers, total floating-point operations scale as $\mathcal{O}(L S d^2 + L S^2 d)$ . Standard Transformers evaluate this full computational graph regardless of the empirical loss gradient or Shannon entropy $H(Y_t \mid Y_{<t}, X)$ at token step $t$ .

When a model processes boilerplate syntax or repetitive boilerplate patterns, passing representations through dense FFN projections and full attention context updates yields negligible predictive improvement while consuming identical energy and memory bandwidth. Mixture-of-Depths replaces this static allocation with a dynamic policy that routes only a critical subset of tokens through computation blocks while streaming the remaining tokens through identity residual pathways.

2. Mathematical Formulation of Mixture-of-Depths

Let $X^l = [x_1^l, x_2^l, \dots, x_S^l]^T \in \mathbb{R}^{S \times d}$ denote the input matrix of token representations at layer $l$ .

2.1 Linear Router Projections

Each MoD block incorporates a lightweight routing mechanism parameterized by a weight vector $w_\theta^l \in \mathbb{R}^d$ . The router projects each token embedding $x_i^l$ into a scalar routing logit $r_i^l$ :

$r_i^l = (w_\theta^l)^T x_i^l \in \mathbb{R}, \quad \forall i \in \{1, \dots, S\}$

The vector of all sequence routing logits at layer $l$ is denoted as $R^l = [r_1^l, r_2^l, \dots, r_S^l]^T$ .

2.2 Capacity Factor and Quantile Selection

Unlike token-choice routing in Mixture of Experts (Shazeer et al., 2017), which allows variable numbers of tokens per expert and requires auxiliary load-balancing losses, MoD employs expert-choice routing across the sequence dimension.

The system defines a sequence capacity $C \in \mathbb{N}^+$ , governed by a user-specified capacity factor $\alpha \in (0, 1]$ :

$C = \lfloor \alpha \cdot S \rfloor$

To select the top- $C$ tokens, the router determines the dynamic selection threshold $P_\beta(R^l)$ corresponding to the $\beta$ -th percentile of $R^l$ , where:

$\beta = 1 - \frac{C}{S} = 1 - \alpha$

The subset of tokens chosen to participate in layer computation is defined by the index set:

$\mathcal{I}_{\text{active}}^l = \left\{ i \in \{1, \dots, S\} \;\middle|\; r_i^l > P_\beta(R^l) \right\}, \quad |\mathcal{I}_{\text{active}}^l| = C$

Let $\tilde{X}^l \in \mathbb{R}^{C \times d}$ represent the gathered matrix containing solely the representations of tokens in $\mathcal{I}_{\text{active}}^l$ .

2.3 Layer Output Update and Residual Bypass

Let $f^l(\cdot)$ denote the composite block transformation comprising multi-head self-attention followed by an MLP:

$f^l(\tilde{X}^l) = \text{MLP}^l\left(\text{LN}\left(\text{MHA}^l\left(\text{LN}(\tilde{X}^l)\right) + \tilde{X}^l\right)\right)$

For tokens outside $\mathcal{I}_{\text{active}}^l$ , computation is skipped entirely, passing the input directly to the output via the residual stream. For selected tokens, the output of $f^l(\tilde{X}^l)$ is scaled by the router weight $r_i^l$ to maintain gradient flow into the router parameters:

$x_i^{l+1} = \begin{cases} r_i^l \cdot f_i^l(\tilde{X}^l) + x_i^l & \text{if } i \in \mathcal{I}_{\text{active}}^l \\ x_i^l & \text{if } i \notin \mathcal{I}_{\text{active}}^l \end{cases}$

Input Sequence X^l (S tokens)
      │
      ├───> Router: r_i = w^T * x_i ───> Top-k Selection (Capacity C = α * S)
      │                                       │
      │                                       ├── Selected Tokens (C tokens) ───> [ Attention & MLP ] ───> (*) Scale by r_i ───┐
      │                                       │                                                                                │
      │                                       └── Bypassed Tokens (S - C tokens) ───────────────────────── Identity Bypass ───(+)
      │                                                                                                                        │
      └────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘
                                                                                                                               │
Output Sequence X^{l+1} (S tokens) <───────────────────────────────────────────────────────────────────────────────────────────┘

2.4 Differentiability and Gradient Routing

A primary challenge in conditional routing is ensuring differentiability across discrete token selection. In MoD, multiplying the block transformation $f_i^l(\tilde{X}^l)$ by the continuous scalar $r_i^l$ places $w_\theta^l$ directly on the computational graph.

The loss gradient with respect to router weights $w_\theta^l$ is:

$\frac{\partial \mathcal{L}}{\partial w_\theta^l} = \sum_{i \in \mathcal{I}_{\text{active}}^l} \frac{\partial \mathcal{L}}{\partial x_i^{l+1}} \cdot f_i^l(\tilde{X}^l) \cdot (x_i^l)^T$

This formulation eliminates the need for score function estimators (REINFORCE) or straight-through gradient estimators (STE). The gradient magnitude received by $w_\theta^l$ is proportional to the contribution of the block computation $f_i^l(\tilde{X}^l)$ toward reducing the overall language modeling loss $\mathcal{L}$ .

3. The Non-Causal Top-k Dilemma and Causal Autoregressive Sampling

During pre-training on complete sequences of length $S$ , calculating the top- $C$ tokens is executed across the entire sequence. However, in autoregressive decoding, tokens are generated sequentially one by one ( $t = 1, 2, \dots$ ). At step $t$ , the router cannot determine if $r_t^l$ ranks in the top $\alpha$ -fraction of the future sequence $R^l = [r_1^l, \dots, r_S^l]$ , because tokens $t+1, \dots, S$ do not yet exist.

To solve this causal barrier, Raposo et al. (2024) evaluated two mechanisms during training:

3.1 Auxiliary Sigmoid Centering Loss

The first approach enforces an auxiliary binary classification objective directly onto the routing logits $r_i^l$ . Let $y_i^l \in \{0, 1\}$ denote ground-truth participation based on sequence-level top- $k$ selection:

$y_i^l = \mathbb{I}\left(i \in \mathcal{I}_{\text{active}}^l\right)$

An auxiliary binary cross-entropy loss is applied over the sigmoid of router logits:

$\mathcal{L}_{\text{aux}}^l = - \frac{1}{S} \sum_{i=1}^S \left[ y_i^l \log \sigma(r_i^l) + (1 - y_i^l) \log (1 - \sigma(r_i^l)) \right]$

This loss pressures the router to output values $r_i^l > 0$ ( $\sigma(r_i^l) > 0.5$ ) for tokens that belong to the top- $k$ set, and $r_i^l < 0$ ( $\sigma(r_i^l) < 0.5$ ) for bypassed tokens. During autoregressive decoding, the model routes a token into the computation block if $\sigma(r_t^l) > 0.5$ , operating independently of future context.

3.2 Auxiliary MLP Predictor Head

The second approach trains a distinct lightweight classifier $g_\phi^l: \mathbb{R}^d \to [0, 1]$ in parallel with the main network. The router receives representation $x_i^l$ with a stop-gradient operator:

$\hat{y}_i^l = g_\phi^l(\text{stop\_gradient}(x_i^l))$

$\mathcal{L}_{\text{pred}}^l = - \frac{1}{S} \sum_{i=1}^S \left[ y_i^l \log \hat{y}_i^l + (1 - y_i^l) \log(1 - \hat{y}_i^l) \right]$

Because $\text{stop\_gradient}$ isolates the backbone from $\mathcal{L}_{\text{pred}}^l$ , the main language modeling representations remain unconstrained by auxiliary objectives.

Empirical evaluations in the DeepMind study demonstrated that the auxiliary predictor reaches greater than 97% routing agreement within initial training steps and exceeds 99% accuracy at convergence. When switching from sequence-level top- $k$ to causal predictor routing at inference, language model perplexity degradation is negligible (less than 0.2%).

4. IsoFLOP Analysis and Pareto Optimal Frontiers

To evaluate efficiency rigorously, conditional computation architectures must be assessed under compute-equivalent regimes (Chinchilla scaling laws, Hoffmann et al., 2022).

Validation Loss (Lower is Better)
  │
  │     Vanilla Transformer IsoFLOP Frontier
  │      \
  │       \    MoD IsoFLOP Frontier (Shifted Down & Right)
  │        \      \
  │         ●      \
  │          \      ● (Matches Baseline Loss with 50-66% Fewer FLOPs/Step)
  │           \      \
  │            ●      ● (IsoFLOP Optimal: Lower Loss, Larger Parameter Count)
  │             \
  │
  └──────────────────────────────────────────────────────── Step Time / FLOPs

4.1 Optimal Capacity Factor and Block Interleaving

DeepMind swept capacity factors $\alpha \in [0.125, 0.95]$ and structural routing patterns across training budgets ranging from $6 \times 10^{18}$ to $1 \times 10^{20}$ total FLOPs.

Key empirical findings include:

Aggressive Capacity Reduction ( $\alpha = 0.125$ ): Restricting computation blocks to process only $12.5\%$ of tokens per sequence ( $\alpha = 0.125$ ) yielded superior loss-to-FLOP trade-offs compared to moderate reductions ( $\alpha = 0.50$ or $\alpha = 0.75$ ).
Alternating Block Interleaving: Models that route every single layer suffered from representation drift. In contrast, interleaved routing—alternating between one full-capacity dense Transformer block ( $\alpha = 1.0$ ) and one MoD routing block ( $\alpha = 0.125$ )—achieved optimal stability. The dense blocks ensure frequent global attention mixing across all tokens, while the MoD blocks execute sparse, high-impact transformations.
Down-and-to-the-Right Frontier Shift: For a fixed training FLOP budget, the optimal MoD configuration contains more total parameters than the optimal dense baseline, yet executes a single forward pass in 50% to 66% of the wall-clock time.

5. Mixture-of-Depths-and-Experts (MoDE)

The routing mechanism of Mixture-of-Depths operates along the depth dimension (determining whether to compute), whereas Mixture-of-Experts (Switch Transformers, Fedus et al., 2022; ST-MoE, Zoph et al., 2022) operates along the width dimension (determining which expert to compute).

These two paradigms integrate into Mixture-of-Depths-and-Experts (MoDE) through two architectures:

5.1 Staged MoDE

In Staged MoDE, a binary MoD router first determines whether token $x_i^l$ participates in the block. If selected, the token passes to self-attention and subsequently enters a standard top- $k_{\text{moe}}$ MoE router that distributes it among $E$ FFN experts:

$x_i^{l+1} = \begin{cases} r_i^l \cdot \left( \sum_{e \in \text{Top-}k} g_e(x_i^l) \cdot \text{Expert}_e(x_i^l) \right) + x_i^l & \text{if } i \in \mathcal{I}_{\text{active}}^l \\ x_i^l & \text{if } i \notin \mathcal{I}_{\text{active}}^l \end{cases}$

5.2 Integrated MoDE

In Integrated MoDE, the MoD residual bypass is incorporated directly into the MoE routing matrix as an explicit $(E+1)$ -th "No-Op" expert:

$\mathcal{E} = \{\text{Expert}_1, \text{Expert}_2, \dots, \text{Expert}_E, \text{Identity}\}$

A unified router outputs a softmax distribution over all $E+1$ choices. DeepMind's empirical comparisons revealed that Integrated MoDE outperforms standard MoE with reduced capacity factors, because tokens explicitly learn to route to the identity operation rather than suffering arbitrary dropping when expert capacities saturate.

6. Structural Comparison Across Sparsity Paradigms

Standard Dense Transformer: Dense computation across all layers; static computation graph; no routing mechanism; full KV cache ( $L \times S \times d$ ); baseline (1.0x) execution.
Mixture-of-Experts (MoE): Width-level sparsity across FFN experts; static computation graph; top- $k$ routing over $E$ experts; full KV cache ( $L \times S \times d$ ); reduces active parameter footprint per token while preserving total model capacity.
Early Exit (CALM / ACT): Depth-level sparsity via prefix halting; dynamic / ragged computation graphs; confidence thresholding routing; reduced KV cache for halted tokens; token halts execution permanently at early layer $l < L$ .
LayerSkip: Depth-level sparsity via layer dropout during training; static / speculative computation graphs; early-exit draft verifier; full multi-layer verification; speculative drafting with early layer exits followed by full-model verification.
Mixture-of-Depths (MoD): Depth-level sparsity across selective intermediate blocks; static computation graph with fixed top- $C$ tensor capacity; top- $C$ quantile routing with causal auxiliary prediction; reduced KV cache for bypassed layers; reduces FLOPs per forward step by up to 50% to 66% while allowing tokens to rejoin later layers.

7. Reference PyTorch Implementation

Below is a complete, standalone PyTorch module demonstrating the capacity-constrained top- $k$ routing mechanism, the auxiliary sigmoid loss, and the residual bypass of a Mixture-of-Depths layer:

import torch
import torch.nn as nn
import torch.nn.functional as F

class MixtureOfDepthsBlock(nn.Module):
    """
    Mixture-of-Depths (MoD) Transformer Block.
    Routes a fixed capacity fraction (alpha) of tokens through Self-Attention & MLP,
    bypassing the remainder via a direct residual connection.
    """
    def __init__(self, d_model: int, n_heads: int, d_ffn: int, alpha: float = 0.125):
        super().__init__()
        self.d_model = d_model
        self.alpha = alpha
        
        # Router projection
        self.router = nn.Linear(d_model, 1, bias=False)
        
        # Core computation block (Self-Attention + MLP)
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ffn),
            nn.GELU(),
            nn.Linear(d_ffn, d_model)
        )
        
    def forward(self, x: torch.Tensor, causal: bool = False):
        """
        Args:
            x: Tensor of shape [batch_size, seq_len, d_model]
            causal: If True, uses causal thresholding (inference mode)
        Returns:
            out: Tensor of shape [batch_size, seq_len, d_model]
            aux_loss: Auxiliary classification loss for router training
        """
        B, S, D = x.shape
        capacity = max(1, int(self.alpha * S))
        
        # 1. Compute scalar routing logits: [B, S]
        router_logits = self.router(x).squeeze(-1)
        
        if not causal:
            # Training Mode: Sequence-level Top-k (non-causal selection)
            # Find top-C routing weights across sequence dimension
            topk_vals, topk_indices = torch.topk(router_logits, k=capacity, dim=-1)
            
            # Binary mask for active tokens: [B, S]
            mask = torch.zeros(B, S, device=x.device, dtype=torch.bool)
            mask.scatter_(dim=1, index=topk_indices, value=True)
            
            # Ground-truth targets for auxiliary router loss: 1 if in top-k, 0 otherwise
            targets = mask.float()
            aux_loss = F.binary_cross_entropy_with_logits(router_logits, targets)
            
        else:
            # Inference Mode: Causal thresholding (sigma(logit) > 0.5 <=> logit > 0)
            mask = router_logits > 0.0
            aux_loss = torch.tensor(0.0, device=x.device)
            
        # 2. Extract selected tokens: [B, capacity, D]
        # For batch implementations, gather active tokens into fixed-size capacity buffer
        out = x.clone()
        
        for b in range(B):
            active_idx = torch.nonzero(mask[b]).squeeze(-1)
            if len(active_idx) == 0:
                continue
                
            # If active tokens exceed capacity during causal inference, truncate to capacity
            if len(active_idx) > capacity:
                active_idx = active_idx[:capacity]
                
            selected_x = x[b:b+1, active_idx, :] # [1, C, D]
            
            # Compute Self-Attention on active tokens
            norm_x = self.norm1(selected_x)
            attn_out, _ = self.attn(norm_x, norm_x, norm_x)
            attn_res = selected_x + attn_out
            
            # Compute MLP on active tokens
            mlp_out = self.mlp(self.norm2(attn_res))
            block_out = attn_res + mlp_out # [1, C, D]
            
            # Scale block output by router weights (places router in gradient path)
            routing_weights = router_logits[b, active_idx].unsqueeze(-1) # [C, 1]
            scaled_out = block_out * routing_weights
            
            # Residual addition
            out[b, active_idx, :] = x[b, active_idx, :] + scaled_out.squeeze(0)
            
        return out, aux_loss

8. Summary and Architectural Takeaways

The Mixture-of-Depths paradigm establishes several foundational principles for efficient autoregressive model design:

Decoupling Parameters from FLOPs: MoD demonstrates that scaling parameter count without increasing per-step FLOPs improves convergence rates and downstream task accuracy on an isoFLOP basis.
Hardware-Harmonious Conditional Compute: Unlike early-exit and adaptive recurrence algorithms that produce ragged tensor executions, MoD preserves fixed capacity factors $C = \lfloor \alpha S \rfloor$ , maintaining static memory graphs and high accelerator compute density.
Interleaved Depth Sparsity: Optimal execution is achieved by interleaving dense global-attention layers with sparse 12.5% capacity MoD layers, ensuring robust long-range context integration while bypassing over 85% of redundant FFN and attention calculations in intermediate stages.
Causal Autoregressive Decoupling: Non-causal top- $k$ pre-training is converted to causal step-by-step inference via auxiliary sigmoid centering losses or stop-gradient classifier predictors with greater than 99% routing fidelity.

Mixture-of-Depths: Mathematical Foundations, Dynamic Compute Routing, Capacity-Constrained Tensors, and IsoFLOP Scaling