Train One SAE Across Multiple Sparsity Budgets

Interpretability, priced per sparsity level

A sparse autoencoder turns an opaque residual-stream activation $x$ into a short list of nameable parts: an encoder scores latents, $\mathbf{l}=\sigma(W_{\mathrm{enc}}x+b_{\mathrm{enc}})$, a TopK activation keeps the $k$ largest, and the reconstruction is a scaled sum of decoder embeddings, $\hat{x}=\sum_{i\in\operatorname{top}_k} l_i\,e_i+b_{\mathrm{dec}}$. Gao and colleagues made this the standard recipe, and BatchTopK from Bussmann and colleagues enforces the same budget across a batch.

The budget $k$ is not a knob you can turn later. Train at $k=128$ and evaluate at 32, and reconstruction quality falls off a cliff; the dictionary was never asked to make its first 32 features stand on their own. So teams keep shelves of models: a sparse one for clean case studies, a dense one for faithful reconstruction, each trained separately.

We measure fidelity as the fraction of unexplained variance, $\operatorname{FVU}=\operatorname{Var}(x-\hat{x})/\operatorname{Var}(x)$, and sparsity as the number of active latents $\ell_0$, which for TopK activations equals $k$ exactly. The question we set out to answer: can one dictionary be honest at every $\ell_0$ at once?

Score every prefix, not just the last

TopK already sorts the surviving latents by activation value, and the standard objective wastes that ordering. We define the prefix reconstruction $\hat{x}_j$ built from only the top $j$ latents and penalize every prefix at once, $\mathcal{L}_{\mathrm{hier}}=\tfrac{1}{|\mathcal{J}|}\sum_{j\in\mathcal{J}}\lVert x-\hat{x}_j\rVert^2$ with $\mathcal{J}=\{1,\dots,K\}$. Under this loss the optimal reconstruction improves monotonically as $j$ grows, a property the vanilla objective simply does not have.

The extra terms cost almost nothing. All prefixes fall out of one cumulative sum in a single forward pass, and $\mathcal{J}$ can be thinned: scoring every 8th prefix, $\{1,8,16,\dots,128\}$, is indistinguishable from the full set while cutting loss terms eightfold. Only aggressive thinning, every 64th term, starts to hurt below $\ell_0=128$.

One dictionary, any budget

We trained a single HierarchicalTopK autoencoder with $K=128$ on layer-12 outputs of Gemma-2 2B, dictionary size $D=65{,}536$, on a billion FineWeb tokens, against TopK and BatchTopK baselines trained separately at $k\in\{32,64,128\}$ and a JumpReLU sweep. Across the whole $\ell_0$ range up to 128, the one model matches or beats the specialists on explained variance: the Pareto frontier belongs to a single set of weights.

Inference-time interpolation is exactly what fixed-$k$ training breaks, and our dense $\ell_0$ grid shows it: the baselines degrade away from their trained points while the hierarchical model glides. The one exception flatters nobody's defaults: BatchTopK, which mixes budgets across samples during training, keeps improving into $\ell_0\in[128,512]$, a range almost never used in practice.

The richer objective is free at step time

The obvious objection is speed, so we fused the loss into the sparse-decoder Triton kernel. Our fused vanilla-TopK kernel runs 5.3 percent faster than the stock implementation from Gao and colleagues, and the full hierarchical objective still comes in 4.1 percent faster, 10.05 milliseconds per step against 10.48. Training across every budget costs, in wall-clock terms, nothing.

Why fixed-k dictionaries rot when you turn the dial

The failure has a mechanism. Call a feature almost dead when it fires less than once per $10^5$ tokens. Evaluate a $k=128$ TopK or BatchTopK model at smaller budgets and almost-dead features multiply quickly; the hierarchical model keeps far more of its dictionary alive, because its strongest features were trained to carry reconstructions alone.

The dictionary's internal ordering tells the same story. In a well-behaved reconstruction, the cosine similarity between successively added decoder embeddings should fall as rank increases; vanilla TopK shows the opposite trend, while ours decreases monotonically. Swapping the activation rule at inference barely moves us either: a plain TopK autoencoder loses explained variance from 0.8577 to 0.8511 under a JumpReLU threshold, and the hierarchical model is virtually unchanged.

Interpretability survives the dial too. On the SAE Bench detection score, our model at $\ell_0=128$ stays nearly as interpretable as itself at $\ell_0=32$, while separately trained TopK and BatchTopK models become measurably less interpretable as they get denser. Density usually taxes interpretability; the hierarchical ordering pays that tax once, at training time.

What we have not shown

Our evidence lives on one model and one corpus, Gemma-2 2B and a FineWeb subsample, and our interpretability numbers are automated proxies rather than human studies. Nested dictionaries have relatives, Matryoshka-style autoencoders among them; ours differs by scoring every prefix inside a single objective. The claim we stand behind is operational: rank the features once, score every prefix, and stop retraining a new dictionary for every sparsity budget.