PaperlayerView exact revision on arXiv

Source-linked AI summary

RACE: Scalable Statistical Estimation of Functional Consistency in LLM Neurons

Runyu Wang, Bo Liu, Xiaxin Zhang, Yu Han, Jiawei Cao, Xiaoye Zhang, Zhe Zhang, Yifan Yang, Peng Ping

arXiv:2608.24758v1cs.AI

TL;DR

Domain-wide neuron analysis needs evidence that captures consistency across diverse inputs without the computational cost of gradient-based procedures. RACE uses forward-pass residual alignment and Bayesian aggregation to rank neurons by stable domain contributions. Perturbation and distributional results show stronger target specificity, while the method remains substantially more efficient than gradient-based approaches.

  • Problem

    Real-world auditing, pruning, and steering require population-level neuron rankings consistent across diverse inputs, beyond instance-level explanations.

  • Method

    RACE decomposes residual-stream updates into neuron contributions, scores their module-output alignment during forward passes, and Bayesian-aggregates the resulting evidence.

  • Results

    Targeted suppression causes severe target-domain distribution collapse while leaving WikiText-2 almost entirely unperturbed, including ∆PPL changes of 77.31% on MBPP+ and 129.04% on MATH-500.

  • Takeaways & Limitations

    RACE identifies neurons whose domain contributions are directionally aligned and statistically stable, and targeted perturbations confirm this population is causally loadbearing.

  • Takeaways & Limitations

    RACE may miss complex non-linear synergies and has limited effectiveness for identifying functionally consistent neurons in attention modules.

Abstract

from arXiv · show

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.

1 Introduction

Existing interpretability methods often explain individual instances, but domain-wide applications require neuron rankings that remain consistent across diverse inputs. RACE addresses this need with forward-pass statistical scoring and Bayesian aggregation, and experiments show selective disruption, stronger specificity, and lower overhead than gradient-based methods.

  • Population-level applications require rankings of task-relevant neurons that remain consistent across diverse inputs.
  • Existing task-neuron pipelines incur prohibitive overhead from iterative gradient calculations and intervention operations, while simplifying contributions into task-level averages.
  • RACE scores every neuron at every layer by aligning forward-pass residual contributions with module output directions, then Bayesian-aggregating per-observation signals.
  • RACE is evaluated on 4B–32B LLMs across code generation, mathematical reasoning, and fine-grained behavioral control.
  • Targeted suppression selectively disrupts target capabilities while preserving non-target behaviors, and the method outperforms gradient-based baselines with lower computational overhead.

2 Method

RACE evaluates neuron consistency from token-position evidence induced by a target domain, combining residual decomposition, module-local alignment, and Bayesian inference. Its posterior estimates capture alignment magnitude and stability, while CAM provides conservative neuron rankings.

  • 2.1 Problem Formulation: A target domain is approximated by sampled inputs and analyzed token positions, forming the observation set used for neuron auditing.
  • 2.1 Problem Formulation: Functional consistency requires large positive signed contributions and low variance across observations, excluding sparse or directionally unstable behavior.
  • 2.3 Evidence Aggregation: RACE models each neuron’s per-observation evidence with latent mean and variance parameters, then performs posterior inference to represent magnitude, direction, and stability.
  • 2.2.2 Module-Local Alignment Evidence: RDA obtains per-observation evidence by projecting each neuron’s weighted output onto its host module’s normalized residual update direction.
  • 2.2.1 Neuron-Level Residual Decomposition: The module update decomposes into additive per-neuron residual-stream contributions, including MLP intermediate channels and attention output-channel contributions.
  • 2.3 Evidence Aggregation: A Normal-Inverse-Gamma aggregation uses token-position evidence counts and closed-form updates to estimate signed alignment strength and uncertainty.
  • 2.4 Uncertainty Quantification: Posterior uncertainty increases with bursty, noisy, or sign-fluctuating evidence and shrinks as stable evidence accumulates.
  • 2.4 Uncertainty Quantification: RACE uses CAM, a one-sided conservative lower credible bound on positive posterior mean alignment, to penalize uncertain or negative evidence.

3 Experiments

The experiments evaluate RACE across domains, models, intervention settings, distributional effects, fine-grained behavior, and sample-size regimes. Results show domain-selective suppression, module-dependent effects, targeted behavioral steering, and benefits from Bayesian uncertainty modeling.

  • Experimental Setup: Experiments span code, mathematics, and fine-grained behavioral control across Qwen3-4B-it, OLMo-3.1-32B-it, and Llama-3.1-8B-it.Each domain uses a scoring set, same-domain OOD benchmark, and non-target benchmarks for retention.
  • Depth-Wise Organization of CAM Scores: RACE scores shift toward higher positive CAM density in later layers, while middle layers contain fewer high-scoring neurons, consistently across code and math.This depth-wise organization motivates ranking within each layer and module with fixed per-layer budgets rather than globally sorting neurons.
  • Domain-Specific Intervention: RACE-selected neurons outperform gradient-based methods, while stronger OOD degradation supports their use as target-domain proxies.The experiments also test whether Bayesian CAM scoring outperforms deterministic scoring heuristics.
  • Domain-Specific Intervention: Under RSF, suppressing RACE-selected neurons harms target-domain performance more than non-target performance, with MLP interventions producing stronger same-domain OOD degradation than attention interventions.RSF also improves suppression tolerance as intervention budgets increase, indicating better isolation of domain-specific neurons.
  • Distributional Verification: MLP suppression causes target-domain distribution collapse while leaving WikiText-2 nearly unchanged, with ∆PPL reaching 77.31% on MBPP+ and 129.04% on MATH-500.Attention interventions show weaker distributional contrasts and sensitivity to perturbations.
  • Fine-Grained Behavioral Steering: Suppressing PyComp-1K neurons reduces comprehension usage by 70.9% on MBPP+ and 61.1% on HumanEval+ while largely preserving functional correctness.The result supports steering a specific behavior rather than causing mere broad disruption.
  • Sample-Size Analysis: RACE’s low-data advantage arises from CAM’s uncertainty penalty, which avoids over-ranking weak fluctuations that can reduce intervention specificity.The sample-size analysis distinguishes calibrated uncertainty from a different mean estimator.

4 Related Work

Mechanistic interpretability has largely emphasized instance-level analysis, while dataset-level evaluation remains difficult to scale. RACE addresses this gap with a statistical framework for model auditing.

  • Dataset-level evaluation remains computationally challenging for methods focused on individual instances, including gradient analysis, residual-stream analysis, circuit discovery, and sparse autoencoders.The limitation concerns scaling detailed analyses from individual instances or features to whole datasets.
  • RACE addresses the scaling and robustness gap with a statistical framework for model auditing.

5 Conclusion

RACE treats functional consistency as a population-level statistical inference problem over neuron alignment distributions. The approach uses linear-complexity computation and targeted perturbations to identify neurons associated with domain functions.

  • RACE infers functional consistency from latent alignment distributions rather than single observations.Its posterior isolates neurons with directionally aligned and statistically stable domain contributions.
  • RACE-selected neurons are validated as causally loadbearing through targeted perturbations.
  • RACE is designed with linear computational complexity for practical application.

Limitations

The discussion motivates residual alignment as a layer-local measure of neuron contribution, while identifying limits from nonlinear interactions and attention-module behavior. Final-output relevance is therefore evaluated separately from local alignment.

  • Limitations: RACE’s linear residual-stream projections may miss nonlinear synergies among neurons and distributed polysemantic features.
  • Limitations: RACE and other evaluated baselines show limited effectiveness for identifying functionally consistent attention neurons.The paper suggests attention may encode information differently from MLPs, motivating attention-specific audit strategies.
  • Residual alignment: RDA evaluates a neuron’s contribution to its module’s own residual update rather than directly proving final-output causality.Final behavioral relevance is tested separately through targeted suppression and reference-set filtering.
  • Residual alignment: RDA scores partition the module output norm into aligned and orthogonal components, with the orthogonal component representing cancelled superposition interference.
  • Residual alignment: A neuron’s residual write, rather than its scalar preactivation, is treated as the natural unit of functional contribution.Raw activation magnitude does not indicate whether the resulting vector supports, opposes, or is orthogonal to the module computation.
  • Population inference: Bayesian aggregation promotes a consistently positive alignment share into a population-level functional role while penalizing cross-observation instability.Large variation or sign changes increase dispersion and reduce the resulting consistency score.
  • Residual alignment: The module increment is used instead of the full residual because the full residual includes computation accumulated from earlier layers.This keeps the measure tied to the work performed at the neuron’s own depth.

D Cross-Domain Consistency of RACE-Selected Neurons

The cross-domain analysis compares RACE-selected neurons across mathematical reasoning, code generation, and general language modeling. Attention neurons are substantially more shared across domains than MLP neurons, while intervention results show domain-specific effects primarily for MLPs.

  • Setup: The analysis evaluates mathematical reasoning, code generation, and general language modeling across ATTN and MLP modules.For each layer and module, neurons are ranked by CAM at top-k thresholds of 1%, 5%, and 10%.
  • Cross-domain overlap: At Top-1%, attention neurons achieve an all-domain Jaccard of 0.264, compared with 0.085 for MLP neurons.The reported values correspond to roughly 26% versus 8.5% of selected neurons shared across all three domains.
  • Cross-domain overlap: The attention–MLP sharing gap persists across both scoring metrics and all top-k thresholds.As k increases, MLP overlap rises moderately, but attention remains consistently more shared.
  • Interpretation: The results suggest attention output channels contain more reusable cross-domain routing or integration neurons than MLP down-projection neurons.MLP neurons show narrower response patterns under RACE scoring.

F Model, Generation, and Evaluation Details

The evaluation protocol reports benchmark scores using EvalScope, with unspecified parameters left to evaluator or backend defaults. Model and decoding settings, along with architectural dimensions, are documented in the accompanying tables.

  • Evaluation protocol: Benchmark scores follow the official metric implementation exposed by EvalScope 1.5.0, using vLLM 0.15.1 as the inference backend.Unless otherwise specified, the resulting benchmark score is reported as Domain Accuracy (DA) in the main tables.
  • Evaluation protocol: Table 10 records model and decoding settings, while unlisted parameters remain at evaluator or backend defaults.An em dash indicates that a parameter is unset.
  • Model configurations: Table 11 documents the architectural dimensions of all evaluated models.

G Computational Architecture

The experiments use specified GPU environments, benchmark and corpus metadata, and a controlled 64-input/16-output profiling protocol. Figure 4 compares layer-wise cross-domain neuron overlap, while PyComp-1K is constructed from parsed Python source statements.

  • Compute environment: Qwen3-4B-it experiments ran on eight NVIDIA RTX 4090 GPUs, while Llama-3.1-8B-it and OLMo-3.1-32B-it ran on eight NVIDIA A100 80GB GPUs.
  • Cross-domain overlap: Figure 4 heatmaps report Jaccard similarity between domain pairs across transformer layers for selected modules, attribution metrics, and top-k thresholds.Higher values indicate greater cross-domain consistency in the identified important neurons.
  • Compute measurement: The profiling reference uses the same 64-input/16-output window and costs 130.100 GFLOPs.The reference includes the prefill-to-generation increment and selected lm_head projections.
  • Evaluation control: All methods share the WikiText-2 token window, target positions, target token ids, and target modules, with one lookahead token defining the next-token target.GxAct and AttnLRP batch all target layers together by default.
  • Corpus construction: PyComp-1K contains 1,000 Python statements extracted by streaming The Stack and parsing files with Python’s ast module.Extraction targets ListComp, SetComp, DictComp, and GeneratorExp nodes, while unparsable rows are skipped.

I.2 Dataset Statistics

PyComp-1K is a 1,000-statement corpus for Python comprehension behavior, with examples paired with source metadata and AST labels. Qualitative perturbation examples contrast correct unmodified outputs with incorrect outputs after suppressing RACE-selected neurons.

  • Dataset composition: 1,000 Python statements comprise the PyComp-1K dataset, and each statement may contain multiple comprehensions.The dataset distribution is reported in Table 18.
  • Dataset composition: 754 ListComp, 181 GeneratorExp, 76 DictComp, and 9 SetComp occurrences are reported in PyComp-1K.These counts describe comprehension types, and a single statement can contribute to multiple categories.
  • Dataset records: Each dataset row includes the extracted statement, original line numbers, AST comprehension types, and The Stack source metadata.Metadata includes commit hash, repository name, and file path.
  • Qualitative perturbations: Five representative perturbation samples compare correct unmodified model outputs with incorrect outputs after suppressing RACE-selected PyComp-1K neurons.
  • Qualitative perturbations: The examples cover squaring list elements, sorting a dictionary, forming a sorted tuple union, extracting strings by length, and retrieving a k-th array element.The displayed outputs include both original and perturbed code blocks for these tasks.
  • Qualitative perturbations: The extraction example filters strings whose lengths equal the requested size, whereas the perturbed output contains an invalid list-comprehension form.

K Robustness to Prior Settings and Confidence Levels

RACE-selected neurons remain stable across broad prior settings, while confidence primarily filters marginal candidates rather than changing top-k rankings. Distributional and benchmark evaluations assess perturbation effects across code and math domains.

  • Prior robustness: Jaccard similarity between top-1% selected neurons remains at least 0.98 across λ0, β0 ∈ [10^-3, 10^3] and α0 ∈ [0.5, 50].The posterior is described as dominated by empirical statistics, making prior-parameter choice largely inconsequential.
  • Confidence robustness: Tightening γ from 0.05 to 0.001 shrinks valid positive-CAM neurons from 81.75% to 62.24% at N = 10 without displacing the most prominent neurons.Confidence acts primarily as an absolute verification threshold rather than changing relative top-candidate rankings.
  • Distributional evaluation: Token-distribution disruption is measured with relative perplexity degradation ∆PPL and mean forward KL divergence D̄KL.∆PPL captures predictive-quality deterioration, while D̄KL measures mean per-token distributional shift.
  • Distributional evaluation: Qwen3-4B-it evaluations suppress top-k = 5 RACE-selected neurons per layer and compare the first 100 samples per dataset against the unmodified model.The resulting ∆PPL and D̄KL measurements are summarized in Table 19.
  • Domain evaluations: Additional evaluations report code-domain accuracy and ISI for RMBPP+\WikiText-2, and math-domain accuracy for RMATH-500.The code results are summarized in Figure 2 and Table 20, while math results appear in Tables 21 and 22.

P Bayesian Variance Regularization in Low-Data Regimes

The appendix extends the variance-regularization argument for RACE's Bayesian update and shows that CAM preserves a finite uncertainty margin when evidence is limited. This margin decreases as evidence grows and becomes negligible in large-evidence regimes.

  • Setup: N = |Dc| counts input samples, while n = |Tc| counts induced token-position evidence for a fixed scoring set.These quantities parameterize the NIG update under the default RACE prior µ0 = 0, λ0 = 1, α0 = 1, and β0 = 1.
  • Lower bound: The posterior scale for mean evidence remains positive even when empirical dispersion collapses to SSj = 0.The appendix derives this property from αn = 1 + n/2 and λn = 1 + n.
  • Implication: CAM retains a finite conservative margin against weak or sparsity-induced evidence at finite n.This prevents near-zero empirical variance from eliminating the uncertainty penalty.
  • Implication: As n grows, the lower bound decays, making the prior negligible in large-evidence regimes.This matches the appendix's stated connection to the main-text observation about large-evidence settings.
  • Empirical context: Table 22 reports accuracy (%) after suppressing top-1% ATTN or MLP target-selected neurons per layer in the Math domain.The table uses RMATH-500\WikiText-2 on Qwen3-4B-it.
Loading 2608.24758v1…