RACE: Scalable Statistical Estimation of Functional Consistency in LLM Neurons

AI summary

Overview

Research area: Mechanistic interpretability of Transformer-based large language models, specifically statistical methods for identifying neurons that perform consistent, domain-specific functions.

Technical level: Advanced. The paper combines mechanistic interpretability, Bayesian conjugate inference, and controlled ablation-style interventions on models from 4B to 32B parameters.

Scope: The paper introduces RACE (Residual Alignment for Consistency Estimation), a forward-pass statistical framework that scores every neuron in every layer and module of a Transformer for functional consistency with a target domain, then validates those scores through targeted neuron suppression.

What This Paper Is About

Existing mechanistic interpretability methods explain model behavior one input at a time or require expensive gradient computations, which makes it hard to characterize which neurons matter across an entire domain of inputs. RACE addresses this by decomposing each module's residual-stream update into per-neuron contributions, measuring how well each neuron's contribution aligns with its module's output direction on a single forward pass, and then using Bayesian aggregation to turn those noisy per-observation signals into population-level neuron rankings. The goal is a scalable way to identify task-relevant neurons whose behavior is stable across diverse inputs, so that those neurons can be reliably suppressed, pruned, or steered.

Key Contributions

1. Residual-Direction Alignment (RDA): A forward-pass method that decomposes each attention or MLP module's residual-stream update into per-neuron writes and projects each write onto the observation-specific normalized module-update direction, producing a signed scalar evidence value per neuron per observation at a cost of roughly one-hundredth of gradient-based methods.

2. Bayesian evidence aggregation with CAM scoring: A Normal-Inverse-Gamma conjugate model that distills per-observation evidence into a posterior over each neuron's mean alignment and variance, yielding closed-form posterior means and Conservative Alignment Magnitude (CAM) scores that act as one-sided conservative lower credible bounds on positive module-output alignment.

3. Reference-Set Filtering (RSF): An explicit filtration step that removes neurons ranked highly on a general reference corpus from the target-domain ranking, intended to disentangle domain-specific neurons from broadly reused general-purpose neurons during interventions.

4. Empirical validation across code, math, and fine-grained behavior: Targeted suppression experiments on Qwen3-4B-it, OLMo-3.1-32B-it, and Llama-3.1-8B-it, using MBPP+, HumanEval+, MATH-500, AMC, MMLU-Redux, GPQA, and a newly constructed PyComp-1K set of 1,000 AST-verified Python comprehension-containing statements from bigcode/the-stack.

Main Findings

• RACE outperforms gradient-based baselines on target-domain selectivity: On Qwen3-4B-it code-domain suppression with top-K_sel = 5 neurons per layer, RACE reached an Intervention Specificity Index (ISI) of 2.91 for MLP and 1.62 for ATTN, compared with AttnLRP (1.15 MLP, 2.00 ATTN), GxAct (0.04 MLP, 0.57 ATTN), and the Act. Mean control (2.22 MLP, 1.54 ATTN).

• RACE scales to 32B models: On OLMo-3.1-32B-it with R_MATH-500∖WikiText-2, RACE's MLP ISI was 1.65 at a top-1% budget and 1.73 at top-5%, versus Emp. SNR (0.00 and 0.20), Emp. Mean (1.35 and 1.50), Act. Mean (1.47 and 1.50), and Neg. CAM (0.00 and 0.00). ATTN ISI values for all methods at this scale were near zero.

• Large suppression budgets degrade generation: Under vanilla selection, neurons selected by these methods cause severe overall generation degradation once the budget reaches K_sel ≥ 10 per layer, which is why the code-domain experiment reports K_sel = 5.

• Module asymmetry: MLP interventions produce substantial same-domain out-of-distribution drops, while attention interventions show weaker domain-specific degradation. The authors attribute this to attention's high-scoring neurons being sparse but broadly reusable, consistent with known sparsity in W_O.

• Severe distributional disruption localized to the target domain: Suppressing RACE-selected MLP neurons on Qwen3-4B-it raised relative perplexity by 77.31% on MBPP+ and 129.04% on MATH-500 while changing WikiText-2 perplexity by only 2.76% and 2.60% respectively. Mean forward KL divergence for MLP was 0.58 (MBPP+), 0.22 (MATH-500), and 0.02 (WikiText-2); for the MATH ranking, 0.22 (MBPP+), 0.86 (MATH-500), and 0.03 (WikiText-2). ATTN interventions showed much weaker contrasts (+6.17%, +6.18%, −2.49% perplexity change, with KL values of 0.10, 0.08, and 0.03).

• Fine-grained behavioral steering without catastrophic forgetting: Using R_PyComp-1K∖WikiText-2 with K_sel = 10 neurons per layer, Python comprehension usage dropped by 70.9% in record count (16 records of 18) and 69.5% in total instances (18 of 59) on MBPP+, and by 61.1% (14 of 19) and 57.8% (19 of 45) on HumanEval+, while pass@1 on the comprehension-containing subsets remained at 87.50% and 85.71%.

• Depth-wise organization of CAM scores: Positive CAM density shifts rightward with depth, concentrating in later layers, while middle layers have a sparser distribution of high-scoring neurons, especially in MLP modules. The same pattern appears on both MBPP+ and MATH-500, implying CAM scores are not directly comparable across layers and motivating a stratified per-layer selection protocol.

• Token-position protocol matters: Under R_MBPP+∖WikiText-2 with per-layer top-1% MLP suppression on Qwen3-4B-it, collecting observations over all generated-token positions (the default) collapsed MBPP+ to 3.17% and HumanEval+ to 1.22%, whereas a first-token-only protocol left accuracy at 70.37% and 76.83%, against a no-suppression baseline of 82.28% and 83.54%.

• Localized suppression preserves non-target capability: Across the reported tables, non-target benchmarks such as MMLU-Redux and GPQA generally remain close to baseline after RACE-selected suppression, which is the intended effect of combining CAM scoring with RSF.

Methodology in Plain English

The method rests on two observations: specific LLM capabilities depend on sparse subsets of neurons, and those neurons respond differently to semantically coherent inputs than to unrelated ones.

For every input in a target domain, RACE runs the model once. At each layer and module, the module's contribution to the residual stream is written as a sum of per-neuron pieces, each piece being an activation value times a weight vector. RACE normalizes the module's total update to get a direction, projects each neuron's weight vector onto that direction, and multiplies by the neuron's activation. The result is a single signed number per neuron per input, recording whether that neuron pushed the module's output in the same direction as the module went, or against it.

Because a single input produces a noisy number, RACE treats the collection of numbers for a neuron across all inputs and token positions as noisy samples from a Gaussian, places a Normal-Inverse-Gamma prior over that Gaussian's mean and variance, and derives closed-form posterior updates that require only the running mean and sum of squares. This costs O(1) per neuron per observation.

The posterior mean estimates how strongly and in which direction a neuron consistently aligns with its module's output. The posterior uncertainty penalizes neurons whose evidence fluctuates or is scarce. RACE's headline score, CAM, subtracts a confidence radius from the posterior mean, so a neuron only scores highly if its mean alignment is positive and the evidence is tight. Optionally, Reference-Set Filtering removes neurons that also score highly on a general corpus such as WikiText-2, leaving neurons whose high ranking is specific to the target domain.

Validation works by suppression: set the selected neurons' activations to zero during inference and measure accuracy drops on target benchmarks, same-domain out-of-distribution benchmarks, and non-target benchmarks. A single Intervention Specificity Index combines the target drop and non-target drop into one number.

Why This Matters

Impact on research: The paper reframes neuron importance as a population-level statistical estimation problem rather than an instance-level attribution problem, and shows that a forward-pass alignment signal plus Bayesian uncertainty can outperform gradient attribution at roughly one-hundredth of the computational cost. It also provides a concrete argument that neuron scores are not comparable across layers, which affects how anyone building task-neuron maps should select neurons.

Real-world applications:

• Model capability auditing: Producing per-layer, per-module rankings of neurons tied to a domain gives auditors a population-level view of where a capability lives, rather than a per-example explanation.
• Domain-specific pruning: Because suppressing RACE-selected MLP neurons collapsed target-domain perplexity by 77.31% on MBPP+ and 129.04% on MATH-500 while barely touching WikiText-2, the framework offers a candidate signal for safely removing capacity tied to particular capabilities.
• Behavioral steering and style control: The PyComp-1K experiment shows comprehension usage can be reduced by 70.9% and 61.1% on two code benchmarks while functional correctness stayed near 87.5% and 85.7%, indicating suppression can shape style rather than break capability.
• Safety and selective capability removal: ISI-based measurement of target drop versus non-target retention gives a practical metric for evaluating interventions intended to disable one capability without degrading general performance.

Industry relevance: The cost argument matters for production-scale models. Gradient attribution over large observation sets is expensive; a single forward pass per observation plus closed-form arithmetic makes domain-wide auditing feasible for models in the 4B–32B range tested here, which is the size class commonly deployed or fine-tuned in practice.

Future Directions

• Extending validation beyond the configurations shown: The provided content reports the Llama-3.1-8B-it results as residing in Appendix E and full model details in Appendix F; those numbers are not present in the material summarized here.

• Resolving the attention-module weakness: ATTN ISI values near zero on OLMo-3.1-32B-it and the weaker distributional contrasts for ATTN indicate that the framework currently localizes MLP behavior far more cleanly than attention behavior; a better treatment of broadly reused attention neurons is an open problem.

• Making CAM comparable across layers: The depth-wise shift in CAM density forced a stratified selection protocol. The paper does not report a normalization scheme that would allow a single global ranking across depth.

• Robustness to design choices: The paper states that robustness to prior settings and confidence levels is discussed in Appendix K, and that cross-domain overlap analysis appears in Appendix D, but the specific results are not included in the content summarized here.

Target Audience

Researchers and engineers working on mechanistic interpretability, model editing, pruning, or behavioral control of large language models will benefit most. The paper is also relevant to practitioners who need scalable, low-cost tools for auditing which parts of a deployed model carry a specific capability, and to readers interested in Bayesian uncertainty quantification applied to model internals. Familiarity with residual streams, attention and MLP module structure, and conjugate Bayesian models is helpful, since the method section is written at that level.

Authors’ abstract

Discovering stable neuron behavior across entire domains remains a challenge in mechanistic interpretability. Existing methods often rely on instance-level point estimates or computationally expensive procedures, which either obscure population-level variability or limit scalable domain-wide analysis. We present RACE (Residual Alignment for Consistency Estimation), a forward-pass statistical framework that evaluates the domain-wide functional consistency of Transformer neurons. Perturbation experiments demonstrate that RACE achieves superior domain specificity compared to gradient-based point estimates. Meanwhile, token-distribution-level results verify the association between the selected neurons and the target domain. Furthermore, its computational overhead is two orders of magnitude lower than that of gradient-based methods.