Constructive Circuit Amplification: Targeted Sub-Network Updates to Boost Math Reasoning in LLMs

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

Summary

Constructive Circuit Amplification (CCA) is a mechanistic-interpretability method that improves mathematical reasoning in large language models by updating only a small, targeted subset of internal components. Building on findings that LLM behavior relies on sparse circuits and that fine-tuning tends to reinforce existing circuits, CCA operates in three stages. First, it builds an error-localization dataset by comparing correct and incorrect reasoning traces, using a branching procedure to identify pivotal tokens where greedy decoding diverges between right and wrong answers. Second, a Desiderata-based Component Masking (DCM) approach learns binary masks over attention heads (Q, K, V projections) and MLP neurons, amplifying components that promote the desired token while L1 regularization enforces sparsity. Third, gradient updates are applied only to the selected components, maximizing the logit difference between desired and undesired tokens. Experiments on GSM-Symbolic with Gemma-2-9b-Instruct report accuracy gains of up to 11.4% while modifying only about 1.59% of model components, with negligible impact on general benchmarks such as MMLU, TriviaQA, and TruthfulQA. CCA is modular and task-agnostic, offering a precise alternative to full-model fine-tuning for improving multi-step reasoning.

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

Paper: Prakash, N., et al. (2025). arXiv:2512.16914v1

Overview

Large language models (LLMs) possess general reasoning abilities but frequently make small logical errors on math tasks. Prior research in mechanistic interpretability shows that LLM behavior depends on sparse sub-networks (circuits), and that fine-tuning typically reinforces existing circuits rather than creating new ones. Circuits also compete internally: some promote correct reasoning while others inject noise.

Constructive Circuit Amplification (CCA) exploits these insights: instead of updating the whole model, it identifies components that promote correct reasoning paths and amplifies only those, minimizing interference with other capabilities.

The CCA Method: Three Stages

1. Error Localization

Generate correct traces via greedy decoding and incorrect traces via non-greedy decoding on GSM-Symbolic instances.
A naive Prefix method finds the first non-shared token, but is sometimes imprecise.
The Branching method iteratively greedy-decodes the prefix and defines a pivotal token: one whose addition flips the final answer from correct to incorrect (or vice versa). Formally, $$y_k$$ is pivotal if $f(y_1,...,y_{k-1}) \in A_{corr}$ but $f(y_1,...,y_k) \in A_{incorr}$ , where $f(\cdot)$ is the greedy final answer.
The dataset records the prefix, desired token, and undesired token.

2. Component Localization via DCM

Desiderata-based Component Masking (DCM, implemented with NNsight) learns binary masks over attention heads (Q, K, V) and MLP neurons. The amplification uses:

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

with loss:

L = −(logit_{desired\ token} − logit_{undesired\ token}) + \lambda \sum m

$$m_i = 1$$ doubles the component output to favor the desired token; L1 regularization ( $\lambda$ ) enforces sparsity.
Training: Adam, 50 epochs, lr = 5e-3, batch = 8, masks clamped to [0,1], early stopping if 20% of batches show unchanged masks.

3. Targeted Parameter Updates

Gradient descent for 50 steps on the full dataset using the negative logit-difference loss, updating only the selected components.
Evaluation on GSym-Val every 2–10 steps selects the best checkpoint; learning rate swept from 1e-2 to 1e-5.

Experiments

Benchmark: 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.
Model: Gemma-2-9b-Instruct.
Interference checks: MMLU, TriviaQA, TruthfulQA.

Results

Up to +11.4% accuracy on GSM-Symbolic while updating only 1.59% of components.
Near-zero degradation on general benchmarks (MMLU, TriviaQA, TruthfulQA), demonstrating the precision of the targeted updates.

Significance and Limitations

CCA is modular and pluggable for any task (only hyperparameters like learning rate and sparsity weight need adjustment). Compared with chain-of-thought prompting (Wang et al., 2022b) or generic fine-tuning, it offers a mechanistically grounded, minimally invasive intervention. Limitations include the need for reasoning traces and relatively small localized datasets; future work targets more complex multi-step reasoning.

References

1. Prakash, N., et al. (2025). Constructive Circuit Amplification: Improving Math Reasoning in LLMs via Targeted Sub-Network Updates. arXiv:2512.16914v1. 2. Wang, C., et al. (2022a). Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small. 3. Hanna, M., et al. (2023). How does GPT-2 compute greater-than?: Interpreting mathematical reasoning in a 16-layer transformer. 4. Mirzadeh, I., et al. (2025). GSM-Symbolic: Understanding the Limitations of Mathematical Reasoning in Large Language Models. 5. Jain, S., et al. (2023). Mechanistically analyzing the effects of fine-tuning on procedurally defined tasks.