Towards Isolated Interventions via Almost Orthogonal Features in Language Models

Technical Summary: Towards Isolated Interventions via Almost Orthogonal Features in Language Models

Problem Statement
A central premise of mechanistic interpretability is that meaningful concepts in language models (LMs) correspond to linear features in activation space. For these features to support reliable causal interventions, manipulating one feature should not substantially alter the effects of others. However, in practice, learned features are often entangled. This "feature interference" occurs when features are strongly aligned in representation space, causing localized interventions (e.g., steering a model) to have unintended downstream effects, blurring causal contributions, and limiting intervenability. This issue is exacerbated by the phenomenon of superposition, where models represent more features than they have dimensions, leading to interference patterns that are difficult to control.

Methodology
Motivated by the Independent Causal Mechanisms (ICM) principle, which posits that complex systems can be decomposed into autonomous modules that do not influence each other, the authors propose constraining internal features to be almost orthogonal. The methodology involves a two-step fine-tuning pipeline:

1. Orthogonal SAE Training: The authors train a Sparse Autoencoder (SAE) on the residual stream activations of a language model. Unlike standard SAEs, they introduce an explicit orthogonality regularization penalty on the decoder matrix $D$. The loss function includes a term $\lambda \| \text{tril}(D^\top D) \|_F^2$, which penalizes the off-diagonal elements of the Gram matrix, thereby enforcing low self-coherence (near-orthogonality) among feature atoms.
2. LM Adaptation: Once the SAE is trained and its weights are frozen, the authors insert it into a pre-trained Transformer (specifically a Gemma 2 2B model) at an intermediate layer (layer 12 of 26). They then perform low-rank adaptation (LoRA) on the language model to adapt it to the new feature geometry while preserving overall performance.

Theoretical Framework
The paper formalizes the relationship between feature interference and intervenability. By modeling the LM's representation as a sparse linear combination of features ($X = DZ$) and assuming a linear hypothesis for the output mapping, the authors derive an upper bound on the expected discrepancy between the true output and the model output following an intervention.

• They show that the post-intervention mismatch is controlled by a baseline term and an intervention-specific term proportional to the self-coherence of the feature dictionary, denoted as $\mu(D)$.
• Specifically, the expected discrepancy is upper-bounded by a term involving $\delta^2 (d-1)\mu(D)$, where $\delta$ is the magnitude of the intervention. This theoretically demonstrates that reducing self-coherence (increasing orthogonality) directly limits the propagation of interference, thereby improving the isolation of interventions.

Key Contributions

1. Theoretical Bound: The authors provide a theoretical upper bound relating intervenability in the residual stream to feature interference, explicitly linking the discrepancy in model outputs to the self-coherence of the feature dictionary.
2. Orthogonal Fine-Tuning Pipeline: They demonstrate a pipeline that fine-tunes an LM around an SAE with an almost orthogonal decoder. They show that this approach maintains performance and interpretability scores comparable to non-penalized SAEs on mathematical reasoning tasks.
3. Isolated Interventions: The authors perform the first isolated interventions on specific concepts within an SAE by explicitly limiting spill-over effects. They successfully swap concepts (e.g., changing the name "Jerry" to "Aquaman") while preserving the model's reasoning capabilities and minimizing unintended changes to other features.

Results

• Orthogonality vs. Performance: Empirical evaluation on the GSM8K dataset shows that increasing the orthogonality penalty $\lambda$ monotonically decreases the orthogonality loss (improving feature separation) without significantly degrading the model's accuracy on mathematical reasoning tasks. The performance of models with strict orthogonality penalties remains competitive with the non-penalized baseline.
• Intervention Efficacy: In experiments involving the substitution of first names (e.g., "Jerry" to "Aquaman") and other concepts (animals, locations), models with higher orthogonality penalties ($\lambda \ge 10^{-6}$) demonstrated significantly better isolation.
  • Correct Substitution: As the penalty increased, the model more frequently included the correct substituted name in its generation (e.g., recovering "Aquaman" ~70.9% of the time with $\lambda=10^{-4}$ vs. ~60.1% without penalty).
  • Reduced Spill-over: The models with orthogonality penalties were more successful at dropping the original concept without erroneously retaining it in the output.
• Interpretability: The interpretability of the features (measured by the ability of an LLM to correctly match feature explanations to text snippets) remained constant across different orthogonality penalties, indicating that enforcing orthogonality does not sacrifice the human-understandability of the features.

Significance and Claims
The paper claims that feature entanglement is a primary inhibitor of isolated interventions in language models. By formally relating intervenability to the ICM principle and demonstrating that orthogonality regularization reduces feature interference, the authors suggest a principled route to improving the identifiability and atomicity of SAE features without sacrificing downstream utility.

The authors modestly frame their work as a step toward more reliable and predictable AI systems. They argue that their approach allows for precise, localized modifications of internal representations, which is crucial for understanding how specific components contribute to outputs. They acknowledge that while their method improves the structure and controllability of learned representations, it introduces potential dual-use concerns regarding the precise steering of model outputs. However, they note that their approach requires direct access to model internals and technical expertise, which limits accessibility and reduces the likelihood of misuse. The work does not claim to solve all issues of superposition or adversarial robustness but suggests that encouraging orthogonality may act as a strong architectural bias toward more robust models.