Constructive Circuit Amplification: Surgical LLM Tuning with Just 1.59% of Parameters Boosts Math Reasoning by Up to 11.4%
In the world of large language models, the most frustrating failures are not "I don't know" but "I know it, but I wander off." On a simple word problem, a model can correctly compute totals and write subtraction, then suddenly read "6 boys are absent" as "no mention of absent boys" and confidently produce a wrong answer, as if a bright student was distracted by a bird outside the window.
The paper "Constructive Circuit Amplification: Improving Math Reasoning in LLMs via Targeted Sub-Network Updates" (Prakash et al., 2025) proposes what sounds like sci-fi surgery: instead of large-scale retraining, first locate the exact token where reasoning first goes off-track, then identify the attention heads and MLP neurons that push the model toward the correct path, and finally update only this tiny subset (as little as 0.17%, at most 1.59% of parameters). The result: math reasoning accuracy improves by up to +11.4%, while general capabilities on MMLU, TriviaQA, TruthfulQA remain nearly intact.
This deep dive focuses on Desiderata-based Component Masking (DCM) and targeted sub-network updates, explaining how CCA finds "which neurons to touch," why such minimal edits work, what the experiments show, and where the boundaries lie.
---
Key points
1. Why "touching very little" can beat a full-body massage
The approach rests on two empirical facts from mechanistic interpretability:- Capabilities are not uniformly distributed. Specific behaviors are dominated by sparse sub-networks called circuits, small sets of attention heads and MLP neurons (Wang et al., 2022a; Hanna et al., 2023; Prakash et al., 2024).
- Fine-tuning usually reinforces existing circuits rather than inventing new mechanisms (Jain et al., 2023; Prakash et al., 2024; Chhabra et al., 2025). The model is "almost capable"; the correct internal pathway is just weak or outcompeted by noisy pathways (Rai et al., 2025; Ortu et al., 2024).
- Error-Localization dataset. Each sample has a shared prefix up to the intervention token, a desired_token (correct branch's next token), and an undesired_token (wrong branch's next token). DCM does not learn the final answer; it learns a local preference for the next token.
- Mask intervention (Eq. 1):
- Loss with L1 sparsity (Eq. 2):
- What gets masked. Q/K/V heads (with grouped-attention key/value head counting) and MLP neurons. In Gemma-2-2B-It with Branching token selection, ~3969±398 MLP neurons are selected versus hundreds of heads, suggesting attention handles routing while MLPs perform rule transformations.
- Objective stays local: same negative logit-difference between desired and undesired tokens. DCM phase changes no weights; the update phase fixes the mask and only updates masked components.
- Training recipe: 50 gradient steps, evaluated every 2 steps (up to step 10) then every 10 steps, best validation checkpoint tested; lr sweep over {1e-2, 5e-3, 1e-3, 5e-4, 1e-4, 5e-5, 1e-5}; full-batch gradients rather than mini-batches because the dataset is small and few steps are expected.
- Circuit: the DCM-selected sparse component set.
- Amplification: both the 2x DCM probe and the final updates strengthen these components.
- Constructive: they specifically promote correct reasoning, since the objective contrasts desired vs undesired tokens rather than imitating any single answer, unlike standard fine-tuning which may reinforce correlated but unreliable shortcuts.
- Gemma-2-9B-Instruct: 0.807 → 0.881 ±0.015 (+7.4%), 0.17% of parameters updated
- Gemma-2-2B-Instruct: 0.411 → 0.525 ±0.010 (+11.4%), 1.59% updated (largest share)
- OLMo-2-1124-13B-Instruct: 0.742 → 0.786 ±0.005 (+4.4%), 0.44% updated
- OLMo-2-1124-7B-Instruct: 0.739 → 0.794 ±0.005 (+5.5%), 0.25% updated
- Validated mainly on math reasoning (GSM-Symbolic); generalization to code, scientific reasoning, and multimodal tasks is untested.
- Building Error-Localization datasets requires paired correct/incorrect trajectories and a correctness signal, which is expensive or unavailable for many real tasks.
- Only a single round of targeted amplification is studied; multi-round continual learning and catastrophic forgetting remain unexplored.
Think of internal computation as parallel "candidate reasoning paths": some push toward the right answer, some introduce plausible but wrong shortcuts. CCA aims to make the correct path louder so it wins the next-token logit competition.
2. The three-step CCA pipeline
1. Token Localization. On paired correct/incorrect reasoning trajectories, find the pivotal token where the wrong trajectory first diverges and select the preceding token as the intervention point. 2. Model Component Localization. DCM learns a sparse binary mask over attention heads and MLP neurons, picking those whose amplification most increases the desired-token logit relative to the undesired one. 3. Targeted Parameter Updates. Only the masked components receive gradient updates (the rest are frozen), for ~50 steps of "minimally invasive enhancement."3. DCM masking: turning "the next correct word" into an optimizable wish list
If \(m_i=1\), that component's output is doubled; if $0$, unchanged. The fixed 2x gain is chosen for identifiability rather than fine control, making it easier to tell which components push the desired direction.
The first term maximizes the desired-vs-undesired logit gap; the second enforces sparsity. Training uses Adam at lr 5e-3, batch size 8, 50 epochs with early stopping when 20% of a batch shows no mask changes, masks clamped to [0,1] after each step, implemented via NNsight.
4. Sparse parameter updates: where the actual "surgery" happens
5. Why it is called "Constructive Circuit Amplification"
6. Experimental results (GSM-Symbolic, "CCA w mask" with Branching token selection)
The story is "interpretability-friendly": tiny structural change, reliably visible gain.
7. Masked vs unmasked ablation
In some settings CCA without a mask (CCA w/o mask) is marginally higher, e.g., Gemma-2-2B Branching: 0.532 vs 0.525 with mask. The mask's value is therefore not always peak accuracy but fewer edits, closer adherence to minimum-intervention, and reduced risk to general capabilities, moving the method from empirical fine-tuning toward mechanism-aligned targeted enhancement.8. General-capability preservation
On five general benchmarks (Table 2), most absolute changes are within ±1 percentage point, with the worst case around -4.0 on TriviaQA for Gemma-2-9B Prefix w/o mask; the masked variants avoid the most extreme drops. Minimal edits = minimal disturbance to other behaviors.9. How DCM + sparse updates differ from LoRA
LoRA solves "how many new parameters to add for task adaptation." CCA solves "which internal mechanisms of the original model must change." DCM is the imaging scan that locates the spot; sparse updates are the treatment that fixes only that spot.10. Limitations (faithfully from the paper's Discussion)
---