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

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

Summary

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

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