This is a cross-post from my blog. It is a follow-up to the methods I developed in a previous post on replacement-aware training.

---

Summary

A replacement model (Ameisen et al. 2025) is a modification to an LLM in which some subset of internal activations are replaced with ones computed as a sum of more interpretable features, such as those found by sparse auto-encoders (SAEs) (Cunningham et al. 2023) or cross-layer transcoders (CLTs) (Lindsey et al. 2024). When multiple components are thus re-encoded, methods inspired by structural equation modeling can be applied to construct feature circuits (Marks et al. 2024) that attempt to explain some aspect of the model's behavior in terms of these features. Because the re-encoding process is inherently lossy, errors from earlier components compound in later ones, resulting in severely damaged performance. In fact, applying current SAEs to just a few layers generally results in a replacement model that is no longer recognizable as a language model; it is unable to generate coherent output at all. To mitigate this effect, Marks et al. (2024) introduce error nodes into their recovered causal model, with values set such that they exactly cancel the re-encoding error of the corresponding SAE. I find this approach deeply unsatisfying, as I outline in Error nodes.

In this post, I present an alternative to error nodes: train SAEs that are robust to upstream errors. I call this approach replacement-aware training because it involves introducing a loss term for each SAE that penalizes distortions to the next layer's features, more closely matching the replacement model use case. I previously applied this method to a toy language model, and have now scaled it up to Gemma-2-2B (Gemma Team et al. 2024). The result is a suite of residual stream SAEs that can be used as-is in a full replacement model, which despite degraded performance, still retains language capabilities. I have made the code and weights for these SAEs publicly available; instructions for downloading and running them are available here.

I consider the methods outlined in this post an existence proof of the claim asserted in the title; they are by no means optimal. My overall approach was to keep adding to a bag of stackable tricks such that I could attain the goal of running a replacement model with mostly retained capabilities. While it does appear that no one innovation is sufficient on its own, I suspect that simplifications can be made, which I leave to future work. My core findings are:

1. Replacement-aware training greatly reduces KL divergence between the replacement model and base model as compared to standard methods, without sacrificing faithfulness;
2. A short (~1 million tokens) encoder-only fine-tuning stage can further reduce KL divergence, and this is effective even on replacement-naive SAEs;
3. Using a more powerful encoder based on LISTA (Gregor and LeCun 2010) also reduces KL divergence, and the mechanism for this seems to be improved robustness rather than raw reconstruction fidelity;
4. These reductions in KL divergence matter qualitatively. My replacement models retain some degree of capability as measured by benchmarks like MMLU, whereas alternatives do not remain coherent enough to even reliably give valid answers on these tasks.

Figure 1: Replacement-aware training produces SAE replacement models with significantly less divergence from the base model (Gemma-2-2B) as compared to standard methods. Values are calculated over 1 million validation set tokens.

Figure 2: Replacement models still suffer quite a bit of degradation on tasks like MMLU; no method I tested exceeded chance accuracy (red line) when training on the Pile alone (a). Mixing in 1 million tokens from a separate question-answering dataset (CommonSenseQA) during the KL fine-tuning phase (KL fine-tuning), representing less than 2% of the total training data, and applying PriDe debiasing (Zheng et al. 2023) recovers more performance (b). Even without this fine-tuning, using replacement-aware SAEs results in significantly more probability on the correct answers (c) and a vastly higher proportion of valid responses (d).

Table 1: Summary of results.

Method	Validation KL	Validation CE	MMLU NLL	MMLU ACC	MMLU ACC (finetune + debias)	MMLU Valid
Standard	1.00631	1.79459	4.65993	0.0746804	0.269217	0.345561
Replacement-aware	0.53557	1.13958	3.65254	0.193136	0.306251	0.781562
Replacement-aware with LISTA encoder	0.378151	0.758738	1.67595	0.245579	0.315444	0.965593
Baseline	-	0.387681	1.03801	0.541762	0.539835	0.999912

Background

In this section I set up a notational framework for explaining my methods and how they differ from existing approaches. I also briefly outline my case against error nodes.

Replacement models

A replacement model is a modification to a transformer model in which some subset of internal activations are re-encoded as a sparse sum of interpretable features. In this post I will focus on replacement models using SAEs that operate on the residual stream, though these methods should also be applicable to finer-grained replacement models. For precision I will use the term -replacement model, denoted , where , to specify which layers' residuals have been re-encoded. Note that this distinction between different replacement models is the defining difference between replacement-aware training and existing methods. I'll refer to the -replacement model as the full replacement model, while the -replacement model is simply the base model. Formally, letting be the function converting the final residual to probabilities, the number of layers, an SAE trained on the layer's residual, the transformer layer, and the embedding of an arbitrary input, define the -replacement model as

In this post I will explore variant encoder designs while keeping the decoder step a simple affine transformation. Hence, with the encoder function mapping activations to features, we have

where and are the decoder (or dictionary) weights and biases for the layer SAE.

We'll also find it convenient to refer to the activations of the layer in a replacement model

as well as the SAE features

Note that the features for SAE are defined by Equation 4 even when . Note also that these definitions imply that activations and features for layer are not affected by the inclusion or exclusion of in the replacement set:

so I will use these interchangeably when notationally convenient. The relationships among these terms are visualized as a causal graph below.

Figure 3: Replacement model of a 3-layer transformer. Features are inferred from activations by the SAE encoders reconstructed by the decoder and passed through the transformer block

Error nodes

In the introduction, I claimed that replacement models using more than a few SAEs become entirely incoherent. It might not be intuitively obvious how bad the problem is from just looking at plots of KL divergence. Thus, I think it's illustrative to compare actual rollouts from different replacement models. Note especially that the naive Gemma Scope output is not recognizable as language:

Effect of sleep quality on memory, executive function, and language performance in patients with refractory focal epilepsy and controlled epilepsy versus healthy controls - A prospective study. We aimed to evaluate the effect of sleep quality on memory, executive function, and language performance ... Shorter total sleep time, poorer sleep efficiency, and prolonged sleep latencies were observed to be associated with poor memory and executive function in patients with refractory epilepsy.

Gemma Scope

//_//_// //ing

Gemma Scope with encoder tuning

The results of this study suggest that the effect of

Replacement-aware

The results suggest that the sleep quality of patients with

Base model

<b>BACKGROUND:</b> The aim of this study

True continuation

Our study strongly suggests that sleep disturbances, mainly

Table 2: Greedy (temperature 0) rollouts for the specified full replacement models, limited to 10 new tokens. The prompt is the text of the first example from the pile-uncopyrighted validation split with the "PubMed Abstract" source, with the final sentence removed (reproduced partially above, with snipping indicated by ellipsis). The replacement models (first three outputs) receive geometric mean KL scores vs the base model on the full example of 11.0625, 1.1953125, and 0.359375, respectively. It's worth emphasizing that even the weakest method I discuss in this post (Gemma Scope with encoder tuning; see Headline results) recovers fluency, with no error nodes required.

Recognizing this issue, Marks et al. (2024) account for the errors introduced by SAEs in the replacement model by adding an additional term that exactly cancels them out. They refer to these terms as error nodes, as they are incorporated into the computational graph that is the basis of their circuit discovery method. Translated to the nomenclature of this post, this looks like:

Figure 4: Replacement model with error nodes which exactly cancel out the SAE reconstruction error. I have deliberately placed these causally upstream of the features rather than re-using the reconstructions which have the same value, as they are used in Marks et al. (2024). Not doing this would make feature-level interventions impossible, because their effects would be undone by the error correction (nodes labeled "+"). Note that the features and activations at layer are equivalent to those that would be observed in the -replacement model on input , since the error correction forces the SAE input to be equal to

This style of error node has become somewhat of a standard practice in feature-based circuits work; see for example (Ge et al. 2024; Ameisen et al. 2025; Hanna et al. 2025; Yang et al. 2026). One widely acknowledged limitation of this approach is that causal effects attributed to such nodes are inherently uninterpretable, and in some cases dominate the recovered circuits (Ameisen et al. 2025). A further drawback is that this is incompatible with learning-based approaches to circuit discovery, since the error node removes any incentive for more accurate reconstruction; see the discussion in § 2.3.3 of Caples et al. (2025). However, I believe there is a more fundamental issue with this approach: while it successfully builds consistent representations for observational data, it breaks down once we start trying to apply it to causal interventions. I relegate the full argument to Appendix C, since it is immaterial to my empirical results, but to briefly summarize:

1. Because each error node value is precisely tuned to cancel out the reconstruction error, its effects are different depending on whether we have intervened on any features.
2. This results, post-intervention, in a computational pathway with values depending on serial reconstructions, the very situation the error nodes were trying to prevent.
3. This is in fact an unavoidable consequence of feature space interventions, because such a pathway is the only possible mechanism by which these interventions could have any counterfactual implications.

The situation is more complicated for the attribution graphs of Ameisen et al. (2025) due to the use of CLTs rather than SAEs, resulting in many additional pathways that could potentially overwhelm the effect described above. Even so, most of these paths contain serial reconstructions of varying length: , on average, for the effect of an intervention at layer on the final layer. ^[1] This might explain the observation that effective steering magnitudes had to be guessed and checked, rather than falling directly out of the model. If this argument holds, then to get feature circuits that are capable of making accurate causal predictions such as the effects of steering, we must address the problem of compounding error head-on. The remainder of this post is my attempt to do exactly that.

Methods

Three core innovations underlie my improved replacement model: replacement-aware training, encoder fine-tuning, and a more powerful encoder inspired by LISTA (Gregor and LeCun 2010). I outline each below.

Replacement-aware training

Standard SAE training methods deal with each SAE in isolation. Even though end-to-end training (Braun et al. 2024) includes terms for downstream reconstruction errors, it still uses only one SAE in its replacement model. This approach can be glossed in my terminology as ^[2]

where is a batch of inputs.

In contrast, replacement-aware training uses a replacement model in which additional downstream layers are also replaced with the corresponding SAE outputs. Note that this implies a serial ordering to SAE training; we must have trained layers before we train the SAE for layer . In the most general form, the loss function for replacement-aware training is given by

where are functions of the SAE features in the replacement model starting at the current layer versus those that would have been computed from base model activations. The MSE on the SAE reconstructions is one such function, but in informal experiments I got slightly better results using the cosine distance on the features themselves. I also found that simply using the next layer (i.e. letting for ) rather than all downstream layers was sufficient for greatly improved performance. The intuition behind only including the next layer is that, due to the serial layout of the transformer blocks, the output of the next layer screens off any influence the current layer could have on the later blocks ^[3] . I view this term as a regularization that penalizes the SAE for destroying information that is useful downstream, even if it may not be relevant for the current layer. Since the next layer's features were also trained with this objective, this argument hypothetically chains all the way forward to the logits, so we are effectively implementing an approximation to end-to-end training of the full replacement model.

Additionally, to save on compute, I only include the term on the final layer, treating the logits as a virtual "next layer". For earlier layers, I run a separate KL fine-tuning stage for the final 20% of training tokens, similarly to Karvonen (2025). This allows its impact to be easily disentangled from that of replacement-aware training itself. I tried a few variants of this stage using different subsets of the downstream replacement model, which I'll describe below.

For all results in this post labeled "replacement-aware training", the main phase of training uses the loss function

which is simply the standard SAE loss function rendered in my notation, with an additional term penalizing distortions to the next layer's SAE features, or, in the case of the final layer, KL divergence from the base model. Note that the weights of are frozen before starting training on After Karvonen (2025), the weights are set per-batch such that the total magnitudes of the corresponding terms match. Standard training is equivalent to setting to 0.

KL fine-tuning

For the fine-tuning stage, we'll consider three variants: standard KL fine-tuning (Karvonen 2025), next-layer fine-tuning, and full replacement fine-tuning. Standard KL fine-tuning is simply

Next-layer fine-tuning is identical to Equation 7 for the final layer; for other layers a third term for the -replacement KL is added:

Finally, full-replacement fine-tuning is equivalent to Equation 9, but the full replacement model starting at layer is used. Note that this method is substantially more expensive to run in terms of both memory and compute. I also found it to be much less stable; training runs diverged unless I performed encoder tuning (Encoder tuning) beforehand, which makes it harder to compare to other methods. However, I include it here since it resulted in my best-performing replacement model:

As before, the weights are adjusted per-batch so that all terms have the same magnitude.

Encoder tuning

Extracting features from model activations is fundamentally an inference problem. The way we go about this need not be grounded in concrete model mechanisms, because by hypothesis the model's actual computation takes place in the smaller, superimposed space. In other words, while the decoding step of an SAE must be constrained by observable model behavior, we have flexibility in defining the encoder. We can use this freedom to build in robustness to the errors introduced by upstream SAEs in the replacement model.

Specifically, I propose an additional training stage that takes each layer's SAE decoder as fixed, tuning the encoder such that its output in the replacement model more closely matches its output on the base model, while also minimizing distortions to the next layer's features. Let be a set of SAEs defined as in Equation 2, and initialize a matching set of with these parameters. Then, proceeding in ascending layer order ^[4] , optimize the encoder parameters of each according to

The rationale for freezing the decoders is that the dictionaries learned at each layer were optimized for base model activations, which is a property we want to preserve; that's the entire reason we expect them to have any bearing on reality! As for the ordering: in the full replacement model, we are now running the SAEs off-distribution, with the exception of layer 0, since there is nothing to distort its input. By first tuning the layer 0 encoder, we bring the layer 1 encoder's input closer to its training distribution. Fixing the result and tuning the layer 1 encoder then aligns the input of the layer 2 encoder; and so on through the rest of the model.

Note that this procedure can be run on any set of SAEs, replacement-aware or not, and this offers substantially improved performance in the replacement model even with very short training runs, on the order of 1 million tokens. ^[5] I applied this procedure to all the SAEs discussed in this post to provide a fairer comparison.

Encoder variants

BatchTopK

The standard SAE computes features with a single linear operation followed by a nonlinear activation function. Several activation functions are in common use, but for this post I will be using BatchTopK (Bussmann et al. 2024), which is effectively a JumpReLU (Rajamanoharan et al. 2024) with a single global threshold. During training, a BatchTopK encoder computes features for a batch of tokens by keeping the top values after the linear transform and setting the rest to 0. My implementation is a very slight tweak to this, in that I always operate the function as a JumpReLU but with the threshold adaptively chosen to match the chosen sparsity as closely as possible, while enforcing non-negativity ^[6] :

is a hyperparameter controlling the desired sparsity and and are the encoder weights and bias. At inference time, this can be instead used as a JumpReLU:

where is a threshold learned during training ^[7] such that the expected sparsity matches the target.

LISTA

LISTA (Gregor and LeCun 2010) is an approach to sparse coding based on "unrolling" the classic ISTA algorithm (Daubechies et al. 2004) into a neural network. LISTA treats the matrices used to iteratively update the estimated sparse representation as learnable parameters; each iteration of the original algorithm is then represented by a single layer of a neural network. Translating the original presentation into the parlance of the current post, LISTA with iterations computes the following features:

where is termed the mutual inhibition matrix and is an activation function, which we will allow to vary per-iteration and per-layer. ^[8] Note that when this is exactly equivalent to a vanilla SAE with .

One major drawback of applying this directly as an SAE encoder is the number of parameters, which is now much larger than a standard SAE with since typically Fortunately, subsequent work brings this down to be in line with standard SAEs. Chen et al. (2018) argue for a weight-coupling scheme that effectively removes the parameter , but also add separate weights for each iteration, resulting in total parameters. Liu et al. (2019) build on this further, showing that the per-iteration weights can actually be taken to be scalar multiples of each other, resulting in a final parameterization that is just ; this is just the standard SAE encoder weights, plus scalars, which is negligible in comparison.

Chen et al. (2018) also propose what is effectively a "firm shrinkage" (Gao and Bruce 1997) activation function in place of the standard soft shrinkage. Specifically, the top inputs (in absolute value) are passed through unchanged, while the rest are subject to the standard soft shrinkage function. If we drop the soft shrinkage part, this is just the familiar TopK activation function! Putting this all together, swapping in BatchTopK in place of TopK, and adding the "Onsager correction" term from Borgerding and Schniter (2016) for improved stability, results in

For this post, I used , i.e. stepping up in equal increments, specifically [20, 40, 60, 80, 100]. This is a similar scheme to that used in from Chen et al. (2018). is the standard SAE decoder from Equation 2, and are the per-iteration scalar factors from Liu et al. (2019). For improved stability---I often observed feature values exploding once the SAE was off-distribution, as happens when using it in the replacement model---I additionally apply PyTorch's spectral norm parameterization to the encoder weight matrix and soft-cap the scalars to a max of 1 with tanh. This probably isn't the best way of dealing with the instability, but it was sufficient to get consistently well-behaved training runs for a 5-iteration variant. I leave optimizing the exact setup to future work.

I think once we write down LISTA this way, the way it works is fairly intuitive: we simply pass the current residual plus the Onsager correction (which is effectively a skip connection from the previous iteration) into the encoder, repeatedly. Because increases with each iteration, this means that additional features are selected conditional on the ones that have been already added. This could allow the SAE to handle hierarchical structure within the features, as has been proposed in prior work (Costa et al. 2025; Bussmann et al. 2025). Indeed, with sufficient hand-waving one could view this process as an approximation to (a learned version of) matching pursuit (Mallat and Zhang 1993), where due to the increasing for the BatchTopK activation functions, we are selecting new features per iteration instead of just 1.

Results

In this section I first present my core results, then investigate the individual contributions of each of my methods. I describe the training process in more detail in Training setup. Source code is available on GitHub. Data, including all SAE weights, validation and benchmark results, and tables and plots included in this post are available via huggingface bucket; see previous GitHub link for an explanation of the directory layout.

All results presented here are based on complete sets of SAEs I trained on all 26 layers of Gemma-2-2B in various configurations, using between 50 and 100 million tokens from the pile-uncopyrighted (Gao et al. 2020) training split per layer. The core training pipeline consists of 3 stages: initial training (Replacement-aware training), KL fine-tuning (KL fine-tuning), and encoder tuning (Encoder tuning). I found that performing KL fine-tuning with the full replacement model (Equation 10) was much less stable than using the this-and-next-layer-only replacement model (Equation 9), in addition to requiring much more memory (and thus more expensive hardware to achieve reasonable training times for 26 SAEs). While interchanging the order of this step with encoder tuning did mitigate this somewhat, doing so also makes it impossible to perform some of the analysis I want to show later in the results. Thus, I restrict the use of Equation 10 to a single "flagship" model, which I will exclude from the comparisons it is incompatible with.

To investigate how changes to the training setup and encoder design impact the end-to-end performance of the full replacement model, I trained using all configurations from the cross product

(standard main phase, next-layer main phase)
(no KL fine-tuning, KL fine-tuning on 10 million tokens)
("vanilla" SAE encoder, 5-iteration LISTA encoder)

For KL fine-tuning, I matched the loss to the corresponding main phase method, so that results labeled "standard" used Equation 8 and ones labeled "replacement-aware" used Equation 9. The amount of total training data varied because after my initial runs of 100 million tokens (which happened to be for vanilla next layer), looking at the loss history revealed severely diminishing returns after about 50 million tokens, so I shortened the runs that came after that analysis to save on compute costs. Below is a summary table of the SAEs, along with links to their weights:

Weights	Main training phase	KL fine-tune phase	Encoder
Encoder-tuned / Pre-tuning	Standard, 5e7 tokens	None	Vanilla BatchTopK, k=100
Encoder-tuned / Pre-tuning	Standard, 5e7 tokens	1e7 tokens	Vanilla BatchTopK, k=100
Encoder-tuned / Pre-tuning	Standard, 5e7 tokens	None	LISTA BatchTopK, k=[20, 40, 60, 80, 100]
Encoder-tuned / Pre-tuning	Standard, 5e7 tokens	1e7 tokens	LISTA BatchTopK, k=[20, 40, 60, 80, 100]
Encoder-tuned / Pre-tuning	Next layer, 1e8 tokens	None	Vanilla BatchTopK, k=100
Encoder-tuned / Pre-tuning	Next layer, 9e7 tokens	1e7 tokens	Vanilla BatchTopK, k=100
Encoder-tuned / Pre-tuning	Next layer, 5e7 tokens	None	LISTA BatchTopK, k=[20, 40, 60, 80, 100]
Encoder-tuned / Pre-tuning	Next layer, 5e7 tokens	1e7 tokens	LISTA BatchTopK, k=[20, 40, 60, 80, 100]
Encoder-tuned (no pre-tuned version available)	Next layer, 5e7 tokens	1e7 tokens, full replacement (Equation 10)	LISTA BatchTopK, k=[20, 40, 60, 80, 100]

Table 3: Summary of the SAE suite, with links to weights.

Since these SAEs were trained using my own choice of hyperparameters and code, which could be buggy, as a sanity check I also compared them to the Gemma Scope (Lieberum et al. 2024) residual stream SAEs, using encoder tuning (Encoder tuning) to match their sparsity to my SAEs. See Gemma Scope encoder tuning for further details.

Headline results

Applying my full method results in replacement models with greatly improved end-to-end performance compared to standard approaches. Below, I plot an expanded version of Figure 1 on 1 million tokens from the pile-uncopyrighted (Gao et al. 2020) validation split for the full suite of SAEs that I studied. Note that I report geometric means, rather than arithmetic means, because this is a more appropriate summary statistic given the per-token distribution of these values. ^[9] Note also that all SAEs in this chart are post encoder tuning, due to this having a disproportionate impact on the worst-performing replacement models. I explore this effect in more detail in Encoder tuning is effective.

Figure 5: Comparison of KL divergence and CE loss on 1 million validation tokens. The ordering was chosen to give a descending order to the CE loss, for easier comparison with benchmark results later in this section. There is a clear trend of replacement-awareness, KL fine-tuning, and LISTA encoders improving performance; I will explore their individual contributions in subsequent sections. All methods include encoder tuning. I found this step to have a disproportionate impact on the worst-performing methods, so this provides a fairer comparison. See Encoder tuning is effective for further elaboration.

Table 4: Validation KL and CE values over 1 million tokens

Method	Validation KL	Validation CE
Gemma Scope	1.66864	2.714
Standard	1.00631	1.79459
Standard LISTA	0.876545	1.57223
Standard + KL fine-tune	0.78684	1.51442
Replacement-aware (5e7 tokens)*	0.749334	1.43642
Replacement-aware	0.702321	1.39068
Replacement-aware LISTA	0.578471	1.14886
Replacement-aware + KL fine-tune	0.53557	1.13958
Standard LISTA + KL fine-tune	0.54007	1.10139
Replacement-aware LISTA + KL fine-tune	0.448847	0.953865
Replacement-aware LISTA + Full replacement fine-tune	0.378151	0.758738
Baseline	-	0.387681

Table 4: Validation KL and CE values over 1 million tokens. Note that I have added results for the 5e7-token checkpoints of the vanilla encoder replacement-aware SAEs, since the unmarked ones used 1e8 training tokens as explained above. There's an asterisk here, because due to the warm start (see Warm-starting from the next layer SAE modestly boosts performance, 5e7 tokens is probably enough), this is likely slightly over-estimating what the true performance would have been on a training run of this length. Still, this doesn't change the rank-ordering of methods.

To put these numbers into context, I also compared performance on the question-answering benchmarks MMLU (Hendrycks et al. 2020), CommonsenseQA (CQA) (Talmor et al. 2018), and the ARC Easy set (ARC-E) (Clark et al. 2018). I chose these benchmarks primarily for ease of implementation, since their data are readily available and evaluation only requires computing logits for single tokens, and because they offer a range of difficulties for the base model (54% accuracy for MMLU, 64% for CQA, 82% for ARC-E). See Benchmark setup for details on the exact prompting setup I used.

Figure 6: Top-logit accuracy (i.e. scoring each question as correct iff the argmax logit matches the token for the expected answer) for full replacement models across different multiple-choice question-answering datasets, with n-shot prompting as indicated (see also Benchmark setup). The ordering of methods is the same as in the previous plots (by descending validation CE), which reveals that benchmark accuracy is not simply a monotonic function of validation performance.

Table 5: Benchmark accuracy by method.

Method	MMLU ACC	CQA ACC	ARC-E ACC
Gemma Scope	0.0167221	0	0.00127389
Standard	0.0746804	0.0165153	0.0577495
Standard LISTA	0.0196988	0.0495458	0.0178344
Standard + KL fine-tune	0.124059	0.00165153	0.0318471
Replacement-aware	0.143933	0.0123865	0.0951168
Replacement-aware LISTA	0.260812	0.198183	0.243737
Replacement-aware + KL fine-tune	0.193136	0.100743	0.184713
Standard LISTA + KL fine-tune	0.242252	0.165153	0.251805
Replacement-aware LISTA + KL fine-tune	0.254859	0.184145	0.251805
Replacement-aware LISTA + Full replacement fine-tune	0.245579	0.184971	0.257749
Baseline	0.541762	0.641618	0.824204

Disappointingly, no method consistently breaks chance performance, despite clear improvements resulting from the combination of replacement-awareness, KL fine-tuning, and using LISTA encoders. Investigating further, I found that these differences are largely driven by some replacement models' (in)ability to even give a valid answer to a multiple-choice question, and a marked propensity for particular responses even when they are valid:

Figure 7: Proportion of questions where the top-logit answer was a valid multiple-choice answer ("A"-"D" in MMLU and ARC-E, "A"-"E" in CQA).

Table 6: Distribution of responses by benchmark and option label

Method	MMLU A	MMLU B	MMLU C	MMLU D	MMLU Other	CQA A	CQA B	CQA C	CQA D	CQA E	CQA Other	ARC-E A	ARC-E B	ARC-E C	ARC-E D	ARC-E Other
Gemma Scope	0.0429872	0.0267029	0.000963054	0	0.929347	0	0	0	0	0	1	0.00127389	0.000849257	0.000849257	0	0.997028
Standard	0.285677	0.0309928	0.0284539	0.000437752	0.654439	0.0776218	0.000825764	0.00825764	0	0	0.913295	0.185987	0.0203822	0.0169851	0	0.776645
Standard LISTA	0.0812467	0.00253896	0	0	0.916214	0.289017	0.000825764	0	0	0	0.710157	0.0615711	0	0.000424628	0	0.938004
Standard + KL fine-tune	0.214411	0.259849	0.0446507	0.00385222	0.477237	0.00495458	0.00247729	0	0	0	0.992568	0.0738854	0.0602972	0.0059448	0	0.859873
Replacement-aware	0.143407	0.39958	0.03038	0.0189109	0.407722	0.00495458	0.0363336	0.0165153	0	0	0.942197	0.0246285	0.362633	0.0157113	0	0.597028
Replacement-aware LISTA	0.141569	0.39958	0.229557	0.226668	0.00262651	0.000825764	0.199009	0.355904	0.443435	0.000825764	0	0.000424628	0.20552	0.394904	0.399151	0
Replacement-aware + KL fine-tune	0.464805	0.24803	0.0588338	0.00989319	0.218438	0.440132	0.0726672	0.00165153	0.000825764	0	0.484723	0.359236	0.366879	0.0322718	0.00212314	0.23949
Standard LISTA + KL fine-tune	0.250306	0.695325	0.0418491	0.00735423	0.00516547	0.465731	0.330306	0.102395	0.00247729	0	0.0990917	0.00764331	0.946072	0.044586	0.00127389	0.000424628
Replacement-aware LISTA + KL fine-tune	0.0750306	0.00884258	0.539223	0.375066	0.00183856	0.0619323	0.00330306	0.470685	0.464079	0	0	0.000849257	0.000849257	0.464544	0.533758	0
Replacement-aware LISTA + Full replacement fine-tune	0.150849	0.192085	0.217125	0.405533	0.0344073	0.0809249	0.146986	0.0701899	0.655656	0.0462428	0	0.036518	0.126539	0.529087	0.294268	0.0135881
Baseline	0.200228	0.220714	0.305463	0.273507	8.75503e-05	0.260941	0.176713	0.211396	0.204789	0.14616	0	0.216561	0.24586	0.268365	0.269214	0
True answer	0.227981	0.250044	0.253984	0.267992	0	0.193229	0.208918	0.197358	0.206441	0.194055	0	0.249682	0.24586	0.267516	0.236943	0

Table 6: Distribution of responses by benchmark and option label, with "Other" aggregating any invalid top-logit responses. The true answer distributions, which are not perfectly balanced across labels, are reported in the bottom row for comparison. Though the specific bias varies with the training method, all replacement models choose certain responses disproportionately to their true distribution for most data sets. This bias doesn't necessarily match the (smaller) bias shown by the base model.

So far, it appears that improved replacement models are able to recover some degree of in-context learning, to the extent that they can at least consistently recognize when a multiple-choice answer is required, but then they hit a ceiling. These tasks are likely being made more difficult by the relatively poor generalization of SAEs outside of the distributions they were trained on (Kissane et al. 2024; Heindrich et al. 2025). To mitigate this, I re-ran KL fine-tuning (KL fine-tuning) for the methods that included that step while interleaving a mix of questions and answers from the CQA train split. ^[10] Specifically, 10% of tokens were from variable-length examples matching the n-shot template used in benchmarking (see CQA fine-tuning) and 90% were from pile-uncopyrighted (as in all other training runs). Since the fine-tuning step is 20% of the total token budget, this affects 2% of the total training data. Additionally, to mitigate the bias observed in Table 6, I applied PriDe (Zheng et al. 2023), which estimates and adjusts for a question-independent prior probability on each label by augmenting the dataset with additional copies of a subset of examples with their options permuted. Note that this implicitly uses the probabilities of responses conditioned on the answer being valid, so this isn't quite directly comparable to the previous plot, at least for the Standard replacement model. However, since fine-tuning saturated the valid response fraction for the other methods (see Figure 9), this is directly comparable for them.

Figure 8: Fine-tuning on the CQA train split and applying PriDe to debias the model's responses recovers performance to well above chance. This generalizes across different question-answering benchmarks.

Figure 9: Fine-tuning on the CQA train split results in nearly 100% valid answers, except for Standard training on MMLU and CQA.

Table 7: Accuracy and valid answer fraction for CQA-finetuned replacement models

Method	MMLU ACC	CQA ACC	ARC-E ACC	MMLU Valid	CQA Valid	ARC-E Valid
Standard + Finetune with CQA	0.269217	0.245252	0.346921	0.877605	0.933939	0.997028
Standard LISTA + Finetune with CQA	0.299685	0.380677	0.491295	0.999737	0.995871	1
Replacement-aware + Finetune with CQA	0.306251	0.477291	0.50276	0.979075	0.992568	0.994055
Replacement-aware LISTA + Finetune with CQA	0.315444	0.464905	0.506157	0.999912	1	1
Baseline	0.539835	0.646573	0.827176	0.999912	1	1

Table 7: Accuracy and valid answer fraction for CQA-finetuned replacement models. Responses are scored as correct when the PriDe-adjusted probability of the expected answer is the highest among valid response tokens. Responses are marked as valid, as before, based on the unadjusted argmax logit.

These plots show a clear advantage for both replacement-aware training and the LISTA encoder, but no apparent additional benefit for their combination. In the following, I will explore the individual effects of each of these variants.

Replacement-aware training does not sacrifice faithfulness

Another important metric to consider is the quality of the SAE reconstructions. Because SAEs are learned, it is possible that they pick up on patterns that aren't actually used by the underlying model. Their utility as explanations of model behavior is thus contingent on their representations being grounded in the true activations; they should be faithful. To measure this, I report the relative reconstruction error (RRE) ^[11] , defined as , averaged across 1 million validation tokens. Note that I am plotting trimmed means (arithmetic means calculated with values in the percentile and above discarded) to deal with extreme outliers for some methods. Note also that methods using the LISTA encoder don't have this outlier problem, which may be noteworthy in and of itself; see the table of un-trimmed arithmetic means below (Table 9).

Figure 10: Comparison of RRE for each layer and method in the full replacement models, reported as 99th-percentile-trimmed mean over 1 million validation tokens. As in the previous charts, all methods include encoder tuning. I'll break this down into more comprehensible comparisons later on, but for now the key things to note are the nearly identical overall shapes and the fact that the lowest curves on the plot swap near layer 9.

Table 8 (trimmed means):

Layer	Gemma Scope	Standard	Standard LISTA	Standard + KL fine-tune	Replacement-aware	Replacement-aware LISTA	Replacement-aware + KL fine-tune	Standard LISTA + KL fine-tune	Replacement-aware LISTA + KL fine-tune	Replacement-aware LISTA + Full replacement fine-tune
0	0.193942	0.177245	0.231459	0.199465	0.208982	0.225631	0.211795	0.248802	0.233452	0.222346
1	0.268334	0.251331	0.303519	0.268541	0.275568	0.295287	0.278979	0.31596	0.303184	0.296952
2	0.330986	0.311834	0.37207	0.329103	0.333678	0.360544	0.337907	0.382992	0.368993	0.357601
3	0.363535	0.345263	0.403357	0.36031	0.364085	0.390228	0.368741	0.411264	0.398023	0.385819
4	0.384425	0.362876	0.423958	0.376016	0.375815	0.402421	0.381132	0.426811	0.408316	0.397075
5	0.41011	0.389561	0.446271	0.402402	0.407185	0.429146	0.412789	0.449097	0.434949	0.422925
6	0.437652	0.41251	0.464583	0.42311	0.425959	0.442553	0.430163	0.462777	0.447603	0.435399
7	0.449956	0.419771	0.467312	0.428928	0.430136	0.44132	0.433336	0.464627	0.444996	0.434418
8	0.484028	0.447975	0.489743	0.456143	0.454322	0.459569	0.455919	0.485343	0.463562	0.453304
9	0.509427	0.470928	0.504223	0.478142	0.474306	0.475266	0.473732	0.500662	0.478755	0.46882
10	0.532535	0.499275	0.520775	0.505812	0.497585	0.493186	0.494816	0.516641	0.496991	0.486211
11	0.544106	0.51074	0.527727	0.516968	0.511812	0.502562	0.505889	0.524361	0.506198	0.496472
12	0.570581	0.547172	0.544888	0.553911	0.538551	0.519254	0.527013	0.541353	0.522935	0.513989
13	0.558897	0.534425	0.536446	0.537007	0.53373	0.511878	0.519638	0.532618	0.515275	0.506552
14	0.5698	0.539259	0.542394	0.543115	0.547311	0.518117	0.528455	0.537851	0.522107	0.514402
15	0.565108	0.529197	0.533348	0.529574	0.536784	0.505664	0.515842	0.525883	0.507875	0.501403
16	0.574498	0.536073	0.540352	0.534911	0.543467	0.512318	0.523451	0.532257	0.513726	0.509224
17	0.581819	0.540464	0.540796	0.539851	0.549293	0.512235	0.526735	0.532906	0.513501	0.50926
18	0.596491	0.552398	0.550831	0.549245	0.553883	0.517598	0.532112	0.539012	0.517019	0.514015
19	0.611524	0.565073	0.561546	0.560709	0.564117	0.528256	0.542173	0.547873	0.527173	0.523992
20	0.614504	0.566902	0.560578	0.562412	0.561919	0.528131	0.542332	0.546141	0.526591	0.523626
21	0.622302	0.574336	0.565371	0.569478	0.567015	0.534524	0.548035	0.550566	0.532583	0.530597
22	0.646236	0.595201	0.581693	0.589048	0.585487	0.552871	0.566766	0.567643	0.550386	0.548515
23	0.660624	0.608895	0.592959	0.604015	0.599388	0.567679	0.581178	0.581343	0.565703	0.563014
24	0.660773	0.607013	0.586682	0.602311	0.595343	0.564122	0.57854	0.577547	0.561304	0.561596
25	0.695859	0.635952	0.613045	0.628771	0.634854	0.653385	0.6168	0.605532	0.637196	0.628912

Table 9 (arithmetic means):

Layer	Gemma Scope	Standard	Standard LISTA	Standard + KL fine-tune	Replacement-aware	Replacement-aware LISTA	Replacement-aware + KL fine-tune	Standard LISTA + KL fine-tune	Replacement-aware LISTA + KL fine-tune	Replacement-aware LISTA + Full replacement fine-tune
0	0.197375	0.181367	0.234401	0.203511	0.213328	0.228263	0.216219	0.25174	0.236098	0.224954
1	0.272075	0.255268	0.306807	0.272463	0.279442	0.298186	0.282987	0.319248	0.306134	0.299973
2	0.334898	0.31601	0.375714	0.333247	0.337637	0.363738	0.342002	0.386563	0.372217	0.360893
3	0.367865	0.349734	0.407466	0.364777	0.368322	0.39403	0.373097	0.41526	0.401814	0.3897
4	0.388987	0.367535	0.42787	0.380633	0.379997	0.40596	0.385389	0.430493	0.41179	0.400618
5	0.417579	0.3948	0.449933	0.407684	0.411218	0.432402	0.416843	0.452513	0.438122	0.426114
6	0.478615	0.420728	0.468045	0.431769	0.430339	0.445658	0.434287	0.466022	0.450629	0.438485
7	0.783516	0.440558	0.47072	0.451718	0.436192	0.444394	0.437968	0.467843	0.447992	0.437434
8	2.53808	0.529498	0.493221	0.548096	0.467963	0.462664	0.462072	0.488596	0.46657	0.456371
9	7.47082	0.716758	0.507778	0.755365	0.526854	0.478424	0.484614	0.504014	0.481846	0.471917
10	5.87248	1.28125	0.524351	1.31222	0.775646	0.496404	0.520131	0.520004	0.500114	0.489362
11	4.95913	0.982886	0.531393	0.909446	2.03486	0.50584	0.570588	0.527782	0.509387	0.499734
12	3.4119	1.28867	0.548488	1.26147	3.21678	0.522569	0.750758	0.544776	0.526175	0.517219
13	2.66282	1.03302	0.539915	0.992539	1.82443	0.515087	0.945644	0.535892	0.518395	0.509676
14	5.50588	1.14064	0.546012	1.17881	2.28648	0.521444	1.51361	0.541297	0.525367	0.517707
15	12.5383	1.87798	0.537195	1.79988	1.76152	0.509234	1.82228	0.529565	0.511375	0.504876
16	30.1178	4.08077	0.544361	3.89248	1.90042	0.51612	2.67934	0.536117	0.517475	0.512819
17	62.3586	11.6783	0.544943	10.9598	2.20115	0.516189	3.82413	0.53695	0.517433	0.512922
18	130.176	31.8794	0.555305	26.3438	2.14597	0.52193	4.01894	0.543447	0.521358	0.518128
19	293.178	62.3939	0.565969	50.0874	2.43569	0.532574	4.82449	0.552324	0.531508	0.528159
20	500.817	93.4269	0.565099	75.9513	2.08928	0.532611	3.37958	0.550756	0.531071	0.527941
21	824.562	83.2858	0.570109	68.4853	2.07705	0.539164	3.23138	0.555387	0.537262	0.535094
22	1612.39	80.2373	0.5863	70.8471	2.37336	0.557418	3.6818	0.572334	0.554961	0.553041
23	2903.25	74.5901	0.597467	70.4339	2.41424	0.572087	4.01207	0.585901	0.570176	0.567434
24	3278.04	69.9742	0.591371	72.1284	2.63609	0.568316	4.44726	0.58197	0.565534	0.566004
25	3235.99	41.0329	0.619658	55.0914	2.14274	0.657815	3.57273	0.611568	0.641884	0.636062

I notice a couple of noteworthy aspects of this chart. First, the overall shape is very similar across all methods. We'll investigate the specific shape shortly, but I'd like to also call attention to the region between layers 5 and 10. At earlier layers, standard training outperforms replacement-aware methods, but the lines cross over somewhere in this region and stay that way until the final layer. To make this clearer, here's the same plot limited to just standard training and the best-performing replacement-aware version:

Figure 11: Subset of Figure 10 for just standard training and the best-performing replacement-aware method. While standard training starts off with better reconstruction, it is overtaken after 9 layers.

In fact, replacement-aware training seems to be comparable or even strictly better by this metric when compared with its replacement-naive counterparts:

Standard encoder

Figure 12

KL fine-tuning

Figure 13

LISTA

Figure 14

LISTA + KL fine-tuning

Figure 15

With the exception of the non-finetuned standard encoder, where the values are essentially equal past the early layers, we actually see lower RRE when using replacement-aware training. In other words, replacement-aware training is not sacrificing faithfulness to achieve its improved end-to-end performance! ^[12]

A simple model of reconstruction error

The reconstruction error of each SAE, considered in isolation, has been a primary metric for almost all prior work on SAEs for mechanistic interpretability. But I don't think it makes sense to consider the SAEs in isolation; a good representation needs to capture all relevant information, where "relevance" includes sufficiency for constructing accurate downstream representations. As I just showed in the previous section, replacement-aware training is superior to standard training in this regard. As might be expected from the form of Equation 7, which is the standard SAE objective plus an additional term, this trades off against the traditional metric. As I'll show in this section, the primary advantage of replacement-aware training is a reduction in the compounding effect of SAE error, which grows with model depth.

Let's now return to the shape of Figure 10. A straightforward model of reconstruction error is to consider two sources: error caused by an intrinsic "difficulty level" of sparse coding at each layer, and a second-order effect caused by the reconstruction target itself being offset from the true activations. While we could in principle quantify the first effect by estimating the entropy of the high-dimensional joint distribution of activations at each layer, we already have a much more convenient proxy. Prior to encoder tuning, every SAE was trained with an objective that included the MSE for the current layer against the true activations. We should thus expect the reconstruction performance in the single-SAE-on-baseline-activations case to be inversely proportional to the underlying difficulty, modulo some offset due to MSE trading off against other objectives for some methods. And indeed, we see essentially the same shape across all the methods I tested:

Figure 16: Comparison of RRE with base model inputs, i.e. using the -replacement model rather than full replacement model. Trimming outliers was unnecessary, because all SAEs are now fully on their training distributions. Note that for this and all subsequent single-layer replacement model charts I use the pre-encoder-tuned versions of the corresponding SAEs.

Figure 17: Offset of each method from the single-layer RRE of the "Standard" method from the previous chart, ignoring the final layer. Values are clearly clustered around the mean difference for their respective series (dashed lines), with standard deviation of about 0.006, or less than 2% of the mean standard RRE (see Table 10).

Table 10: Offset stats.

Method	Offset mean	Offset std	Offset std (relative)
Gemma Scope	0.0137297	0.00284853	0.00837836
Standard + KL fine-tuning	0.0156311	0.00238957	0.00702843
Standard with LISTA	0.0416262	0.0059518	0.017506
Standard with LISTA + KL fine-tuning	0.055298	0.00645358	0.0189819
Replacement-aware	0.0391248	0.00474276	0.0139498
Replacement-aware + KL fine-tuning	0.0428781	0.00441566	0.0129878
Replacement-aware with LISTA	0.0514623	0.00325095	0.009562
Replacement-aware with LISTA + KL fine-tuning	0.0573609	0.00301493	0.00886779

Table 10: Offset stats.

To account for the second-order effect, a naive model would be to assume that the additional error is additive and independent. This would imply growth proportional to the square root of depth, since we're in high enough dimensions that randomly selected error directions will be orthogonal to each other. That is, if is the error added at layer , then the magnitude of the total error under this model for the layer would be given by

While this is almost certainly overly simplistic, if we plot the excess error in the full replacement model over the corresponding single-SAE error along with a best-fit square root, we can see that it fits remarkably well, at least once we get about a third of the way through the model:

Figure 18: Excess RRE at each layer in the full replacement model, defined as , where the full-replacement values are calculated on the encoder tuned model and the single-layer-replacement values are calculated on the corresponding pre-tuned SAEs. Dashed lines represent the best fit of a function .

Table 11: Excess RRE by layer.

Layer	Gemma Scope	Standard	Standard + KL fine-tuning	Standard with LISTA	Standard with LISTA + KL fine-tuning	Replacement-aware	Replacement-aware + KL fine-tuning	Replacement-aware LISTA	Replacement-aware LISTA + KL fine-tuning
0	-0.00242255	0.000983289	-0.000223161	-0.000218571	7.66213e-05	0.00105875	0.000264328	-0.000402221	-0.000749353
1	0.0178952	0.0149172	0.0140644	0.0194792	0.017919	0.0105392	0.0103037	0.0133132	0.0136862
2	0.0324831	0.030531	0.0283181	0.0380928	0.0333498	0.0206257	0.0203133	0.0261448	0.0261439
3	0.0472641	0.0430369	0.0402385	0.0506506	0.0442795	0.0288111	0.028704	0.0331605	0.0336458
4	0.0666221	0.059897	0.0555842	0.067598	0.057513	0.0406585	0.0401989	0.0455684	0.0452023
5	0.0652779	0.0616114	0.0570409	0.0697489	0.0588096	0.041644	0.0410515	0.0452413	0.0447312
6	0.0935088	0.0843675	0.078883	0.09196	0.0776041	0.0591979	0.0584546	0.0615519	0.0612641
7	0.105015	0.0928344	0.0864981	0.0985934	0.0826806	0.0664898	0.0653264	0.0659528	0.0644095
8	0.129427	0.110851	0.103548	0.1117	0.0938798	0.0807505	0.0780954	0.0746422	0.0727049
9	0.140789	0.119953	0.111692	0.115038	0.0983438	0.087522	0.0842835	0.0775538	0.0757272
10	0.15179	0.135156	0.126603	0.120764	0.103219	0.0960675	0.0904931	0.0817296	0.0795933
11	0.161751	0.142461	0.133334	0.124102	0.106929	0.104151	0.0955343	0.085278	0.0832042
12	0.178247	0.167395	0.159146	0.130977	0.114128	0.118621	0.104559	0.0907316	0.0887763
13	0.177647	0.163677	0.151429	0.132269	0.114196	0.121055	0.104387	0.0921847	0.0895181
14	0.189193	0.167592	0.156875	0.137273	0.119303	0.13212	0.110393	0.0972689	0.0944938
15	0.206009	0.180362	0.166561	0.148928	0.127937	0.145298	0.122023	0.106646	0.103165
16	0.206829	0.179436	0.164804	0.146839	0.125842	0.145339	0.12236	0.10537	0.101492
17	0.223131	0.191966	0.177548	0.152906	0.131652	0.157322	0.13149	0.110882	0.106391
18	0.241619	0.209326	0.192862	0.165275	0.140935	0.167868	0.142789	0.120012	0.115183
19	0.248978	0.215935	0.198807	0.170163	0.143871	0.171785	0.146301	0.123601	0.117702
20	0.248221	0.215025	0.197736	0.168686	0.142405	0.167281	0.144189	0.122181	0.116654
21	0.254549	0.220201	0.202429	0.173422	0.146243	0.169811	0.146892	0.126445	0.119295
22	0.260715	0.221268	0.20113	0.169464	0.141778	0.166773	0.143996	0.124868	0.118422
23	0.259375	0.217415	0.196263	0.163751	0.136195	0.161319	0.138757	0.122838	0.11494
24	0.267597	0.219931	0.197042	0.165925	0.136735	0.16101	0.138606	0.127767	0.122254
25	0.321842	0.251133	0.224448	0.194088	0.160472	0.161595	0.143484	0.152207	0.140878

Table 11: Excess RRE by layer.

Table 12: Square-root fit parameters.

Method	sqrt scale	sqrt offset	MSE of sqrt fit
Gemma Scope	0.0670694	-0.0557713	0.00702397
Standard	0.0556639	-0.0403634	0.00411874
Standard + KL fine-tuning	0.0505582	-0.034878	0.00357311
Standard with LISTA	0.0392629	-0.00902911	0.00143824
Standard with LISTA + KL fine-tuning	0.0327684	-0.00576087	0.00124694
Replacement-aware	0.0425507	-0.033041	0.00445681
Replacement-aware + KL fine-tuning	0.0355778	-0.0228236	0.00211101
Replacement-aware LISTA	0.0301347	-0.0133773	0.000748628
Replacement-aware LISTA + KL fine-tuning	0.0282006	-0.0102384	0.000562596

Table 12: Square-root fit parameters.

I propose, then, that one way to understand the effectiveness of an SAE training method is to consider this limiting constant as a measure of how quickly re-encoding error compounds. To what extent it is worth trading off individual layer accuracy to decrease the bound depends on the depth of the model in question. This certainly makes it more challenging to validate potential improvements, since they may not even be measurable until you've already trained many SAEs with them, and toy models may be too shallow to even observe a difference! Empirically, as we saw in Figure 11, in Gemma-2-2B this trade-off only becomes beneficial after 9 layers, even for my best-performing method. I suspect that truly optimizing end-to-end replacement model performance will require a deeper understanding of how and why certain representations generalize better, such that one can iterate on the training setup without having to train a full replacement model as I needed to here.

Visualizing the error bound versus the single-layer RRE in a Pareto frontier-ish ^[13] way shows that replacement-awareness, KL fine-tuning, and LISTA all push in the direction of making more of the trade-off of individual layer accuracy for slower error compounding:

Figure 19: Excess error bound, i.e. the limiting constant from Equation 16, against single-layer RRE.

Table 13: Error bound vs single-layer RRE.

Method	RRE	Error bound
Gemma Scope	0.35245	0.0670694
Standard	0.339987	0.0556639
Standard + KL fine-tuning	0.355871	0.0505582
Standard with LISTA	0.381038	0.0392629
Standard with LISTA + KL fine-tuning	0.395433	0.0327684
Replacement-aware	0.381334	0.0425507
Replacement-aware + KL fine-tuning	0.385008	0.0355778
Replacement-aware LISTA	0.393836	0.0301347
Replacement-aware LISTA + KL fine-tuning	0.399214	0.0282006

Table 13: Error bound vs single-layer RRE.

Encoder tuning is effective

Encoder tuning (Encoder tuning) appreciably improves replacement model performance. While the impact is greatest on replacement-naive SAEs, this effect does stack with the improvements from replacement-aware training:

Figure 20: Replacement model KL before and after encoder tuning. Tuning took place over 1 million training set tokens, which represents just under 2% and 0.025% of the total training budget for my SAEs and for Gemma Scope, respectively. All SAEs were tuned with a target sparsity of . In the case of Gemma Scope, this was achieved by adding a global offset to the JumpReLU thresholds. See Gemma Scope encoder tuning for more details. See the end of this section for a table of these values.

It's worth exploring how such a short fine-tuning stage can have such a large impact. One potential cause is the fact that error induced by the replacement model in early layers has a big effect on the statistics of the activations reaching later layers. The SAEs used here are especially sensitive to this due to the threshold parameters in their activation functions. Although replacement-aware training attempts to minimize such distortions, each SAE was still trained with inputs coming entirely from the base model distribution. One major function of encoder tuning is to adapt these thresholds to the new statistics, which we can see is effective by comparing feature stats before and after this phase:

Figure 21: Mean feature on 1 million validation set tokens by layer before and after encoder tuning with a target sparsity of . Note that in the pre-tuned replacement models there are two separate failure modes depending on the training method: exploding or vanishing , which are driven by changes in the mean activation magnitudes. Encoder tuning fails to exactly hit the target for all layers, but the result is generally quite close, within about 5% for most methods and layers.

Does this fully explain the effectiveness of encoder tuning, or is there additional impact of fine-tuning the other encoder parameters? Because these SAEs use the BatchTopK activation function ^[14] , we can check this directly by comparing performance for the tuned variants versus the pre-tuned versions with their adaptive "training mode" thresholds. Evidently, there is some additional effect of tuning the full set of parameters, though this ranges from a quite large effect for Gemma Scope to a negligible one for the LISTA and replacement-aware SAEs:

Figure 22: Replacement model KL with no tuning, adaptive thresholding, and encoder tuning. There is a reasonably large effect of merely adapting the thresholds, but encoder tuning does offer additional improvement.

This isn't definitive, but these results could be evidence that the LISTA/replacement-aware SAEs were already mostly adapted to the replacement model, modulo the activation statistics. In contrast, the standard SAEs had more room for improvement and so fine-tuning the encoder parameters had a much greater impact. It might be worth investigating encoder tuning as a standalone method, since it is applicable independently of any initial training.

Table of KL values for charts in this section:

Method	Pre-tuned KL	Adaptive KL	Post-tuned KL
Gemma scope	11.9297	8.72424	1.66864
Standard	9.30826	12.0985	1.00631
Standard + KL tuning	11.9398	12.0227	0.78684
Standard LISTA	2.58798	1.17143	0.876545
Standard LISTA + KL tuning	1.73374	0.727763	0.54007
Replacement-aware	4.63118	1.18139	0.53557
Replacement-aware + KL tuning	4.63118	1.18139	0.53557
Replacement-aware LISTA	2.30692	0.661828	0.578471
Replacement-aware LISTA + KL tuning	2.21987	0.693616	0.448847

Table 14: KL values for the charts in this section.

KL fine-tuning doesn't degrade reconstruction quality

Since I use KL fine-tuning extensively, it's worth quickly sanity checking that it doesn't degrade reconstruction quality. If we pair methods by whether or not they included this phase, we see that they are indeed very close:

Standard, vanilla encoder

Figure 23

Replacement aware, vanilla encoder

Figure 24

Standard, LISTA

Figure 25

Replacement aware, LISTA

Figure 26

LISTA encoders improve robustness to upstream error

We have already seen that using the LISTA encoder improves end-to-end performance. Naively, one might expect that this would be due to improved reconstruction fidelity across the board. However, if we revisit the plots of single-layer replacement model reconstruction errors, we actually see reduced performance for the LISTA SAEs when we pair them with their vanilla encoder equivalents:

Standard training

Figure 27

Standard training + KL fine-tuning

Figure 28

Replacement-aware training

Figure 29

Replacement-aware training + KL fine-tuning

Figure 30

The situation is much different when we make the same comparison within the full replacement model. Much as when we compared replacement-aware SAEs to replacement-naive ones, the lines cross over in the middle layers:

Standard training

Figure 31

Standard training + KL fine-tuning

Figure 32

Replacement-aware training

Figure 33

Replacement-aware training + KL fine-tuning

Figure 34

Thus, it appears that the mechanism by which the LISTA encoder improves end-to-end performance is its better adaptation to the full replacement model. In other words, the LISTA encoder makes more of the trade-off discussed in A simple model of reconstruction error, which is apparent if we recreate Figure 19 with lines connecting methods differing in LISTA-ness:

Figure 35: Replacement models with the LISTA encoder trade off single layer reconstruction accuracy for more slowly compounding error.

Discussion

In this post, I've presented several methods for improving the end-to-end performance of SAE replacement models. More so than any specific innovation, I want the take-away from this to be that more robust sparse representations are not only possible, but relatively unexplored. The LISTA encoder and replacement-aware training were literally the first things I tried once I became convinced of the inadequacy of the error node approach. Single-layer reconstruction error is simply not the correct optimization target if we want to explain how feature interactions determine model behavior. Taking a holistic view of the entire replacement model is an angle I expect to continue to be fruitful. In particular, I think directly tackling the problem of compounding reconstruction error (A simple model of reconstruction error), as I have attempted to do here, is a neglected and important area of future research.

The main technical drawbacks of replacement-aware training are modestly increased computational and memory costs, forced sequential ordering of training runs, and the fact that it is not compatible with activation shuffling. ^[15] The sequential ordering may not be too much of a burden in practice, since I found that warm-starting the layer SAE from the layer weights results in a moderate boost in performance for a given token budget (Warm-starting from the next layer SAE modestly boosts performance), which could potentially be instead exchanged for shorter training runs on the earlier layers.

Replacement-aware training could be thought of as a variant of end-to-end training (Braun et al. 2024) adapted to the full replacement model. I previously found that end-to-end trained SAE replacement models outperform standardly trained ones on TinyStories. Even so, the much less computationally expensive next-layer replacement-aware SAEs outperformed those. I would guess that downstream reconstruction error and feature-space distortion are both capturing a similar signal---something like "sufficiency for the downstream layer to do its job"---but that the latter is more precisely targeted to our actual replacement model. I suspect that the explicit dependence on previously trained downstream SAEs is not strictly necessary. There might be something clever you can do based only on the statistics of the base model's next layer activations, for example with something like information dropout (Achille and Soatto 2018).

Applying LISTA as the SAE encoder is a similar move to other recent works that have explored encoder variants, such as Matryoshka SAEs (Bussmann et al. 2025) and matching pursuit SAEs (Costa et al. 2025). While LISTA has been combined with Variational Auto-encoders (VAEs) in vision (Xiao et al. 2023), to my knowledge I'm the first to try it in this setting. Interestingly, I did not find improved performance of LISTA SAEs according to standard metrics (LISTA encoders improve robustness to upstream error); their benefit appears to be realized only in the context of the full replacement model. It might be worth re-analyzing previous mediocre or negative results for variant SAEs (Baker and Li (2025) comes to mind; their framing, not mine!) through this lens. Another obvious angle for follow-up work would be to try out different numbers of iterations and sparsity schedules. I explored several variants of LISTA that didn't make it into this post due to their lack of training stability. ^[16] I think my implementation has a lot of room for improvement, particularly around finding a good parameter regularization scheme to ensure feature activation levels stay in a sane range.

Finally, encoder tuning (Encoder tuning) may be worth investigating in its own right, possibly even as an alternative to replacement-aware training. It would definitely have to be refined somehow; I found that performance actually started decreasing with longer training runs, possibly resulting from the influence of the outliers mentioned in Replacement-aware training does not sacrifice faithfulness. ^[17] A more stable version could potentially be used as an inexpensive way to generate more robust features out of off-the-shelf SAEs.

Acknowledgements

I'd like to thank Curt Tigges for his mentorship during SPAR and for practical advice on proving out an early version of the LISTA-inspired encoder, and Billy Martin for helpful conversations and providing some of the GPU time used in these experiments.

References

Achille, Alessandro, and Stefano Soatto. 2018. "Information Dropout: Learning optimal representations through noisy computation." IEEE Transactions on Pattern Analysis and Machine Intelligence 40 (12): 2897–905.

Ameisen, Emmanuel, Jack Lindsey, Adam Pearce, et al. 2025. Circuit tracing: Revealing computational graphs in language models. https://transformer-circuits.pub/2025/attribution-graphs/methods.html.

Baker, Zachary, and Yuxiao Li. 2025. "Analysis of Variational Sparse Autoencoders." arXiv [cs.LG], September.

Bloom, Joseph, Curt Tigges, Anthony Duong, and David Chanin. 2024. SAELens. https://github.com/decoderesearch/SAELens.

Borgerding, Mark, and Philip Schniter. 2016. "Onsager-corrected deep learning for sparse linear inverse problems." arXiv [cs.IT], July.

Braun, Dan, Jordan Taylor, Nicholas Goldowsky-Dill, and Lee Sharkey. 2024. "Identifying functionally important features with end-to-end sparse dictionary learning." arXiv [cs.LG], May.

Bussmann, Bart, Patrick Leask, and Neel Nanda. 2024. "BatchTopK Sparse Autoencoders." arXiv [cs.LG], December.

Bussmann, Bart, Noa Nabeshima, Adam Karvonen, and Neel Nanda. 2025. "Learning multi-level features with Matryoshka sparse autoencoders." arXiv [cs.LG], March.

Caples, Diego, Jatin Nainani, Collum McDougall, and Rob Neuhaus. 2025. Scaling sparse feature circuit finding to Gemma 9B. https://www.lesswrong.com/posts/PkeB4TLxgaNnSmddg/scaling-sparse-feature-circuit-finding-to-gemma-9b.

Chen, Xiaohan, Jialin Liu, Zhangyang Wang, and Wotao Yin. 2018. "Theoretical linear convergence of unfolded ISTA and its practical weights and thresholds." arXiv [cs.LG], August.

Clark, Peter, Isaac Cowhey, Oren Etzioni, et al. 2018. "Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge." arXiv [cs.AI], March.

Costa, Valérie, Thomas Fel, Ekdeep Singh Lubana, Bahareh Tolooshams, and Demba Ba. 2025. "From flat to hierarchical: Extracting sparse representations with matching pursuit." arXiv [cs.LG], June.

Cunningham, Hoagy, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. 2023. "Sparse autoencoders find highly interpretable features in language models." arXiv [cs.LG], September.

Daubechies, I, M Defrise, and C De Mol. 2004. "An iterative thresholding algorithm for linear inverse problems with a sparsity constraint." Communications on Pure and Applied Mathematics 57 (11): 1413–57.

Elhage, Nelson, Tristan Hume, Catherine Olsson, et al. 2022. Toy Models of Superposition. https://transformer-circuits.pub/2022/toy_model/index.html.

Gao, Hong-Ye, and Andrew G Bruce. 1997. "Waveshrink with firm shrinkage." Statistica Sinica, 855–74.

Gao, Leo, Stella Biderman, Sid Black, et al. 2020. "The Pile: An 800GB dataset of diverse text for language modeling." arXiv [cs.CL], December.

Ge, Xuyang, Fukang Zhu, Wentao Shu, Junxuan Wang, Zhengfu He, and Xipeng Qiu. 2024. "Automatically identifying local and global circuits with linear computation graphs." arXiv [cs.LG], May.

Gemma Team, Morgane Riviere, Shreya Pathak, et al. 2024. "Gemma 2: Improving open language models at a practical size." arXiv [cs.CL], July.

Gorton, Liv. 2024. "The missing curve detectors of InceptionV1: Applying sparse autoencoders to InceptionV1 early vision." arXiv [cs.LG], June.

Gregor, Karol, and Yann LeCun. 2010. "Learning fast approximations of sparse coding." International Conference on Machine Learning, June, 399–406.

Hanna, Michael, Mateusz Piotrowski, Jack Lindsey, and Emmanuel Ameisen. 2025. "Circuit-tracer: A new library for finding feature circuits." Proceedings of the 8th BlackboxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP (Stroudsburg, PA, USA), November, 239–49.

Heindrich, Lovis, Philip Torr, Fazl Barez, and Veronika Thost. 2025. "Do sparse autoencoders generalize? A case study of answerability." arXiv [cs.LG], February.

Hendrycks, Dan, Collin Burns, Steven Basart, et al. 2020. "Measuring massive multitask language understanding." arXiv [cs.CY], September.

Karvonen, Adam. 2025. "Revisiting end-to-end sparse autoencoder training: A short finetune is all you need." arXiv [cs.LG], March.

Kissane, Connor, Neel Nanda, and Arthur Conmy. 2024. "SAEs are highly dataset dependent: a case study on the refusal direction." AI Alignment Forum, June.

Lieberum, Tom, Senthooran Rajamanoharan, Arthur Conmy, et al. 2024. "Gemma Scope: Open sparse autoencoders everywhere all at once on Gemma 2." arXiv [cs.LG], August.

Lindsey, Jack, Adly Templeton, Marcus Jonathan, Thomas Conerly, Joshua Batson, and Christopher Olah. 2024. Sparse crosscoders for cross-layer features and model diffing. https://transformer-circuits.pub/2024/crosscoders/index.html.

Liu, Jialin, Xiahan Chen, Zhangyang Wang, and Wotao Yin. 2019. "ALISTA: Analytic Weights Are As Good As Learned Weights in LISTA."

Mallat, S G, and Zhifeng Zhang. 1993. "Matching pursuits with time-frequency dictionaries." IEEE Transactions on Signal Processing 41 (12): 3397–415.

Marks, Samuel, Can Rager, Eric J Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. 2024. "Sparse feature circuits: Discovering and editing interpretable causal graphs in language models." arXiv [cs.LG], March.

Rajamanoharan, Senthooran, Tom Lieberum, Nicolas Sonnerat, et al. 2024. "Jumping ahead: Improving reconstruction fidelity with JumpReLU sparse autoencoders." arXiv [cs.LG], July.

Talmor, Alon, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. 2018. "CommonsenseQA: A question answering challenge targeting commonsense knowledge." arXiv [cs.CL], November.

Xiao, Pan, Peijie Qiu, Sungmin Ha, Abdalla Bani, Shuang Zhou, and Aristeidis Sotiras. 2023. "SC-VAE: Sparse Coding-based Variational Autoencoder with Learned ISTA." arXiv [cs.CV], March.

Yang, Jingcheng, Tianhu Xiong, Shengyi Qian, Klara Nahrstedt, and Mingyuan Wu. 2026. "Circuit tracing in vision-language models: Understanding the internal mechanisms of multimodal thinking." arXiv [cs.CV], February.

Zheng, Chujie, Hao Zhou, Fandong Meng, Jie Zhou, and Minlie Huang. 2023. "Large language models are not robust multiple choice selectors." arXiv [cs.CL], September.

Appendix A: Implementation details

Training setup

I implemented the SAEs and SAE training pipeline used in this post without dependence on existing SAE libraries. I use SAELens (Bloom et al. 2024) only to load the weights for GemmaScope, which is then converted into my module format for evaluation consistent with the rest of my analysis. All source code is available on GitHub. The base model, Gemma-2-2B, is run using the huggingface transformers library. To implement replacement models, my code wraps a bare transformers model by replacing each transformer block with a module that first calls the original block, then replaces its output with the output of an attached SAE. For better performance, I include logic that stops the transformer computation as early as possible to obtain the needed activations.

The training pipeline uses the huggingface datasets library to fetch training data. All SAEs were trained using 50-100 million tokens from the pile-uncopyrighted (Gao et al. 2020) training split, as indicated in Table 3. KL fine-tuning and encoder tuning also use this dataset. All metrics presented in this post were run on 1 million tokens from the validation split of the same dataset. For efficiency, a consistent context length of 1024 tokens (as used in Gemma Scope) was used. Longer examples were truncated to that length, while smaller examples were concatenated greedily to attempt to fill each batch with as few padding tokens as possible. This was accomplished by looking ahead by up to 256 examples to fill in any gaps on a best-effort basis. Appropriate attention masks were set to account for multiple examples per batch row. Additionally, for each example I define a token mask to filter out padding and any other special tokens; my SAEs pass through the underlying transformer block's output for these positions, and they are skipped when calculating loss and validation metrics.

All optimization is done using the PyTorch implementation of the Adam optimizer without weight decay. Following Gemma Scope (Lieberum et al. 2024), I use beta=[0, 0.999]. I found (by informal testing) that a learning rate of 1e-4 was effective, though I performed no formal sweep. BatchTopK thresholds were learned independently of the main optimizer using an exponential moving average with learning rate 1e-2. Following Karvonen (2025), I adaptively set the relative weights of each term in the loss functions such that they would have the same magnitude on a per-batch basis. I used a batch size of 2 complete context windows, for a total of 2048 tokens. When computing KL divergence loss terms, I aggregate using geometric mean rather than arithmetic mean, as I found this performed better in my earlier experiments with TinyStories.

All SAEs used a width of 16384 (approximately 7x the hidden dimension of 2304) to match with Gemma Scope. SAEs used either the BatchTopK (Equation 12) or LISTA (Equation 15) encoders, with sparsity parameters set such that the final encoder output would have 100 expected non-zero entries. For BatchTopK this is simply ; for LISTA, which I exclusively ran with 5 iterations, I used an increasing schedule after each iteration of [20, 40, 60, 80, 100].

For LISTA SAEs, I constrain the columns of the decoder to have unit norms by dividing them by their magnitudes after each training step. Additionally, the encoder weights are parametrized via PyTorch's spectral_norm, and the per-iteration scalar values are soft-capped to 1 with tanh.

For memory and compute efficiency, the base model weights were cast to bfloat16, and PyTorch's autocast feature is used wherever possible. SAE weights were trained in float32, but cast to bfloat16 once they were no longer needed in the replacement model used during training. Checkpoints are of course saved in full precision. All validations and benchmarks were done using the lower-precision weights and autocasting.

Encoder tuning (Encoder tuning) was done over 1 million tokens from the same training set used in base training. The decoder weights are frozen for the entirety of this step. I found it most effective to apply Equation 11 to pairs of layers; the encoders for both SAE and are included in the optimizer parameters. To mitigate the outlier token problem (Replacement-aware training does not sacrifice faithfulness) and improve training stability, I apply a tanh softcap to the replacement model features of 10 times the maximum feature magnitude of the target baseline features.

Benchmark setup

To run the benchmarks reported in this post, I used the linked huggingface datasets MMLU (Hendrycks et al. 2020), CommonsenseQA (CQA) (Talmor et al. 2018), and the ARC Easy set (ARC-E) (Clark et al. 2018). For all benchmarks I use the n-shot template (line breaks added for clarity; exact spacing is as indicated with newline characters \n)

Question: {question}\n\n
A. {option[0]}\n
B. {option[1]}\n
C. {option[2]}\n
D. {option[3]}\n
E. {option[4]}\n
Answer: {answer_label}\n\n
Question: ...

Following the n-shot prefix, for the question being evaluated, the same template is filled in right until the colon after "Answer", and the top logit for the predicted next token is taken as the model's response.

For the "Valid answer" metric reported in various plots, I use the total proportion of times that the top-logit response corresponded to a valid option. Note that MMLU and ARC-E have only 4 valid options; option "E" is omitted from the n-shot template as well.

To implement PriDe (Zheng et al. 2023) (used only where specified), for each benchmark I take the first 5% of questions (including the n-shot prefix) as a debiasing sample. For each question in the sample, I copy the example with permuted options corresponding to every cyclic permutation (i.e. A, B, C, D -> B, C, D, A -> C, D, A, B -> D, A, B, C), which guarantees that every option appears once with each label. I then compute the mean probability (conditional on valid answer) assigned to each label over this sample as the prior probability for that label. Where I report "debiased" answers, these are the highest-probability valid option after dividing each label's probability (again, conditioned on valid answer) by the corresponding prior. All of these probability calculations are performed in log-space (with, e.g. logsumexp operations) for better numerical precision.

Benchmark-specific details are highlighted below.

MMLU

For each MMLU subsection I use the corresponding "dev" split to generate a 5-shot prefix for each question. I evaluate over all examples from the "test" split of every subsection, excluding ones that exceed the maximum context length. This excludes the high_school_european_history, high_school_us_history, high_school_world_history, professional_law, and security_studies subsections. The resulting filtered set includes 11422 questions from 52 subsections.

CQA

CQA benchmark results are reported over the "validation" split. I use the first 10 examples to generate the 10-shot prompt, resulting in a total of 1211 questions. Note that I use the "train" split when performing CQA fine-tuning (CQA fine-tuning) so that results are not contaminated when evaluating CQA-finetuned replacement models.

ARC-E

I use the "test" split of the "Arc-Easy" subset of ARC. As in CQA, I take the first 10 examples to generate the 10-shot prompt. This dataset includes a small number of questions with a variable number of options, so to simplify analysis I discard all questions that do not have exactly 4 options. This leaves a total of 2355 questions.

CQA fine-tuning

For results that are marked as "CQA fine-tuned," I repeat the appropriate version of KL fine-tuning (KL fine-tuning) with the standard SAE training dataset interleaved with questions and answers from the CQA "train" split in a 90%-10% mixture. Each CQA sample is formatted exactly as an n-shot prefix would be (Benchmark setup), with n drawn uniformly from the range (5, 25), the idea being that I didn't want the question-answering-specific features to be overfit to specific token positions. The CQA dataset is allowed to repeat to guarantee that enough tokens can be drawn from it to reach the target fraction.

Gemma Scope encoder tuning

The Gemma Scope SAEs are implemented with the JumpReLU activation function, which is very similar to the inference-time BatchTopK except that the threshold varies for each feature. To adapt these SAEs into my framework, I add an additional parameter representing a global offset to the activation threshold, while keeping the base thresholds fixed. This avoids the need to set up the notoriously finicky kernel density estimation needed to tune the full set of thresholds. I learn this threshold offset in exactly the same way as I do for the BatchTopK SAEs, by exponential moving average. As can be seen in Figure 21, this is successful in adapting these SAEs to the desired sparsity of k=100. To make sure that this was as fair a comparison as possible, I also checked that tuning these SAEs to their "canonical" values (which varied by layer, but were generally slightly less than 100), to which they were plausibly better adapted, did not improve performance. Indeed, the SAEs perform slightly better when tuned to =100 at all layers:

Figure 36

Appendix B: Miscellaneous analyses

Warm-starting from the next layer SAE modestly boosts performance

Replacement-aware training imposes a serial ordering on SAE training, since it requires that later layers already have their weights fixed. Given this, I thought it natural to warm start each SAE with the weights from the next layer, since it seems likely that adjacent layers will use similar representations. This may also encourage SAEs to retain features that are useful in multiple layers, since they will no longer have to discover them from scratch. As, to my knowledge, this is a novel initialization strategy for SAE training, I thought it would be worth a very brief examination of how effective it is. Below, I plot the single-layer RRE for standard training with or without this warm start:

Figure 37: Two training runs using standard SAE training, one with warm start ("Standard") and one without ("Standard, Fresh Init"), both over 5e7 tokens. Gemma Scope (4e9 tokens), with JumpReLU thresholds adjusted to enforce equivalent sparsity (k=100) is plotted for comparison.

Layer	Gemma Scope	Standard	Standard, Fresh Init	Diff	Relative Diff
0	0.19932	0.180271	0.213263	-0.0329925	-0.154703
1	0.252744	0.239315	0.271258	-0.031943	-0.117759
2	0.300137	0.28373	0.318279	-0.0345488	-0.108549
3	0.318368	0.304365	0.333042	-0.0286764	-0.0861047
4	0.319223	0.304775	0.334593	-0.0298186	-0.089119
5	0.346149	0.329507	0.356045	-0.0265382	-0.0745362
6	0.344983	0.329273	0.355097	-0.0258237	-0.0727231
7	0.345515	0.327804	0.353016	-0.025212	-0.071419
8	0.355005	0.337783	0.36325	-0.0254666	-0.0701075
9	0.368929	0.351463	0.375667	-0.0242041	-0.0644297
10	0.381003	0.364561	0.387954	-0.0233939	-0.0603006
11	0.38254	0.368601	0.38963	-0.0210293	-0.0539725
12	0.392577	0.380148	0.400244	-0.0200962	-0.0502099
13	0.381505	0.371095	0.388986	-0.0178911	-0.0459942
14	0.380856	0.372011	0.388194	-0.0161831	-0.0416881
15	0.359406	0.349193	0.364015	-0.0148221	-0.0407183
16	0.367982	0.356994	0.372763	-0.0157693	-0.0423038
17	0.358951	0.348789	0.361992	-0.0132034	-0.0364743
18	0.355288	0.3435	0.357843	-0.0143428	-0.0400814
19	0.362968	0.349597	0.364034	-0.0144369	-0.0396582
20	0.366732	0.352347	0.367754	-0.0154073	-0.0418956
21	0.368278	0.354728	0.370463	-0.015735	-0.042474
22	0.385937	0.374559	0.39109	-0.0165312	-0.0422695
23	0.401697	0.392172	0.408789	-0.0166171	-0.0406494
24	0.393424	0.387725	0.401764	-0.014039	-0.0349435
25	0.374176	0.385348	0.385316	3.15309e-05	8.18311e-05

Table 15: Warm start vs fresh init single-layer RRE.

This reveals a modest improvement for warm-starting, with a mean difference in RRE of 0.0206, or a mean relative difference of about 6%; the gap increases towards the start of the model (15% for the first layer, only 3.5% for the penultimate). For context, we can see in the chart that Gemma Scope (constrained to equivalent sparsity, see Gemma Scope encoder tuning), which was trained with much more data (4e9 tokens) but from a different distribution, lies somewhere in between the two. This is of course too crude an analysis to conclude much, but it does suggest that the warm start is training SAEs with a greater "effective number of training tokens." Thus, for a given compute budget, an optimal training schedule using this initialization scheme could allocate shorter training runs for earlier layers. This may not be terribly exciting in practice, since earlier layers are already cheaper to run due to early stopping.

5e7 tokens is probably enough

As a sanity check that the single longer training run isn't unfairly skewing the results, as well as a cursory investigation into what the training dynamics of replacement-aware training might look like, I compare validation statistics from the full 1e8 token next_layer training run with those from its 5e7 token checkpoints. Note that due to the warm start initialization (Warm-starting from the next layer SAE modestly boosts performance), this isn't quite identical to what a true 5e7 token training run would look like; the layer SAE was initialized from the layer SAE at the 1e8 token checkpoint, so even the earlier checkpoints for each layer have effectively seen more tokens. Still, there are interesting differences between the two replacement models:

Figure 38: The earlier 5e7 token checkpoint actually gets better RRE in the single-layer condition.

While the shorter run gets better RRE in the single-layer condition, the longer training run is marginally better in the full replacement condition, though only just:

Figure 39

Repeating the analysis from A simple model of reconstruction error on these two replacement models seems to show that the longer training effectively increased the error-compounding-for-single-layer-accuracy trade-off:

Figure 40

Figure 41

Yet, since the actual replacement model RRE barely improved, this could be an indication that the trade-off is sub-optimal; all else being equal, I think it's reasonable to want the single-layer-replacement reconstruction to be as good as possible. It should be possible to tune this by changing the ratio of the current-layer and next-layer terms of Equation 7, which might be worth following up on.

Appendix C: Against error nodes; or, what's the point of a feature circuit, anyway?

What's the point of a feature circuit?

Aside from the concerns outlined in Error nodes, I have a more fundamental issue with using error nodes in circuit discovery: the mechanisms this approach is capable of identifying are not expressible in terms of features alone, and that's actually a major problem! This lack of representational sufficiency limits the validity of the circuit to the locality of a specific input. While this is fine when it comes to observational data, once we start trying to predict the results of feature-space interventions (e.g., steering), we are forced to step outside of this region of validity. The result is that any interventional predictions we try to make are subject to the same compounding error that the error nodes were meant to avoid. To explain, allow me to first sketch out my mental model of what we're trying to accomplish with feature circuits in the first place.

If we take the superposition hypothesis (Elhage et al. 2022) seriously, then the activations we observe in a language model are in fact projections from a much higher dimensional space. Under this view, there is some unobserved "true" set of representations that evolve as computation progresses through the network ^[18] ; the model has learned a compressed approximation of these representations. We can draw this schematically as

Figure 42: Superposition model of neural activations in a three-layer network. The unobserved "true" features have been implicitly learned by the model through their projections . The right-pointing arrows represent the functional relationship between features across layers, such that

Note that since by hypothesis the are higher dimensional than the activations , the operations represented by the downward arrows are not invertible; there are in general many possible configurations of features that could result in the observed activations. However, if we assume that the vector representations of the features are sparse, we can nevertheless infer them through the combination of sparse dictionary learning and sparse coding. This is the standard case for training SAEs as an interpretability tool.

I think that we should be more ambitious than merely identifying the features that are relevant to the model's computation. In vision models, SAEs learn features corresponding to simple shapes like edges and textures in early layers, gradually converging on recognizable objects in the later layers (see, e.g., Gorton (2024)); there is a clear logic to this progression, and in principle it could be reverse-engineered into an interpretable algorithm. This should also be possible for language models! In my view, the point of a feature circuit is to use the inferred features to identify the feature-space mechanisms I have labeled . As I will now show, a full replacement model is a proxy on which we can validate arbitrary hypotheses about , including causal ones, but this property is destroyed by the addition of error nodes.

Replacement model as proxy

Recall the schematic depiction of the full replacement model from Figure 3. Because all of the labeled edges are deterministic, we can validly collapse some of them by incorporating the transformation into the node. In particular, we can use the fact that to remove the pre-encoded activation nodes. Furthermore, we can "lift" the right-pointing arrows to the feature level if we sacrifice causal equivalence at the activation level, using the fact that .

Figure 43: Replacement model network, manipulated to match the superposition model by collapsing edges into nodes and duplicating the post-SAE activations as leaf nodes.

This model is obviously observationally equivalent to the replacement model on the remaining nodes, because all we have done is move some computations around without changing their values. It is not causally equivalent with respect to interventions on the activations since these are now leaf nodes, but it is causally equivalent with respect to interventions on the features. Since we are only interested in feature-level mechanisms, this is totally acceptable.

The point of these manipulations is of course that Figure 43 now matches the structure of Figure 42. The only difference is that the observed activations are now the replacement model activations instead of the base model activations . Thus, this model is faithful to the original to the extent that these activations match.

To be clear, I am not claiming that the true mechanisms of the base model are literally "the composition of SAE encoder, transformer block, and SAE decoder," which would be just as uninterpretable as when we started. What I am claiming is that this model is a proxy with which we can search for and validate simpler, actually interpretable functional relationships between the features. That is, by construction, when we make the intervention in the replacement model, ablating the base model at layer with the predicted activations results in exactly the features at layer as would be computed by the proxy, regardless of the original input .

Against error nodes

What happens when we add in error nodes that cancel out the SAE reconstruction error? We can employ similar manipulations as before to Figure 4 to get a graph supporting direct feature-to-feature interactions, though we still have an irreducible dependence on the base model activations:

Figure 44: Mapping the replacement model with error nodes to the superposition model, with the same trick of collapsing edges, duplicating replacement model activations as leaf nodes, and lifting edges to feature-level. The feature-level mechanisms are .

The values of the error nodes are given by

i.e. the values that exactly cancel the SAE reconstruction error. The features are given by

By construction, this model computes the single-layer-replacement features ; it is a useful proxy for purely observational data. However, consider the causal intervention when the original input is . To clean up the notation for the sake of this argument, let refer to the "natural" (pre-intervention) values that the corresponding more ornamented nodes would take in the graph for this input, and let refer to their post-intervention values. Additionally, let refer to the base model activations for the input . Then we have

Notably, the result is no longer equivalent to the encoder applied to the layer 1 base model activations; we have taken a step away from that activation in the direction defined by the difference between the counterfactual layer 0 features and the ones we actually observed on this input. And of course we have, because that's what a causal intervention is! However, the important thing to note here is that the result now contains two sequential uses of SAE encoders; the outer , and the inner , since . Note that I have underlined terms according to the number of sequential encoders they depend on. Proceeding to the next layer, we have

which now depends on three sequential encodings: applied to a value that depends on , which itself depended on two sequential encodings. Obviously, this pattern continues for any remaining layers in the model; the final layer features depend on sequential encodings, exactly as in the full replacement model.

I would like to emphasize that this dependence on serial reconstructions is not an accidental property of the way I have chosen to set up the causal graphs. The fundamental reason error nodes cannot help in this situation is that interventions on features must have downstream consequences or else be meaningless. And how else could we measure the downstream impact on a later layer's features, than by applying that layer's encoder? Thus, it seems to me that we are forced to contend with compounding reconstruction error; improving the accuracy of causal predictions depends on reducing it, not on some clever scheme to bypass it.

1. Each CLT writes directly to every CLT following it. For any intermediate layer between and , it must appear in exactly half of paths from to , so it contributes 1/2 to the average path length. There are intermediate layers, and an additional edge to must be included in every path, so the average path length is . ↩︎

2. Note that I have deliberately omitted a sparsity-promoting regularization term on the SAE features, as in this post I will be exclusively using SAE architectures that achieve sparsity by other means. ↩︎

3. The use of the KL term is admittedly questionable under this logic, at least for layers other than the final one! Its inclusion is purely pragmatic, as I did find it necessary in order to achieve the level of performance I was aiming for. I hope future work can find a more principled alternative. I report results both with and without the use of this term. ↩︎

4. I did attempt a version of this that tuned all encoders in parallel, but I got much better results training them serially as written. ↩︎

5. My best guess is that this is largely driven by adjusting the BatchTopK/JumpReLU thresholds to match the new statistics of the replacement model, as tuning those parameters in isolation also helps quite a bit, but I still saw additional gains when tuning the full set of encoder parameters. ↩︎

6. This simplified the code and gave me marginally better performance on the hardware I was using. ↩︎

7. I use an exponential moving average of the adaptive threshold in Equation 12, following the implementation in SAELens (Bloom et al. 2024). ↩︎

8. The original paper used the soft shrinkage function , with threshold parameters learned along with the update matrices and shared across iterations. ↩︎

9. i.e., they are constrained to be non-negative and vary by many orders of magnitude. I also previously found them to be approximately log-normally distributed. ↩︎

10. All reported results on CQA are over the validation split, so this won't contaminate them. ↩︎

11. I use this measure rather than the more common MSE because MSE varies by orders of magnitude across layers. The normalization factor makes these values directly comparable. ↩︎

12. The sudden jump for the final layer is likely due to using the KL term of Equation 7 for the entire training run, rather than just for the last 20% of tokens, so there is potentially some trade-off happening there that isn't present on the earlier layers. ↩︎

13. "ish", because the implicit objective here, the mean replacement model RRE, is actually one dimensional. ↩︎

14. Gemma Scope uses JumpReLU, but we can adapt it in exactly the same way by adding a global offset to the individual feature thresholds, which is also what I did during the encoder tuning phase for those SAEs. ↩︎

15. The replacement-aware loss function depends on having a complete example in context, since we have to run the next transformer layer on the SAE output. ↩︎

16. For example, I attempted a run with a 10-iteration encoder, only to have the activations vanish in the replacement model even after encoder tuning. Due to the expense of this kind of training (each LISTA iteration uses the same compute as a vanilla encoder!), I gave up on trying to debug this. ↩︎

17. I tried gradient clipping, using trimmed MSE as the loss function, and even outright dropping tokens generating extreme feature values, but none of these were successful. ↩︎

18. I focus here on the cross-layer dynamics because that feels most natural for the residual stream SAEs discussed in the post, but it should be possible to extend this to be explicitly across token positions by adding nodes for attention-level features. ↩︎