Constructive Circuit Amplification: Improving Math Reasoning in LLMs via Targeted Sub-Network Updates

Forum topic · ✨步子哥 · 2025-12-21

Summary

This post introduces Constructive Circuit Amplification (CCA), a method from a paper by Nikhil Prakash et al. that improves mathematical reasoning in large language models by updating only a tiny, targeted set of internal components. The approach rests on two findings from mechanistic interpretability: LLM tasks are executed by sparse sub-networks (circuits), and fine-tuning tends to reinforce existing circuits rather than create new ones. CCA works in three stages: (1) it localizes reasoning errors by generating correct and incorrect reasoning traces via greedy vs. non-greedy decoding and identifying pivotal tokens where answers diverge; (2) it trains binary masks (via a Desiderata-based Component Masking procedure) over attention heads (Q, K, V) and MLP neurons to find components promoting the desired token; (3) it performs targeted gradient updates on only those selected components. Reported results show up to 11.4% accuracy improvement on GSM-Symbolic while modifying only 1.59% of components, with minimal degradation on MMLU, TriviaQA, and TruthfulQA. The post explains the math, training details, and implications in an accessible way.

Overview

This post explains the paper "Constructive Circuit Amplification: Improving Math Reasoning in LLMs via Targeted Sub-Network Updates" (Prakash et al., 2025, arXiv:2512.16914v1), presented through an extended metaphor of neural surgery on an AI's "brain": circuits compete inside the model, and the correct reasoning path is often drowned out by noisy ones.

Key points

Motivation: Prior mechanistic interpretability research shows LLM behavior relies on sparse sub-networks (circuits) of attention heads and MLP neurons, and fine-tuning typically *reinforces* existing circuits rather than building new ones (Jain et al., 2023; Merullo et al., 2023). Circuits also compete, with constructive paths sometimes suppressed by noisy ones.
CCA method (three stages):

Error localization

pivotal token

Component localization via DCM (Desiderata-based Component Masking)

$h_{\text{org}} = m_i * 2 * h_{\text{org}} + (1 - m_i) * h_{\text{org}}$

with loss:

$L = -(\text{logit}_{\text{desired token}} - \text{logit}_{\text{undesired token}}) + \lambda \sum m$

Trained with Adam for 50 epochs, lr = 5e-3, batch size 8, early stopping when masks stabilize on 20% of batches; masks clamped to [0, 1]. L1 regularization enforces sparsity. 3. Targeted update: Gradient descent (50 steps) on only the selected components using the negative logit-difference loss; evaluation on GSym-Val every 2–10 steps selects the best checkpoint; learning rates swept from 1e-2 to 1e-5.

Experimental setup: GSM-Symbolic (Mirzadeh et al., 2025) — 100 templates, 50 instances each, split GSym-Train/Val/Test = 0.52/0.08/0.40, filtering templates with mean accuracy < 0.8. Models tested include Gemma-2-9b-Instruct. Side-effect benchmarks: MMLU, TriviaQA, TruthfulQA.
Results: Up to +11.4% accuracy on GSM-Symbolic while modifying only 1.59% of components, with near-negligible impact on general abilities (MMLU, TriviaQA, TruthfulQA).
Limitations & outlook: CCA requires reasoning traces and relatively small datasets; future work targets more complex reasoning. The method is modular — pluggable to any task by adjusting hyperparameters (learning rate, sparsity weight).

Related work referenced

Circuit discovery: Olah et al. (2020); Wang et al. (2022a) on indirect object identification; Hanna et al. (2023) on the "greater-than" circuit.
Fine-tuning and circuits: Jain et al. (2023); Merullo et al. (2023); sparse updates via RL fine-tuning (Mukherjee et al., 2025) and sparse supervised fine-tuning (Liu et al., 2025).
Math reasoning: chain-of-thought prompting (Wang et al., 2022b); Lewkowycz et al. (2022); probe-based correctness prediction (Sun et al., 2025b); ThinkEdit (Sun et al., 2025a).

Takeaway

Rather than retraining the whole model, CCA acts like widening key roads in a busy city: it identifies the few components that promote correct reasoning and amplifies them, leaving the rest of the network untouched — a precise, interpretable alternative to global fine-tuning.

Reference: Prakash, N., et al. (2025). Constructive Circuit Amplification: Improving Math Reasoning in LLMs via Targeted Sub-Network Updates. arXiv:2512.16914v1.