New SAE training method lets Gemma-2-2B replacement models retain language without error nodes

Researchers studying mechanistic interpretability often build 'replacement models' where internal LLM activations are replaced with interpretable features from sparse autoencoders (SAEs). This approach, however, suffers from severe degradation—even a few SAE-replaced layers can make the model incoherent. To fix this, previous work (Marks et al. 2024) introduced explicit error nodes that cancel out the SAE's reconstruction errors, but Lloyd found this method 'deeply unsatisfying' because it obscures the true causal structure.

In a new cross-post on LessWrong, Lloyd presents an alternative: replacement-aware training. Instead of patching errors afterwards, they train SAEs that are robust to upstream distortions by adding a loss term that penalizes damage to the next layer's features. Scaling this approach from a toy model to Gemma-2-2B, they produced a suite of residual stream SAEs that can be used as-is in a full replacement model. While performance on benchmarks like MMLU is still below chance without fine-tuning, the replacement models retain basic language coherence and significantly lower KL divergence from the base model compared to standard SAEs. Additional tricks—a short (~1 million token) encoder fine-tuning and a LISTA-based encoder—further reduce divergence. Lloyd considers these results an existence proof that error nodes are unnecessary if SAEs are designed for the replacement use case.