Activation Sparsity

Instrument a trained ReLU feed-forward block and count how many of its $d_{ff}$ neurons are nonzero for a given token. The answer is around ten percent, and nothing in training asked for that — it is what the block converges to on its own.

The consequence is arithmetic. If $h_i = 0$, then $h_i W_2^{(:,i)}$ contributes nothing, and the $i$-th column of the down-projection did not need to be read from memory. At decode time, when the block is bandwidth-bound and the weights are the traffic, not reading ninety percent of the largest matrix is close to a tenfold reduction in the cost of the block.

The problem is knowing which

$\mathcal{A}(x)$ is only known after computing $h$, which requires reading all of $W_1$. Skipping work you have already done is not a saving.

Deja Vu’s answer is to predict the set. A small learned classifier — a two-layer network reading the block’s input — guesses which neurons will fire, and only those rows of $W_1$ and columns of $W_2$ are fetched. The classifier is cheap, runs a layer ahead so its latency hides behind the previous block, and achieves recall high enough that the tokens produced are near-identical to the dense model’s.

Deja Vu found the same structure in attention — around 20% of heads carry the output for a given token — so the technique applies to both halves of the block, though the FFN is where the memory is.

The activation function decides whether any of this is possible

This is the catch, and it is fatal for most current models. GELU and SiLU produce small values, not zeros. A value of $10^{-4}$ still has to be multiplied, and its row still has to be read.

Llama, Mistral, Qwen and the rest use SwiGLU, so none of them is sparse in the sense required. Two responses exist. Mirzadeh et al. show that relufication — swapping the activation and fine-tuning briefly — recovers the sparsity at little quality cost. Others thresholding the small values, treating anything below $\epsilon$ as zero, which works but converts an exact optimisation into an approximation with an error budget to defend.

import torch
from torch import Tensor, nn


class SparseFFN(nn.Module):
    """Predict the live set, then touch only those rows and columns."""

    def __init__(self, dim: int, hidden: int, predictor: nn.Module):
        super().__init__()
        self.w1 = nn.Parameter(torch.empty(hidden, dim))
        self.w2 = nn.Parameter(torch.empty(dim, hidden))
        self.predictor = predictor                    # cheap; run one layer early

    def forward(self, x: Tensor, top_k: int) -> Tensor:               # x: [D]
        idx = self.predictor(x).topk(top_k).indices                   # [k]

        # The gather is the whole point: these are the only weights read.
        h = torch.relu(self.w1[idx] @ x)                              # [k]
        return self.w2[:, idx] @ h                                    # [D]

The gather is also the reason the win is smaller than the sparsity. Scattered row reads do not coalesce the way a contiguous matmul does, so the effective bandwidth is well below peak — the tenfold reduction in bytes becomes perhaps a twofold reduction in time.

Where it pays unambiguously is when the weights do not fit in memory at all. LLM in a Flash keeps the model on SSD and streams the predicted rows into RAM, which turns a model that could not run on the device into one that runs slowly. Against that baseline the comparison is not a speed-up, it is feasibility.

MoE is the same observation made structural. Rather than discovering after the fact that a tenth of the block fired, partition the block into experts and train a router to choose explicitly. The sparsity is then designed rather than emergent, the memory access is contiguous per expert rather than scattered, and no predictor is needed because the router is the predictor. Activation sparsity is the retrofit for models that were not built that way.