🤖 machine learning

Interpretability-Guided Layer Selection over Subspace Projection: SAEs as Stethoscopes, Not Scalpels, for Raw Task Vector Model Editing

This paper demonstrates that projecting task vectors onto SAE feature subspaces fails as a model editing strategy due to geometric misalignment, and instead proposes using SAEs as diagnostic "stethoscopes" to selectively inject raw task vectors into specific layers, significantly improving mathematical reasoning performance without additional computational cost.

Original authors: Li Lei, Madalina Ciobanu, Qingqing Mao, Ritankar Das

Published 2026-05-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Li Lei, Madalina Ciobanu, Qingqing Mao, Ritankar Das

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Fixing a Brain Without Breaking It

Imagine you have a very smart, general-purpose robot (a Large Language Model) that knows a little bit about everything. You want to teach it to be an expert at Number Theory (a specific type of math) without making it forget how to do other things or slowing it down.

Usually, to teach a robot a new skill, you have to retrain its whole brain. This is expensive and risky; it might "forget" old skills or get confused. The researchers wanted a "surgical" approach: just tweak the specific parts of the brain responsible for Number Theory and leave the rest alone.

The Two Tools They Used

To do this surgery, they used two different tools:

The Task Vector (The "Instruction Manual"): This is a list of tiny changes needed to turn the robot from "General Knowledge" to "Math Expert." Think of it as a stack of sticky notes saying, "Change this gear," "Tighten that bolt," etc.
The SAE (The "Stethoscope"): A Sparse Autoencoder (SAE) is a tool that listens to the robot's internal thoughts. It can tell you which parts of the brain are thinking about math and which are thinking about poetry. Think of it as a doctor's stethoscope that can pinpoint exactly which organ is working hard.

The Mistake: Using the Stethoscope as a Scalpel

The researchers first tried a very logical idea. They thought:

"Let's use the Stethoscope (SAE) to find the math-thinking parts of the brain. Then, let's use the Stethoscope to filter the Instruction Manual (Task Vector). We will only let the instructions pass through if they match the math-thinking parts exactly."

The Result: It failed completely.

The Analogy: Imagine you have a high-resolution blueprint for a house (the Task Vector). You try to photocopy it through a tiny, blurry keyhole (the SAE filter) to see if it fits a specific room. The result is a tiny, distorted, unrecognizable smudge. You lost 97% of the information.
The Science: The "Instruction Manual" lives in Weight Space (the physical structure of the robot's brain). The "Stethoscope" listens in Activation Space (the robot's internal thoughts). Trying to force the physical instructions through the thought-filter caused a geometric mismatch. The signal was almost entirely lost.

The Solution: The Stethoscope is for Diagnosis, Not Surgery

The researchers realized their mistake. They changed their strategy:

"Let's use the Stethoscope (SAE) only to find the right rooms in the brain. Once we know which rooms are for math, we will take the entire, unfiltered Instruction Manual and paste it directly into those rooms."

The Result: It worked beautifully.

The Analogy: Instead of trying to squeeze the blueprint through a keyhole, the Stethoscope simply points to the correct doors. You then walk right up to those doors and hand the full, high-quality blueprint to the workers inside.
The Outcome: On a math test called "Minerva Math," the robot's Number Theory score jumped from 29.6% to 39.4%. It improved in 5 out of 7 math subjects and didn't get worse at any of them.

Key Takeaways

Don't Filter the Signal: Trying to force a "weight-space" fix (changing the brain's structure) through an "activation-space" filter (the robot's thoughts) destroys the signal. It's like trying to pour a gallon of water through a coffee filter; most of it gets stuck.
SAEs are Great Doctors, Bad Surgeons: The SAE is excellent at diagnosing where the problem is (which layers of the brain need help), but it is terrible at deciding what to change.
Keep it Raw: When you fix the brain, use the full, raw instructions. Don't try to "translate" them through the SAE.

Summary

The paper proves that if you want to surgically edit an AI, you should use interpretability tools (like SAEs) to find the right location, but then apply the full, unfiltered changes directly to that location. Trying to filter the changes through the interpretability tool kills the improvement.

Technical Summary: Interpretability-Guided Layer Selection over Subspace Projection

Problem Statement
Large Language Models (LLMs) increasingly require "surgical" model editing to enhance specific domain capabilities without the computational cost or catastrophic forgetting associated with full fine-tuning. Two primary research threads address this: Task Arithmetic, which treats weight differences between fine-tuned and base models as additive "task vectors" in weight space, and Mechanistic Interpretability (MI), which uses Sparse Autoencoders (SAEs) to decompose activation space into interpretable features for intervention.

The prevailing hypothesis suggests a synergistic pipeline: use SAEs to identify domain-relevant features and then inject task vectors only through those specific features (subspace projection). This paper rigorously evaluates this hypothesis on mathematical reasoning tasks using Gemma-3-4B-IT and uncovers a fundamental failure mode in this approach.

Methodology
The paper contrasts two pipelines for SAE-guided model editing (illustrated in Figure 1):

The Failed Approach (SAE Projection):
- Process: Task vectors ( $\Delta W$ ) are projected onto the subspace spanned by domain-specific SAE decoder vectors before injection.
- Mechanism: This attempts to filter weight-space modifications through activation-space directions.
- Outcome: This acts as a severe information bottleneck, discarding approximately 97% of the modification energy. The authors attribute this to a geometric misalignment: SAE decoder vectors span the activation manifold (capturing statistical regularities of representations), while task vectors exist in weight space (encoding computational changes). Projecting weight-space modifications through activation-space directions conflates two distinct geometric structures, resulting in near-total loss of functional signal.
The Proposed Approach (SAE as Stethoscope, Not Scalpel):
- Process: SAEs are used strictly for diagnosis (layer selection), not intervention (filtering).
- Step 1 (Task Vector Extraction): A task vector is extracted via LoRA fine-tuning on a domain-specific dataset (e.g., Number Theory).
- Step 2 (SAE-Guided Layer Selection): A per-layer Specificity Score (SP) is computed. For each layer $l$ , the score is the maximum specificity of its SAE features, where specificity is the ratio of mean activation on target-domain inputs versus other domains. Layers with $SP(l) \geq \tau_{SP}$ (e.g., 4.0) are selected.
- Step 3 (Selective Raw Injection): The unfiltered, raw task vector ( $\Delta W$ ) is injected only into the selected layers, scaled by a global parameter $\alpha$ . No projection or masking is applied to the vector itself.
- Key Insight: This preserves 100% of the modification energy while confining edits to domain-relevant regions. The authors observe an empirical conservation law: as the number of selected layers decreases, the optimal scaling parameter $\alpha$ increases proportionally ( $n_{layers} \times \alpha_{opt} \approx 11.2$ ).

Key Results
Evaluated on the Minerva Math benchmark (7 subjects, including Number Theory, Counting & Probability, Algebra, etc.) using Gemma-3-4B-IT:

Failure of Projection: The SAE projection method (injecting projected vectors into selected layers) yielded no statistically significant improvements across any of the seven math subjects, even when increasing SAE feature width from 16K to 262K. It retained only ~2–3.5% of the modification energy.
Success of Raw Injection: The proposed method (injecting raw vectors into SAE-selected layers) significantly improved performance:
- Number Theory: Accuracy increased from 29.6% to 39.4% ( $z = +3.41, p = 0.0007$ ).
- Overall: 5 of 7 subjects showed statistically significant improvements; none were significantly degraded.
- Comparison: Full LoRA merging (applying $\Delta W$ to all 34 layers) degraded performance in several subjects due to interference.
Robustness: The method is robust to the scaling parameter $\alpha$ within a broad range (0.70–1.10) and is fully deterministic at float64 precision.

Key Contributions

Methodological Framework: Proposes "SAE-Guided Layer Selection with Raw Task Vector Injection," demonstrating that SAEs are effective for identifying where to edit (diagnosis) but ineffective for filtering what to inject (intervention).
Substantive Negative Result: Documents that projecting task vectors onto SAE feature subspaces discards ~97% of modification energy and fails to produce gains. This is presented as the central scientific contribution, challenging the intuitive assumption that activation-space feature selection can directly filter weight-space updates.
Empirical Evidence of Geometric Mismatch: Provides strong evidence that activation-space SAE directions and weight-space task vectors are not interchangeable, supporting recent literature questioning the reliability of SAEs as intervention tools beyond probing.
Extensive Ablation: Includes 30+ configurations testing feature widths, layer selection thresholds, task vector sources, and scaling parameters.

Significance and Claims
The paper claims that the distinction between activation space and weight space is a "load-bearing axis" of the model editing problem. The primary significance lies in redefining the role of interpretability tools:

SAEs as Stethoscopes: They are valuable for diagnosing domain specialization at the layer level.
SAEs are not Scalpels: They should not be used to surgically filter weight-space modifications via projection.

The authors conclude that effective surgical editing requires a separation of roles: use SAEs to identify the target layers, but inject the raw, unfiltered task vector to preserve the functional signal. This approach offers a principled, training-free (post-training), and deterministic framework for enhancing domain capabilities without catastrophic forgetting. The findings suggest that the efficacy of this diagnostic selection is contingent on the "sharpness" of the task vector itself, requiring focused domain training data.