Technical Summary: Residualized Sparse Autoencoders for Multi-Layer Transformer Interventions

Problem Statement

Sparse Autoencoders (SAEs) are a primary tool for mechanistic interpretability, decomposing transformer hidden states into sparse, interpretable features. However, standard practice trains SAEs independently on each layer's raw activations. This approach ignores a fundamental structural property of transformers: the residual stream is strongly coupled across depth. A significant fraction of any layer's activation consists of information carried forward or lightly transformed from earlier layers.

Training independent SAEs on raw activations creates two practical problems for multi-layer interventions:

1. Capacity Redundancy: Different layerwise dictionaries waste capacity relearning the same predictable, carried-forward information.
2. Unpredictable Interactions: When multiple layers are replaced simultaneously (e.g., via feature ablation or reconstruction), the interactions between layers are not additive. If the same information is represented in multiple dictionaries, simultaneous replacement can cause features to duplicate, cancel, or drive activations out of the training distribution, leading to compounded errors and degraded model performance (cross-entropy).

Methodology: Residualized Sparse Autoencoders (ReSAEs)

The authors propose Residualized Sparse Autoencoders (ReSAEs) to decouple layer-specific information from predictable cross-layer structure. The method involves three key stages:

1. Pairwise Affine Residualization

For a selected subset of layers $L = (\ell_1, \ell_2, \dots, \ell_M)$, the method fits an affine map between consecutive selected layers. For a pair $(\ell_m, \ell_{m+1})$, an affine transformation $A_m h^{(\ell_m)} + c_m$ is learned via ridge-regularized ordinary least squares to predict the activation of the later layer from the earlier one.
The residualized activation $r^{(\ell_{m+1})}$ is then defined as the component of the later layer's activation not explained by this affine map:
$$r^{(\ell_{m+1})} = h^{(\ell_{m+1})} - (A_m h^{(\ell_m)} + c_m)$$
The first anchor layer is centered by its empirical mean.

2. Independent Per-Layer SAE Training

Each SAE is trained on the residual rather than the raw activation.

• Input: The residual activation is normalized (via per-block RMS normalization) and fed into a standard BatchTopK SAE.
• Objective: The SAE minimizes the reconstruction loss of the residual target.
• Independence: Gradients do not flow between SAEs of different blocks; each SAE learns only the "new" information introduced at its specific layer.

3. Original-Space Reconstruction and Online Intervention

To evaluate ReSAEs using standard protocols, the reconstructed residual is mapped back to the original activation space through the fitted affine chain.

• Reconstruction: $\hat{h}^{(\ell_{m+1})} = A_m \hat{h}^{(\ell_m)} + c_m + S^{-1}_{m+1} \hat{z}^{(m+1)}$, where $\hat{z}$ is the SAE reconstruction of the residual.
• Online Intervention: During a forward pass with hooks, the residual input for layer $\ell_m$ is computed by subtracting the predicted contribution of the previously reconstructed layer $\hat{h}^{(\ell_{m-1})}$ from the current incoming activation. This ensures that the SAE at layer $\ell_m$ sees an input distribution similar to its training target, even if upstream layers have been modified.

Key Contributions

1. ReSAE Architecture: Introduction of a training paradigm where SAEs are trained on the component of activation not linearly predictable from a previous selected layer, with reconstructions composed via an affine chain.
2. Improved Multi-Layer Intervention: Demonstration that despite reconstructing less raw variance, ReSAEs recover transformer cross-entropy more faithfully under simultaneous multi-layer replacement compared to raw SAEs.
3. Reduced Redundancy: Evidence that residualization produces more orthogonal decoder directions, indicating that capacity is directed toward layer-local information rather than duplicated cross-layer structure.
4. Comprehensive Evaluation: Extensive benchmarks on Pythia-1.4B and Gemma-2-9B covering reconstruction quality, cross-entropy degradation, decoder similarity, sparse probing, targeted probe perturbation (TPP), and spurious correlation removal (SCR).

Results

Experiments were conducted on Pythia-1.4B (layers 0, 6, 12, 18) and Gemma-2-9B (various layer gaps).

• Reconstruction vs. Functionality: ReSAEs exhibit lower Explained Variance (EV) in the original activation space compared to raw SAEs, as they intentionally discard linearly predictable structure. However, this lower EV correlates with lower decoder cosine similarity (less redundancy).
• Cross-Entropy Recovery: Under multi-layer sequential-online replacement (where errors propagate), ReSAEs significantly outperform raw SAEs in recovering the model's cross-entropy, particularly at higher sparsity levels ($k \ge 160$). In teacher-forced settings (no error propagation), ReSAEs consistently achieve the best cross-entropy recovery.
• Overinteraction: Raw SAEs show large negative "overinteraction" (non-additive degradation) when multiple layers are replaced. ReSAEs reduce this metric toward zero, indicating that their layerwise replacements compose more predictably.
• Feature Quality:
  • Sparse Probing & TPP: ReSAEs generally outperform raw SAEs in sparse probing accuracy and targeted probe perturbation, suggesting the learned features are more useful for linear readout and selective ablation.
  • Spurious Correlation Removal (SCR): ReSAEs do not uniformly outperform raw SAEs here; in some settings (particularly per-layer joint ablations on Pythia), raw SAEs score higher. This suggests residualized features may be less selective for removing specific spurious attributes without damaging task signals.
• Scale and Spacing: The benefits of ReSAEs persist in the larger Gemma-2-9B model across different layer gaps, with consistent reductions in decoder redundancy and improvements in cross-entropy recovery.

Significance and Claims

The paper argues that treating the transformer residual stream as a coupled object rather than a collection of isolated layers is crucial for effective multi-layer interventions.

• Capacity Allocation: By removing linearly predictable cross-layer structure, ReSAEs force the sparse dictionary to focus on the components of the activation most relevant to the model's downstream computation at that specific layer.
• Intervention Stability: The method provides a more stable default for multi-layer interventions, reducing the compounding errors and unpredictable interactions that plague standard SAE deployments.
• Trade-off: The authors acknowledge that ReSAEs are not a universal improvement for all interpretability tasks (e.g., SCR results are mixed). However, for the specific goal of multi-layer replacement and feature-based intervention, the trade-off of lower raw variance for higher functional fidelity and reduced redundancy is beneficial.

The work suggests that removing predictable structure via affine regression is a useful preprocessing step for SAEs in multi-layer settings, though the optimal choice of predictor (affine vs. nonlinear) and layer selection remains an open design space.