Replacement-aware SAE Training

Summary

One important use case for SAEs is in scalably identifying feature circuits (Ameisen et al. 2025), based on a replacement model in which transformer component outputs are re-encoded by their corresponding SAEs. A major drawback of this approach is that reconstruction errors tend to compound, since transformer models perform sequential operations on a shared residual stream. Standard residual stream SAE training is maximally myopic with respect to this use case, in that every SAE is trained in isolation. Even end-to-end training (Braun et al. 2024), which additionally looks at overall model performance, still trains each SAE’s parameters individually without consideration of how this will affect the downstream SAEs. We should be able to do better if we explicitly consider the replacement model during training. In this post I consider one of the simplest possible replacement-aware training paradigms, in which a term for the feature-space error of the next layer’s SAE is added to the standard SAE loss function. I present a brief quantitative comparison of SAEs trained by these various approaches on TinyStories (Eldan and Li 2023).

The key ideas of this post are:

1. Multi-SAE replacement models can be improved by considering the impact on downstream SAEs during training, and
2. It is better to measure this impact in SAE feature space rather than in the native activation space.

Motivation

As part of the project I worked on for SPAR, I looked into extracting circuits built out of SAE features in a replacement model (Ameisen et al. 2025) formed by re-encoding each layer’s output with an SAE. This proved to be infeasible without the use of error nodes (Marks et al. 2024), since re-encoding more than a single layer made the model’s output completely incoherent–it became unable to generate plausible English of any kind. In my opinion, this calls into question the validity of this entire approach to circuit extraction. If the “error” terms are so critical that without them the model completely breaks down, I don’t see how we can trust any causal story for the model’s behavior that handwaves them away as noise. While some degradation in performance is to be expected, I think it’s worthwhile to demand that our replacement model retain the qualitative property of “still being a language model!”

One way we might achieve this goal is by simply improving the quality of reconstruction at each layer. The SAE is just one of many possible scalable methods to extract interpretable features out of model activations. We might look to the broader field of sparse coding for inspiration; see for example Lewis Smith’s writeup of trying the gradient pursuit algorithm for this task in a Google DeepMind team blogpost (Nanda et al. 2024). I also have some ideas and a few small experiments I’ve done along this vein loosely based on LISTA (Gregor and LeCun 2010), which I will write up more formally at a later date. However, I am not terribly optimistic that better sparse coding techniques alone will scale to frontier models. This is because errors in the replacement model inherently have a compounding effect, since the outputs of every subsequent layer depend on the previous ones. This problem gets worse the deeper a model gets.

An orthogonal angle of attack is to mitigate the compounding effect itself for a given level of error. This requires some reworking of the training setup. The standard approach treats each transformer layer independently. This has some advantages, such as lower memory and compute requirements and the ability to train all layers in parallel. However, this setup fails to take into account the use case at the start of this section; we generally also care about the overall performance of the model. The way one cares about things in machine learning is by adding them to the loss function; (Braun et al. 2024) achieve this by adding a term for the KL divergence between variants of the model with and without each SAE. While this does, as expected, improve the performance of the replacement model, it is much more computationally and memory intensive. These costs can be mitigated quite substantially with a hybrid approach. (Karvonen 2025) found that following up a standard training run with an end-to-end KL finetuning step, with as few as 5% of the total training tokens, resulted in comparable performance. This step of course requires just as much memory and per-token compute, but running it on many fewer tokens lowers the overall cost.

While the end-to-end training setup does attempt to address overall model performance, it is subtly different from the motivating application. Specifically, it only considers the effect of re-encoding a single layer at a time, whereas for circuit extraction we want to replace all layers simultaneously. In the following I will formalize this distinction and propose an alternative procedure that more closely approximates the replacement model during training.

Replacement models

Let \(M\) represent a GPT-style transformer model, consisting of \(m\) layers \(T_i\) for \(i\in[0..m-1]\). Our goal will be to train a set of residual stream SAEs, one for each layer. For these purposes, we can treat the output of the transformer as the sequential composition of each layer \(T_i\) (where \(h\) outputs the logits from the final residual, and \(x\) is the embedding of the input):

\[ M(x) = h \circ T_{m-1} \circ T_{m-2} \circ \dots \circ T_0(x) \]

Then, the \(r\)-replacement model \(\widehat{M}^r\), where \(r \subseteq[0..m-1]\), is an alternative model that re-encodes some subset of layers with the corresponding SAE:

\[ \begin{gather} \widehat{M}^r(x) = h \circ \widehat{T}^r_{m-1} \circ \widehat{T}^r_{m-2} \circ \dots \circ \widehat{T}^r_0(x) \\ \widehat{T}^r_i=\cases{ SAE_i \circ T_i & if $i \in r$ \\ T_i & otherwise } \end{gather} \]

In this notation, we would designate the original model as \(\widehat{M}^{\emptyset}\), and the full replacement model as \(\widehat{M}^{[0..m-1]}\). For convenience, I will also define the activation of the \(i^{th}\) layer of the \(r\)-replacement model for input embedding \(x\) as

\[ \widehat{A}^r_i(x)=\widehat{T}^r_i\circ\widehat{T}^r_{i-1}\circ \dots \circ \widehat{T}^r_0(x). \]

Standard SAE setup

In this post, we will use the common TopK SAE variant, which is defined as \[ \begin{aligned} SAE_i(a) &= F_i(a)(W^{dec}_i)^T + b^{dec}_i\\ F_i(a)[t, \ell] &= \left\{\begin{array}{lr} Y_i(a)[t, \ell], & \ell \in \textrm{argtopk}(k, Y_i(a)[t, :]) \text{ and}\ Y_i(a)[t, \ell] \geq 0 \\ 0 & \text{ otherwise} \end{array} \right. \\ Y_i(a) &= a(W^{enc}_i)^T + b^{enc}_i . \end{aligned} \tag{1}\]

Note that \(F_i(a)\), the feature vector for that input, is constructed element-wise at each token position \(t\) from a linear transformation of \(y\), where only the top \(k\) largest elements are retained, followed by a ReLU.¹ The parameters of the SAEs are the encoder weights and bias (\(W_i^{enc} \in \mathbb{R}^{d_{model} \times d_{sae}}\) and \(b_i^{enc} \in \mathbb{R}^{d_{sae}}\)) and decoder weights and bias (\(W_i^{dec} \in \mathbb{R}^{d_{sae} \times d_{model}}\) and \(b_i^{dec} \in \mathbb{R}^{d_{model}}\)). \(k\in\mathbb{Z}^+\) is a hyperparameter controlling the sparsity of the feature vectors \(F_i\).

Then, the standard TopK SAE training objective over minibatches \(X\) is defined for each \(SAE_i\) as \[ \mathscr{L}^{std}_i = {1\over{n}} \cdot \sum_{j=0}^{n-1}||\widehat{A}^{\{i\}}_i(X_j) - \widehat{A}^{\emptyset}_i(X_j)||_2^2, \tag{2}\]

i.e. the MSE between the replaced and original \(i^{th}\) layer activations. The SAE for each layer is trained independently to minimize Equation 2.

KL fine-tuning and end-to-end training

(Karvonen 2025) first train each SAE using the standard loss of Equation 2, followed up by a fine-tuning phase that adds in the KL divergence of the replacement model including that layer’s SAE:

\[ \mathscr{L}^{finetune}_i = \mathscr{L}^{std}_i +{\beta\over{n}} \cdot\sum_{j=0}^{n-1} KL(\widehat{M}^{\{i\}}(X_j)\ ||\ \widehat{M}^{\emptyset}(X_j)) . \tag{3}\]

\(\beta\)² is computed for each minibatch such that the magnitudes of the two loss terms match. I also follow this approach in my implementation.

In end-to-end training (Braun et al. 2024), the reconstruction errors for each later layer are combined into a single training objective, with an added term for the KL divergence between the \(\{i\}\)-replacement model and the original. Adapting the notation somewhat for this post³, and computing \(\beta\) as above (dividing by the number of MSE terms so that the KL term is balanced to their average), this loss is \[ \mathscr{L}^{e2e}_i = {1\over{n}} \cdot \sum_{j=0}^{n-1} \left( \begin{align} \sum_{\ell=i+1}^{m-1} ||\widehat{A}^{\{i\}}_\ell(X_j) - \widehat{A}^{\emptyset}_\ell(X_j)||^2_2 \\ + {\beta \over{m-i}}\cdot KL(\widehat{M}^{\{i\}}(X_j)\ ||\ \widehat{M}^{\emptyset}(X_j)) \end{align} \right) \tag{4}\]

Note that this intentionally does not include a term for the current layer’s reconstruction error, only those downstream of it.

Replacement-aware training

The loss functions presented so far have dealt with a quite limited replacement model, in which only a single layer was re-encoded with its SAE. This is somewhat mismatched for applications like circuit extraction where we would like to use the full replacement model. The most obvious way to accommodate this would be to use the full replacement model directly in training, i.e. by training all SAEs simultaneously with the following loss:

\[ \mathscr{L}^{full} = {1\over{n}} \cdot \sum_{j=0}^{n-1} \left( \begin{align} \sum_{\ell=0}^{m-1} ||\widehat{A}^{[0..m-1]}_\ell(X_j) - \widehat{A}^{\emptyset}_\ell(X_j)||^2_2 \\ + {\beta \over{m}} \cdot KL(\widehat{M}^{[0..m-1]}(X_j)\ ||\ \widehat{M}^{\emptyset}(X_j)) \end{align} \right) \tag{5}\]

Note that this is almost identical to Equation 4; the important difference is that it uses the full replacement model in both the reconstruction and KL terms. Implemented naively, this would be even more expensive than end-to-end training, since we would have to have every SAE parameter and its optimizer state in memory, rather than just those for a single layer. This may prove to be impractical for larger models.⁴ However, I suspect that we can get most of the benefits by looking ahead by a single layer (where \(\beta\) is again computed as above to balance the two loss terms):

\[ \mathscr{L}^{next}_i = \mathscr{L}^{std}_i + {\beta\over{n}} \cdot \sum_{j=0}^{n-1} \cases{ KL(\widehat{M}^{\{i\}}(X_j)\ ||\ \widehat{M}^{\emptyset}(X_j)) & $i = m-1$ \\ ||\widehat{A}^{\{i,i+1\}}_{i+1}(X_j) - \widehat{A}^{\emptyset}_{i+1}(X_j)||^2_2 & $i \in [0..m-2]$ } \tag{6}\]

That is, in addition to the standard reconstruction loss, we add a term that for the final layer measures the KL divergence of the replacement model. For all other layers, this term instead depends on the reconstruction error of the next layer when both this layer and the next are replaced by their SAEs (note the two elements of the replacement set).

Why do I think looking ahead by one layer is likely to be sufficient? If we take seriously the idea that SAEs capture exactly the information that is most important to the model, then we ought to be able to treat the next layer’s SAE as if it screens off the influence of the current layer on those even further downstream. (We still need that next SAE, because the one we’re currently training hasn’t learned which features are important yet!) If we train the set of SAEs in reverse layer order, freezing their parameters as we go, we can approximate the effects of the full replacement model while only needing to run two transformer layers and SAEs in each training loop.

After my preliminary implementation of Equation 6, I had an idea that leans even harder into the notion that SAE features are the only important information passed through the residual stream. The most direct way to preserve this information is to reduce disruption to the features themselves. Reconstruction error gets at this only indirectly, since we expect the residual stream to include some noise. If we waste optimization power on recovering this noise, we get less of what we actually want. Hence, we should calculate the next-layer term in feature space. Because features are very sparse vectors, we likely wouldn’t want to do this with Euclidean distance; something like cosine distance is more appropriate. This leads to the following loss:

\[ \mathscr{L}^{feature}_i = \mathscr{L}^{std}_i + {\beta\over{n}} \cdot \sum_{j=0}^{n-1} \cases{ KL(\widehat{M}^{\{i\}}(X_j)\ ||\ \widehat{M}^{\emptyset}(X_j)) & $i = m-1$ \\ 1 - \frac{(F_{i+1}\circ\widehat{T}^{\{i\}}_{i+1}(X_j)) \cdot (F_{i+1}\circ\widehat{T}^{\emptyset}_{i+1}(X_j))}{||F_{i+1}\circ\widehat{T}^{\{i\}}_{i+1}(X_j)||_2 \cdot ||F_{i+1}\circ\widehat{T}^{\emptyset}_{i+1}(X_j)||_2} & $i \in [0..m-2]$ } \tag{7}\]

where the complicated term is the cosine distance between the features computed by the next layer’s SAE given the current layer’s SAE output, versus the features computed by the next layer’s SAE given the baseline activations.

I implemented two variants of this idea for this post. The first, which I call the next-layer method, directly uses Equation 7 as its loss function. The second variant, inspired by (Karvonen 2025), follows this up with a fine-tuning phase that adds in the full replacement model KL loss, i.e. the KL term from Equation 5.

Results

Each method was run using TinyStories-33M as the base model over 10 million tokens from the TinyStories dataset (Eldan and Li 2023), using TopK SAEs with \(k=100\) and an expansion factor of 4. See Appendix A for more implementation details. Figures in this section can be zoomed and scrolled after clicking on them.

I’ll start with the key results. Below, I plot a histogram of the KL divergence at each token for the full replacement models over 1 million tokens from the TinyStories validation set. Note that the x-axis is log-scaled, since the KL values vary by several orders of magnitude and appear to be approximately log-normally distributed:

While the distributions are visually quite similar, we can see in the reported summary statistics that the fine-tuned next-layer method achieves the lowest geometric mean KL divergence. Given the approximate log-normality of these distributions, I would argue that this is the most natural summary statistic. We can also see that end-to-end and KL fine-tuning get essentially identical results to each other, as observed by (Karvonen 2025), while standard training does much worse.

I think it’s also instructive to look at the reconstruction errors for each method, since these can give a rough estimate of SAE faithfulness–large error could be an indication that the SAE is using different mechanisms than the base model. In Figure 2 I plot the distribution of relative reconstruction error, defined as \[ ||\widehat{A}_i^{[i..m-1]}(x) - \widehat{A}^{\emptyset}_i(x)||_2 \over ||\widehat{A}^{\emptyset}_i(x)||_2 \] for each layer and method. I use this metric rather than MSE because norms vary a lot across layers and token positions, and this should make the values more directly comparable.