How latents switch on

Write the pre-activation as $z = W_{\mathrm{enc}}(x - b_{\mathrm{dec}}) + b_{\mathrm{enc}}$. Every architecture below is a different rule for turning $z$ into a sparse code $f$.

Shrinkage, measured

Piece 6 claimed that the L1 penalty shrinks what it keeps. The figure checks this on trained SAEs. Two SAEs are trained on the toy data of piece 7, one with an L1 penalty and one with TopK. Then we take inputs where exactly one true feature is active and ask how far the reconstruction reaches along that feature’s direction.

At $\lambda =$ 0.30 the L1 cloud runs parallel to the diagonal but below it, by a roughly constant amount. Every feature that fires is under-reported, and small ones vanish entirely. The shortfall grows with $\lambda$. The TopK cloud sits on the diagonal: the SAE chooses which latents fire by rank, so it never needs to bias their values to keep the code sparse. Now set TopK’s $k$ to 2 and look at the last row of the table. In two dimensions any two latents can reconstruct any point, so TopK reconstructs perfectly and stops finding the features. Faithful magnitudes and correct features are separate questions.

Three rules, idealised

To see the trade-offs, idealise. Suppose the dictionary is orthonormal, so $z$ is the true code plus noise. Then the L1 objective has a closed-form optimum, the soft threshold $f = \max(0, z - \lambda)$: every latent that survives is lowered by $\lambda$. A hard threshold $f = z\cdot\mathbb{1}[z \gt \theta]$ keeps magnitudes intact. TopK keeps the $k$ largest entries.

With L1 at 0.40, every blue bar on the second row sits below its tick, by exactly $\lambda$. Lower $\lambda$ to remove the bias and spurious latents appear. The penalty cannot both filter noise and leave magnitudes alone. JumpReLU separates the two jobs: the threshold 0.40 filters, and what passes keeps its size. TopK needs no threshold, but $k$ is fixed, so it over-fires when fewer than $k$ latents are truly active and under-fires when more are.

Four activation functions

$$\begin{aligned} \textbf{Gated:}\quad & f = \mathbb{1}[\pi_{\text{gate}}(x) \gt 0] \odot \mathrm{ReLU}(\pi_{\text{mag}}(x)), \qquad W_{\text{mag}} = e^{r_{\text{mag}}} \odot W_{\text{gate}}\\ \textbf{TopK:}\quad & f = \mathrm{TopK}_k(z)\\ \textbf{JumpReLU:}\quad & f_i = z_i\, H(z_i - \theta_i), \qquad \mathcal{L} = \lVert x - \hat x\rVert^2 + \lambda \textstyle\sum_i H(z_i - \theta_i)\\ \textbf{BatchTopK:}\quad & \text{keep the top } kB \text{ entries of } Z \in \mathbb{R}^{B\times m} \end{aligned}$$

Gated separates whether a latent is on from how large it is, and applies the L1 penalty only to the gate. It needs about half as many firing features for comparable reconstruction, in models up to 7B.Rajamanoharan et al., Improving Dictionary Learning with Gated Sparse Autoencoders, 2024.

TopK fixes $L_0 = k$ directly, which is what made clean scaling laws possible. It initialises the encoder to $W_{\mathrm{dec}}^\top$ and adds an auxiliary loss to revive dead latents.Gao et al., 2024.

JumpReLU learns a threshold $\theta_i$ per latent and penalises $L_0$ itself. The step function $H$ has no useful gradient, so $\theta$ is trained with straight-through estimators. On Gemma 2 9B it matched or beat Gated and TopK on reconstruction at a given sparsity, and it is the architecture of Gemma Scope.Rajamanoharan et al., Jumping Ahead, 2024.

BatchTopK keeps a fixed average $L_0$ across a batch, so easy tokens can use few latents and hard ones many. At inference it is replaced by a single learned threshold. It beat TopK on GPT-2 Small and Gemma 2 2B and roughly matched JumpReLU.Bussmann, Leask and Nanda, 2024.

Three changes to the objective

p-annealing replaces $\lambda\lVert f\rVert_1$ with $\lambda\lVert f\rVert_p^p$ and anneals $p$ from 1 toward 0 during training, starting convex and moving toward the count we want. It was tested on models trained on chess and Othello.Karvonen et al., 2024.

End-to-end SAEs stop asking for faithful reconstruction and ask for unchanged behaviour:

$$\mathcal{L}_{\text{e2e}} = \mathrm{KL}\big(p(\cdot \mid x)\,\big\Vert\, p(\cdot \mid \hat x)\big) + \lambda\lVert f\rVert_1 .$$

They need fewer total and fewer active features for the same downstream loss, but the whole model has to be in the training loop, and the dictionary is tied to the layers after it.Braun et al., Apollo Research, 2024. Their motivation: ordinary SAEs may learn “more about the structure of the dataset than the computational structure of the network.”

Matryoshka SAEs require every prefix of the dictionary, from a set of sizes $\mathcal{M}$, to reconstruct the input alone:

$$\mathcal{L} = \sum_{m' \in \mathcal{M}} \big\lVert x - W_{\mathrm{dec}}[:, {:}m']\, f_{:m'}(x) - b_{\mathrm{dec}} \big\rVert^2 + \lambda\,\mathcal{S}(f).$$

General features are forced into the early latents, where specific ones cannot absorb them. This reduces feature absorption (piece 12) at a small cost in reconstruction. Matryoshka SAEs are part of Gemma Scope 2.Bussmann, Nabeshima, Karvonen and Nanda, 2025; independently Zaigrajew, Baniecki and Biecek, 2025, for CLIP.

Every comparison in this piece is a comparison on the sparsity–fidelity frontier. Whether that frontier predicts usefulness is the question of piece 11.