You don't need error nodes, you need better features

A new training method called replacement-aware training produces sparse auto-encoders (SAEs) that retain language capabilities when used in a full replacement model of Gemma-2-2B, unlike standard SAEs which cause catastrophic failure. The approach, developed by Evan Lloyd and detailed in a blog post, introduces a loss term that penalizes distortions to the next layer's features, resulting in a replacement model with significantly lower validation KL divergence (0.53557 vs 1.00631) and improved MMLU accuracy (0.193136 vs 0.0746804) compared to standard methods. The code and weights are publicly available.

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. 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 https://www.lesswrong.com/feed.xml 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 https://elloworld.net/posts/replacement-aware-sae-training/ 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 https://github.com/evan-lloyd/mechinterp-experiments/tree/gemma-blog-post/transformers sae . 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: 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 . | 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 | 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. 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 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, 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 https://transformer-circuits.pub/2025/attribution-graphs/methods.html limitations-reconstruction-error 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 https://www.lesswrong.com/posts/PkeB4TLxgaNnSmddg/scaling-sparse-feature-circuit-finding-to-gemma-9b 2 3 3 Error Nodes 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 https://www.lesswrong.com/feed.xml Appendix C Against error nodes or what s the point of a feature circuit anyway , since it is immaterial to my empirical results, but to briefly summarize: 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 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. 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 https://www.lesswrong.com/feed.xml fn-nd9E4j8WdoJ7ewJPj-2 where is a batch of inputs. In contrast, replacement-aware training https://elloworld.net/posts/replacement-aware-sae-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 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. 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 https://www.lesswrong.com/feed.xml 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. 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. 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 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. 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 https://www.lesswrong.com/feed.xml Appendix A Implementation details . Source code is available on GitHub https://github.com/evan-lloyd/mechinterp-experiments/tree/gemma-blog-post/transformers sae . Data, including all SAE weights, validation and benchmark results, and tables and plots included in this post are available via huggingface bucket https://huggingface.co/buckets/evan-lloyd/replacement-aware-saes ; 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 https://huggingface.co/datasets/monology/pile-uncopyrighted Gao et al. 2020 training split per layer. The core training pipeline consists of 3 stages: initial training Replacement-aware training https://www.lesswrong.com/feed.xml Replacement aware training , KL fine-tuning KL fine-tuning https://www.lesswrong.com/feed.xml KL fine tuning , and encoder tuning Encoder tuning https://www.lesswrong.com/feed.xml 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 | |---|---|---|---| | 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 https://www.lesswrong.com/feed.xml Encoder tuning to match their sparsity to my SAEs. See Gemma Scope encoder tuning https://www.lesswrong.com/feed.xml Appendix A Implementation details for further details. 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 https://huggingface.co/datasets/monology/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 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. | 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 https://www.lesswrong.com/feed.xml Appendix A Implementation details 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. | 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 . | 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 https://www.lesswrong.com/feed.xml 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 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. | 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. 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. | 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 | | 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 https://www.lesswrong.com/feed.xml fn-nd9E4j8WdoJ7ewJPj-12 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 . | 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 . | 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. | 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. | 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 Encoder tuning https://www.lesswrong.com/feed.xml 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. | 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. 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 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 https://www.lesswrong.com/feed.xml 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. 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 https://www.lesswrong.com/feed.xml 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 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 https://elloworld.net/posts/replacement-aware-sae-training/ results 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 https://www.lesswrong.com/feed.xml 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 https://www.lesswrong.com/feed.xml 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 https://www.lesswrong.com/feed.xml 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. 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. 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 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 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 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 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 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. I implemented the SAEs and SAE training pipeline used in this post without dependence on existing SAE libraries. I use SAELens https://github.com/decoderesearch/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 https://github.com/evan-lloyd/mechinterp-experiments/tree/gemma-blog-post/transformers sae . The base model, Gemma-2-2B, is run using the huggingface transformers library https://github.com/huggingface/transformers . 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 https://github.com/huggingface/datasets to fetch training data. All SAEs were trained using 50-100 million tokens from the 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 https://elloworld.net/posts/replacement-aware-sae-training/ . 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 https://www.lesswrong.com/feed.xml 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 https://www.lesswrong.com/feed.xml 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. To run the benchmarks reported in this post, I used the linked huggingface datasets MMLU https://huggingface.co/datasets/cais/mmlu Hendrycks et al. 2020 , CommonsenseQA CQA https://huggingface.co/datasets/tau/commonsense qa Talmor et al. 2018 , and the ARC Easy set ARC-E https://huggingface.co/datasets/allenai/ai2 arc 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. 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 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 https://www.lesswrong.com/feed.xml CQA fine tuning so that results are not contaminated when evaluating CQA-finetuned replacement models. 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. For results that are marked as "CQA fine-tuned," I repeat the appropriate version of KL fine-tuning KL fine-tuning https://www.lesswrong.com/feed.xml 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 https://www.lesswrong.com/feed.xml 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. 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 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 https://www.lesswrong.com/feed.xml 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. 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 https://www.lesswrong.com/feed.xml 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 https://www.lesswrong.com/feed.xml 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. Aside from the concerns outlined in Error nodes https://www.lesswrong.com/feed.xml 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. 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 . 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. 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 . ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-1 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-2 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-3 I did attempt a version of this that tuned all encoders in parallel, but I got much better results training them serially as written. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-4 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-5 This simplified the code and gave me marginally better performance on the hardware I was using. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-6 I use an exponential moving average of the adaptive threshold in Equation 12, following the implementation in SAELens Bloom et al. 2024 . ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-7 The original paper used the soft shrinkage function , with threshold parameters learned along with the update matrices and shared across iterations. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-8 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 https://elloworld.net/posts/replacement-aware-sae-training/ fig-validation-kl . ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-9 All reported results on CQA are over the validation split, so this won't contaminate them. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-10 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-11 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-12 "ish", because the implicit objective here, the mean replacement model RRE, is actually one dimensional. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-13 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-14 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-15 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-16 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-17 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. ↩︎ https://www.lesswrong.com/feed.xml fnref-nd9E4j8WdoJ7ewJPj-18