Identifying and Transferring Reasoning-Critical Neurons: Improving LLM Inference Reliability via Activation Steering

AI summary

Overview

Research area: Natural Language Processing / LLM interpretability and inference-time intervention (activation steering for reasoning reliability).

Technical level: Intermediate. The paper assumes familiarity with Transformer MLP blocks, activation steering, and accuracy/AUROC metrics, but its core idea is explained through intuitive neuron-level analysis.

Scope: The paper introduces AdaRAS (Adaptive Reasoning Activation Steering), a training-free, test-time method that identifies a sparse set of "Reasoning-Critical Neurons" in LLM MLP blocks and selectively steers their activations to fix incorrect reasoning without harming correct reasoning.

What This Paper Is About

Large language models are unreliable on hard reasoning tasks, and most fixes rely on expensive post-training or costly test-time sampling. This paper investigates whether reasoning correctness is encoded in a small number of internal MLP neurons, and if so, whether directly modifying those neurons at inference time can improve reliability. The authors propose AdaRAS, which identifies such neurons by contrasting correct and incorrect reasoning traces, then adaptively steers them only when a failure is expected.

Key Contributions

1. First systematic evidence that reasoning correctness can be predicted and improved through neuron interventions. The authors report a pilot probing study showing that last-token MLP activations predict correctness with AUROC around 0.7, establishing activation steering as a viable tool for reasoning enhancement.
2. AdaRAS, a parameter-free, test-time activation steering framework. It identifies Reasoning-Critical Neurons via a polarity-aware mean-difference criterion and intervenes only when a failure is predicted, with no additional training or sampling during inference.
3. Demonstrated consistency, transferability, and scalability. Experiments on 10 mathematics and coding benchmarks show consistent accuracy gains, including over 13% on AIME-24 and AIME-25, with improvements of around 1% on Qwen3-4B and cross-dataset transfer within 1%.
4. Mechanistic insight into why it works. Trajectory-level analysis using magnitude and angle features shows AdaRAS reduces fluctuations in latent reasoning trajectories without substantially altering semantic modeling.

Main Findings

• A small subset of neurons correlates with reasoning correctness. Probing classifiers trained only on last-token activations achieve AUROC scores around 0.7 across two datasets and two models, reaching 0.7639 on AIME (24+25) with Qwen3-1.7B and 0.7153 with Qwen3-4B. On AMC-12, AUROC is 0.7091 for Qwen3-1.7B and 0.6727 for Qwen3-4B.
• Individual neurons show polarized activation patterns. The 3028-th neuron in layer 27 shows substantially higher activation during successful reasoning, whereas the 246-th and 1908-th neurons in layer 26 exhibit the opposite trend.
• Consistent accuracy gains across 10 benchmarks. AdaRAS improves over CoT prompting on every dataset: AIME-24 47.83 → 60.87 (+13.04), AIME-25 40.91 → 54.55 (+13.64), AIME-Extend 47.33 → 52.67 (+5.34), MATH-500 84.80 → 86.40 (+1.60), GSM8K 88.32 → 89.08 (+0.76), AMC-12 65.93 → 70.33 (+4.40), HumanEval 77.18 → 79.19 (+2.01), HumanEval+ 69.80 → 73.15 (+3.35), MBPP 68.78 → 72.22 (+3.44), MBPP+ 58.20 → 60.58 (+2.38).
• Average improvement of around 5% over CoT. Gains are larger on hard benchmarks: 9.11% average on AIME and AMC, compared to 2.26% on easier benchmarks.
• AdaRAS surpasses comparable post-trained models. On AIME-25, AdaRAS reaches 54.55% versus 40.91% for OpenThinker, 50.00% for Nemotron, and 22.73% for R1-Distill.
• Probing-based steering is unstable. On AIME-24, probing-based steering yields 43.48%, degrading accuracy by 4.35% relative to the unsteered CoT baseline, though it achieves 73.63% on AMC-12, surpassing AdaRAS there.
• Strong cross-dataset and cross-task transferability. Using RCNs identified on AIME instead of dataset-specific RCNs produces small but stable gains within 1% across all evaluated datasets, including coding benchmarks: MATH-500 86.40 → 87.40, GSM8K 89.08 → 89.39, AMC-12 70.33 → 71.43, HumanEval 79.19 → 80.54, HumanEval+ 73.15 → 73.15, MBPP 72.22 → 72.75, MBPP+ 60.58 → 61.11.
• Scales to stronger models. On Qwen3-4B, AdaRAS improves AIME-24 56.52 → 60.87, AIME-25 59.09 → 72.73, AIME-Extend 68.67 → 76.67, MATH-500 90.00 → 91.20, GSM8K 92.72 → 93.93, AMC-12 76.92 → 80.22, HumanEval 91.28 → 92.62, HumanEval+ 81.88 → 84.56, MBPP 79.37 → 82.28, and MBPP+ 68.78 → 69.31.
• Hyperparameter behavior. AdaRAS peaks at α = 0.3 on AIME-24 and α = 0.4 on AIME-25; performance drops to the unsteered baseline only when α = 0.7 on AIME-24. Steering improves as more RCNs are intervened on, peaks at approximately top-50 neurons (about 0.03% of all neurons), and degrades thereafter; intervening on top-2000 neurons (about 1.20%) drops AIME-24 performance by around 6% compared to top-50.
• All components matter. Ablations on Qwen3-1.7B show Random Steering 34.78 on AIME-24 and 27.27 on AIME-25 (below the 47.83/40.91 baselines); removing MD gives 43.48/50.00; removing activation selection gives 52.17/45.45; removing adaptive intervention gives 56.52/45.45; full AdaRAS gives 60.87/54.55. Disabling polarity-based selection causes an approximate 10% accuracy drop on both AIME-24 and AIME-25.
• Failures are predictable. The adaptive intervention module achieves AUROC of 0.8347 on the AIME dataset.
• Steering is concentrated in later layers. Compared to probing-based steering, which changes early layers widely, AdaRAS concentrates interventions on later layers with activation shifts tightly centered around zero.
• Latent trajectories become more stable. Steered samples show significantly lower and more concentrated magnitude values than unsteered samples, while the angle metric remains largely unchanged.

Methodology in Plain English

The approach has three stages.

1. Finding the critical neurons. The authors first need paired examples of the same problem solved correctly and incorrectly. They generate these by sampling multiple reasoning traces per problem at high temperature (temperature 1.0) and keeping only questions with a balanced mix of exactly 4 correct and 4 incorrect traces. Using this contrastive set, they score each MLP neuron by the difference in its mean activation between correct and incorrect traces. This is a global average over the whole reasoning trace, not just the final answer token.

2. Selecting a sparse, discriminative subset. They keep only neurons whose mean activation changes sign between correct and incorrect traces — a polarity-based filter — and then rank the survivors by the magnitude of their importance score, keeping the top-K (K = 50). This yields a sparse steering vector per layer, added inside the MLP block after the SwiGLU activation, scaled by a strength α. The vector is built once from a reference dataset and applied at every decoding step.

3. Gating the intervention. Because steering sometimes hurts samples that were already correct, they train a lightweight attention-based failure predictor on early neuron activations (selected by F-statistic, input dimension 256, architecture Linear(256,256) → ReLU → Dropout(0.3) → Linear(256,1), trained with Adam, BCEWithLogitsLoss, up to 100 epochs with early stopping patience 10). At inference, steering is applied only when a failure is predicted.

The probing classifier itself uses last-token MLP post-activations from all layers concatenated (dimension 131,072 for feature selection), normalized by 1/(10 × std_train), with ANOVA F-statistic feature selection and L1-regularized logistic regression trained with SAGA over 5-fold cross-validation. Contrastive traces for probing were generated by Qwen3-32B; 15 of an initial 60 AIME problems and 13 of an initial 104 AMC-12 problems survived the balance filter, giving 120 and 104 samples respectively. Experiments ran on 4 NVIDIA A100 GPUs and 8 NVIDIA RTX 3090 GPUs, with greedy decoding for all evaluations and EvalPlus used to execute generated code.

Why This Matters

Impact on research: The paper argues that reasoning reliability is a trajectory-level property rather than a single output attribute, and shows that it can be addressed through internal activation interventions instead of post-training or expensive sampling. It also offers a probing-free alternative to classifier-weight-based neuron importance and demonstrates that RCNs identified on a hard task transfer to other tasks and to larger models.

Real-world applications:

• Improving math and coding assistants that must answer competition-level or production-level problems reliably without retraining.
• Deploying a test-time patch for models already in production, since AdaRAS requires no parameter updates and applies the same precomputed steering vector at every decoding step.
• Reducing reliance on costly sampling strategies such as self-consistency or multi-step calibration in latency-sensitive applications.
• Transferring a steering vector derived from one domain (for example, AIME mathematics) to another (for example, code generation), reducing the cost of adapting to new tasks.

Industry relevance: The method is described as parameter-free and plug-and-play, with minimal additional computation cost and no additional training or sampling at inference time. That profile fits organizations that want accuracy gains on existing model checkpoints without post-training pipelines.

Future Directions

• Extending evaluation beyond the Qwen3 series models and STEM benchmarks to broader architectures and complex reasoning settings such as spatial or multi-hop reasoning.
• Relaxing the requirement for contrastive correct/incorrect data pairs, which the authors note may limit applicability to models at capability extremes where such pairs are hard to obtain.
• Incorporating other interpretability techniques such as sparse autoencoders or vocabulary projection to further explain the mechanistic role of RCNs.
• Investigating alternative data construction strategies and further analysis of adaptive intervention and stability, which the paper defers to its appendix.

Target Audience

Researchers and engineers working on LLM reasoning reliability, mechanistic interpretability, and activation steering. It is most useful to readers who want a training-free, test-time method for improving reasoning accuracy, and to practitioners interested in transferring steering vectors across datasets and model scales. Readers without background in Transformer internals will need to acquire familiarity with MLP blocks and activation-level intervention first.

Authors’ abstract

Despite the strong reasoning capabilities of recent large language models (LLMs), achieving reliable performance on challenging tasks often requires post-training or computationally expensive sampling strategies, limiting their practical efficiency. In this work, we first show that a small subset of neurons in LLMs exhibits strong predictive correlations with reasoning correctness. Based on this observation, we propose AdaRAS (Adaptive Reasoning Activation Steering), a lightweight test-time framework that improves reasoning reliability by selectively intervening on neuron activations. AdaRAS identifies Reasoning-Critical Neurons (RCNs) via a polarity-aware mean-difference criterion and adaptively steers their activations during inference, enhancing incorrect reasoning traces while avoiding degradation on already-correct cases. Experiments on 10 mathematics and coding benchmarks demonstrate consistent improvements, including over 13% gains on AIME-24 and AIME-25. Moreover, AdaRAS exhibits strong transferability across datasets and scalability to stronger models, outperforming post-training methods without additional training or sampling cost.