Mixture of DepthsMoD

In the evolution of large language models (LLMs), one persistent challenge has been how to expand model capacity without incurring prohibitive compute costs. The standard Transformer architecture treats every input token identically, passing each token through the same number of layers and applying the same dense matrix operations regardless of semantic complexity. Yet intuition and empirical observation suggest that not all tokens are created equal. Punctuation marks, stop words, and other high-frequency elements often require far less processing than rare entities, long-range dependencies, or syntactically ambiguous phrases.

Mixture of Depths (MoD) emerges as a response to this imbalance. Drawing inspiration from Mixture of Experts (MoE), which dynamically routes tokens to specialized sub-networks, MoD introduces a dynamic compute allocation mechanism that decides, layer by layer, which tokens deserve deeper processing and which can bypass computation entirely. The result is an architecture that allocates compute where it is most needed, improving training and inference efficiency without sacrificing model expressiveness.

What Is Mixture of Depths

Mixture of Depths is a dynamic routing mechanism for Transformer-based networks. In a conventional Transformer block, every token flows through both the attention sub-layer and the feed-forward network (FFN) at every depth. MoD breaks this uniform paradigm. At designated MoD layers, a router evaluates incoming tokens and selects only a small, importance-weighted subset for actual computation. The remaining tokens skip the layer through residual connections, propagating unchanged to the next block.

This selective depth yields two primary benefits:

• Improved compute efficiency: By reducing the number of tokens processed at each layer, the model cuts floating-point operations (FLOPs), accelerating both training and inference.
• Greater effective capacity: Because compute is conserved on skipped tokens, practitioners can scale up parameter counts or add more layers while maintaining a computational budget comparable to a smaller, densely computed model.

In short, MoD allows a model to be "wide" in parameter space but "narrow" in active compute, adapting its depth to the demands of each input sequence.

Core Mechanism

The MoD architecture rests on three components: a lightweight router, a capacity factor that governs selection width, and a recombination step that restores sequence order after sparse computation.

The Router

At each MoD-enabled Transformer layer, a lightweight linear router assigns an importance score to every token. Given an input tensor $X \in \mathbb{R}^{N \times D}$, where $N$ is the token count and $D$ is the hidden dimension, the router computes scores using a projection matrix $W_r \in \mathbb{R}^{D \times 1}$:

$$S = \text{Softmax}(X W_r)$$

The resulting score reflects how much the model estimates a token will benefit from the upcoming attention and FFN operations. Tokens with higher scores are prioritized for computation; tokens with low scores are candidates for bypass.

Capacity Factor and Top-k Selection

The capacity factor, typically denoted as $k$ or expressed as a ratio, determines the fraction of tokens allowed to enter the dense sub-layers at a given MoD block. If the capacity factor is set to 0.25, for example, only the top 25 percent of scored tokens proceed through attention and FFN. The other 75 percent skip the layer entirely, traveling via residual connections to the subsequent block.

This top-k selection is dynamic. The specific tokens chosen can change from layer to layer and from sequence to sequence, allowing the model to route compute toward semantically rich or contextually complex tokens while ignoring trivial ones.

Recombination and Residual Flow

Because selected and skipped tokens follow different paths through the layer, their representations must be recombined before being passed forward. Processed tokens are scattered back into their original positions, while skipped tokens retain their pre-layer values through the residual stream. The output tensor thus preserves the original sequence length and ordering, ensuring that downstream layers receive a complete, coherent representation of the input.

Simplified Implementation

The following PyTorch-style pseudocode illustrates the core logic:

import torch
import torch.nn as nn

class MoDLayer(nn.Module):
    def __init__(self, d_model, capacity_factor=0.25):
        super().__init__()
        self.d_model = d_model
        self.capacity_factor = capacity_factor
        self.router = nn.Linear(d_model, 1)
        self.transformer_block = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model)
        )

    def forward(self, x):
        batch_size, seq_len, d_model = x.shape
        num_tokens = batch_size * seq_len
        k = int(num_tokens * self.capacity_factor)

        scores = self.router(x).squeeze(-1)
        flat_scores = scores.view(-1)

        _, topk_indices = torch.topk(flat_scores, k, sorted=False)

        output = x.clone().view(-1, d_model)
        flat_x = x.view(-1, d_model)
        selected_tokens = flat_x[topk_indices]
        processed_tokens = self.transformer_block(selected_tokens)

        output.scatter_(
            0,
            topk_indices.unsqueeze(1).expand(-1, d_model),
            processed_tokens
        )

        return output.view(batch_size, seq_len, d_model)

This sketch omits production-grade concerns such as gradient routing, load-balancing losses, and fused kernels, but it captures the essential data flow: score, select, compute, and scatter.

Advantages

Faster Training and Inference

Because only a subset of tokens activates the dense weights at each MoD layer, the effective FLOP count drops substantially. In training, this translates to shorter iteration times and higher throughput per GPU hour. In inference, reduced compute per token enables higher concurrent request handling and lower latency, particularly for long-context workloads where quadratic attention costs dominate.

Scalable Capacity

MoD decouples parameter count from active compute. A model can maintain a large embedding dimension, numerous layers, or wide FFN expansions while only engaging a fraction of tokens at each stage. This makes it possible to train larger models under fixed compute constraints, or to achieve the performance of a dense large model with the training budget of a smaller one.

Dynamic Adaptivity

Unlike static architectures that apply the same depth to every input, MoD adjusts compute dynamically. Simple passages—such as boilerplate text or repetitive formatting—receive shallow processing, while complex reasoning steps, rare vocabulary, or long-range dependencies trigger deeper engagement. The model effectively learns to allocate its "attention budget" across the sequence in a data-dependent manner.

Challenges

Hardware and Kernel Efficiency

Dynamic routing produces variable tensor shapes at runtime. Modern GPUs and TPUs