Sparse Models › The Conditional Computation Paradigm

Michael BrenndoerferUpdated December 31, 20258 min read
Language AI HandbookAll chapters
Part 33 of 67Mixture of Experts
Chapter 239Sparse Models › The Conditional Computation Paradigm
Reading section 3 of 10 · The Conditional Computation Paradigm · 8 minSection 3 of 10 · 8 min

Book contents

Section 3 of 10 · 8 min read

Sparse Models › The Conditional Computation ParadigmOne section · Sections in this chapter3 / 10 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The insight behind sparse models is elegant: not all parameters need to contribute to every prediction. Different inputs may benefit from different computational patterns. A question about chemistry might activate different knowledge than a question about history, even within the same model. This observation suggests that the dense model's uniform treatment of all inputs might be wasteful: we're applying computational resources indiscriminately rather than directing them where they're most needed.

Conditional Computation

A computational paradigm where the operations performed depend on the input. Different inputs follow different computational paths through the network, activating different subsets of parameters based on the input's characteristics.

Human experts give an intuitive analogy. When a hospital faces a complex case, they don't have every specialist examine the patient. Instead, a routing process directs the case to relevant experts: perhaps a cardiologist and neurologist for certain symptoms, or an oncologist and radiologist for others. Each expert has deep knowledge in their domain, but only the relevant experts participate in each case. The hospital maintains a large total capacity (many specialists) while each patient receives focused attention from only the most relevant subset.

Sparse models implement this principle mathematically. Instead of a single large FFN layer, imagine dividing those parameters into multiple separate "expert" networks. A learned routing mechanism examines each input and decides which experts should process it. This routing decision is the key innovation: rather than treating all parameters uniformly, we learn to selectively activate the most relevant subset for each input.

The sparse layer computes its output as a weighted combination of expert outputs. This combination allows the model to blend specialized knowledge from different experts based on the input's characteristics. The weighting scheme ensures that more relevant experts contribute more strongly to the final output, while less relevant experts contribute proportionally less or not at all. For an input xx, the output is computed as:

y=∑i=1Ngi(x)⋅Ei(x)y = \sum_{i=1}^{N} g_i(x) \cdot E_i(x)

where:

  • yy: the output of the sparse layer for input xx
  • NN: the total number of expert networks available in the layer
  • Ei(x)E_i(x): the output of expert network ii when processing input xx
  • gi(x)g_i(x): the gating weight for expert ii given input xx, which determines how much expert ii contributes to the final output
  • ∑i=1N\sum_{i=1}^{N}: summation operator that sums over all NN experts (from expert 1 to expert NN)
  • gi(x)⋅Ei(x)g_i(x) \cdot E_i(x): the weighted contribution of expert ii, scaling the expert's output by how relevant it is to this input

This weighted sum allows the model to combine specialized knowledge from different experts in a flexible, input-dependent manner. The key property is that gi(x)g_i(x) is sparse: most values are zero, meaning most experts don't participate for any given input. This sparsity is what makes the computation efficient: we only compute Ei(x)E_i(x) for experts where gi(x)>0g_i(x) > 0. Typically, only the top-k experts with highest gating weights are activated, with gi(x)=0g_i(x) = 0 for all others. This selection mechanism turns a potentially expensive computation over all experts into a much cheaper computation over just a few.

To ensure computational efficiency while maintaining proper probability distributions, we impose two constraints on the gating weights. The first constraint enforces sparsity. This keeps only k experts are active for any given input. The second ensures the weights form a valid probability distribution, so the output is a proper weighted average. These constraints can be formalized as:

gi(x)>0 only for the k experts with highest router scoresgi(x)=0 for all other experts∑i=1Ngi(x)=1(weights sum to 1, forming a proper mixture)\begin{aligned} g_i(x) &> 0 \text{ only for the } k \text{ experts with highest router scores} \\ g_i(x) &= 0 \text{ for all other experts} \\ \sum_{i=1}^{N} g_i(x) &= 1 && \text{(weights sum to 1, forming a proper mixture)} \end{aligned}

These two constraints work together to let efficient sparse computation. The first constraint (sparsity) ensures that most gi(x)g_i(x) values are exactly zero, so we can skip computing Ei(x)E_i(x) for those experts entirely. This is where the computational savings originate: if only 2 out of 64 experts are active, we perform approximately 1/32 of the computation we would need for a dense layer with equivalent total parameters. The second constraint (normalization) ensures that the non-zero weights form a proper probability distribution, making the output a valid weighted average of expert outputs. This normalization preserves the mathematical properties that allow the model to be trained effectively with standard gradient-based optimization.

The routing mechanism implements these constraints through a three-step process that turns continuous router scores into sparse, normalized gating weights. First, compute unnormalized router scores ri(x)r_i(x) for all experts, typically via a learned linear projection followed by softmax. This linear projection learns which features of the input are relevant for selecting each expert. Second, select the top-k experts with highest scores and set all other weights to zero. This creates a sparse mask. This hard selection step is what lets the computational efficiency of sparse models. Third, renormalize only the selected weights so they sum to 1:

gi(x)={ri(x)∑j∈top-krj(x)if i∈top-k experts0otherwiseg_i(x) = \begin{cases} \frac{r_i(x)}{\sum_{j \in \text{top-k}} r_j(x)} & \text{if } i \in \text{top-k experts} \\ 0 & \text{otherwise} \end{cases}

where:

  • ri(x)r_i(x): the unnormalized router score for expert ii on input xx
  • top-k\text{top-k}: the set of kk experts with the highest router scores
  • ∑j∈top-krj(x)\sum_{j \in \text{top-k}} r_j(x): normalization constant that ensures selected weights sum to 1. This process ensures that only k experts perform forward passes, sharply reducing computation while maintaining the mathematical properties of a weighted mixture.

This mechanism breaks the dense model's coupling between capacity and compute. We can have many experts (large total parameter count), but only activate a few per token (small computational cost). The model gains capacity to store diverse knowledge while maintaining efficient inference. Different tokens benefit from different experts, so selective activation can be a better match for how knowledge is applied.

In[6]:
Code
%pip install torch
import torch
import torch.nn as nn
import torch.nn.functional as F


class ConceptualSparseLayer(nn.Module):
    """Demonstrates the sparse computation concept (simplified)."""

    def __init__(self, d_model: int, d_ff: int, n_experts: int, top_k: int):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k

        # Multiple expert FFN networks
        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(d_model, d_ff),
                    nn.ReLU(),
                    nn.Linear(d_ff, d_model),
                )
                for _ in range(n_experts)
            ]
        )

        # Router that decides which experts to use
        self.router = nn.Linear(d_model, n_experts)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, d_model = x.shape

        # Compute routing scores for each token
        router_logits = self.router(x)  # [batch, seq, n_experts]

        # Select top-k experts per token
        top_k_logits, top_k_indices = torch.topk(
            router_logits, self.top_k, dim=-1
        )
        top_k_weights = F.softmax(top_k_logits, dim=-1)

        # Initialize output
        output = torch.zeros_like(x)

        # Only compute outputs for selected experts (sparse computation)
        for i in range(self.top_k):
            expert_idx = top_k_indices[:, :, i]  # [batch, seq]
            weight = top_k_weights[:, :, i : i + 1]  # [batch, seq, 1]

            # Process each expert (simplified - real implementations batch this)
            for e in range(self.n_experts):
                mask = expert_idx == e
                if mask.any():
                    tokens_for_expert = x[mask]
                    expert_output = self.experts[e](tokens_for_expert)
                    output[mask] += (
                        weight[mask].squeeze(-1).unsqueeze(-1) * expert_output
                    )

        return output
Out[7]:
Console
Configuration: 8 experts total, 2 active per token (25.0% sparsity)

Sample routing decisions:
  Batch 0, Token 0: Experts [3, 7]
  Batch 0, Token 1: Experts [5, 0]
  Batch 0, Token 2: Experts [2, 3]
  Batch 0, Token 3: Experts [0, 7]
  Batch 1, Token 0: Experts [6, 2]
  Batch 1, Token 1: Experts [4, 3]
  Batch 1, Token 2: Experts [7, 0]
  Batch 1, Token 3: Experts [4, 2]
Out[8]:
Visualization
Heatmap grid with 8 token rows and 8 expert columns, where active expert assignments are marked with X in dark blue cells and inactive cells appear light blue, showing sparse and varied routing patterns across tokens.
Routing decisions for a batch of 8 tokens using top-2 expert selection. The heatmap displays active experts (dark blue) versus inactive ones (light blue) for each token (rows). The sparse and varied activation pattern illustrates how different tokens utilize different subsets of the available experts.

The configuration shows the core efficiency principle of sparse models. Each token activates only 2 out of 8 experts, using just 25% of available capacity. The routing decisions show that different tokens select different experts, with each making independent routing choices based on the learned router weights. While the layer contains 8 experts worth of parameters, any single token only triggers computation in 2 of them. This reaches the key sparse model trade-off: maintaining high total capacity (all 8 experts available) while keeping per-token compute cost low (only 2 experts activate per forward pass).