Source-linked AI summary

Sparse Reward Subsystem in Large Language Models

Guowei Xu, Mert Yuksekgonul, James Zou

arXiv:2602.00986v2cs.CL

TL;DR

Existing approaches decode reward-related information from full LLM hidden states without showing how it is structured across neurons. This paper uses probing, pruning, and intervention to identify a sparse reward subsystem of value and dopamine neurons, finding that it supports confidence prediction and process reward modeling for inference-time search.

  • Problem

    Existing approaches use black-box probes on full hidden states, leaving the neuron-level structure of reward-related information insufficiently characterized.

  • Method

    The paper combines probes, pruning, and interventions to identify value neurons encoding state value and dopamine neurons encoding step-level TD errors.

  • Results

    Reward information is concentrated in a sparse neuron subset; value neurons are robust and transferable, while dopamine neurons track TD-error patterns during reasoning.

  • Takeaways & Limitations

    Value neurons provide lightweight confidence estimates, and dopamine neurons can serve as an intrinsic process reward model for guiding inference-time search.

  • Takeaways & Limitations

    The superposition-based sparsity explanation is motivational rather than a rigorous proof, and dopamine training and evaluation retain only steps with |δt| > 0.3.

Abstract

from arXiv · show

Recent studies show that LLM hidden states encode reward-related information, such as answer correctness and model confidence. However, existing approaches typically fit black-box probes on the full hidden states, offering little insight into how this information is structured across neurons. In this paper, we show that reward-related information is concentrated in a sparse subset of neurons. Using simple probing, we identify two types of neurons: value neurons, whose activations predict state value, and dopamine neurons, whose activations encode step-level temporal difference (TD) errors. Together, these neurons form a sparse reward subsystem within LLM hidden states. These names are drawn by analogy with neuroscience, where value neurons and dopamine neurons in the biological reward subsystem also encode value and reward prediction errors, respectively. We demonstrate that value neurons are robust and transferable across diverse datasets and models, and provide causal evidence that they encode reward-related information. Finally, we show applications of the reward subsystem: value neurons serve as effective predictors of model confidence, and dopamine neurons can function as a process reward model (PRM) to guide inference-time search.

1 Introduction

The paper argues that reward-related information in LLM hidden states is concentrated in a sparse reward subsystem rather than being characterized only through full-state black-box probes. It identifies value neurons and dopamine neurons and demonstrates their predictive, causal, transferable, and practical roles.

  • Research gap: Existing full-hidden-state probes decode reward-related information but do not characterize how it is structured across neurons.Identifying responsible neurons matters for mechanistic interpretability and lightweight, targeted applications.
  • Reward subsystem: The paper identifies value neurons that predict state value and dopamine neurons that predict step-level temporal-difference errors.Together, these neuron types form a reward subsystem within LLM hidden states.
  • Sparse neurons: Less than 1% of neurons suffice to predict state value, according to pruning experiments.The paper also establishes value-neuron specificity through intervention experiments.
  • Sparse neurons: Dopamine-neuron activations show peaks during unexpected progress and troughs when reasoning errors occur.These patterns support their interpretation as step-level reward-prediction-error signals.
  • Applications: Value neurons predict model confidence before generation, while dopamine neurons can function as an internal process reward model for inference-time search.The reward subsystem is presented as useful for both confidence estimation and search guidance.

2 Motivation

The paper motivates a sparse reward subsystem by viewing LLMs as policies over reasoning trajectories and by arguing that value and TD-error information is low-dimensional. Maximum-entropy reinforcement learning and superposition provide motivation, while the authors explicitly distinguish these arguments from rigorous proof.

  • Value in reasoning trajectories: An autoregressive LLM can be viewed as a policy over partial reasoning trajectories, with value representing the probability that continuation eventually yields a correct answer.The trajectory begins at an initial state, actions update the state, and correctness supplies the terminal reward.
  • Why reward information may exist: Maximum-entropy reinforcement learning motivates a connection between high-quality policies and value functions, suggesting that LLM representations may encode value information.The paper presents this as motivation for why reward-related neurons may exist, not as a proof.
  • Why neurons may be sparse: The posterior probability of terminal correctness is a sufficient scalar summary for binary prediction, so reward-relevant information can be low-dimensional in principle.The argument applies analogously to TD-error prediction because TD error is also scalar.
  • Why neurons may be sparse: The superposition hypothesis motivates realizing frequently useful, high-importance value and TD-error features as sparse representations.The authors caution that this hypothesis does not prove these neurons must be sparse.

3 Sparse Value Neurons in Large Language Models

The paper identifies sparse value neurons whose activations predict state value, tests their predictive sufficiency and reward specificity, and finds them robust across datasets, models, and layers. Intervention and overlap analyses further support a functionally coherent reward subsystem.

  • Value probe: A two-layer MLP value probe predicts scalar reward from layer-wise hidden states and is trained with a temporal-difference objective.The probe evaluates hidden states for each reasoning state and uses TD errors during optimization.
  • Sparse value neurons: Less than 1% of neurons can preserve value-prediction performance under pruning, defining the retained subset as value neurons.Stable AUC as pruning increases indicates that the probe relies on a very small fraction of neurons.
  • Robustness: Value neurons remain predictive across multiple benchmarks, models, and layers, while final-position analysis shows AUC mostly remains above 0.8.The reported benchmarks include MATH500, ARC, MBPP+, and IFEval across four models; the initial-position evaluation is described as challenging.
  • Intervention evidence: Zeroing the top 1% of value neurons causes catastrophic degradation beyond Magnitude and Wanda, whereas NTP-neuron ablation has no effect and overlaps only 0.7%.The intervention is evaluated on Qwen-2.5-7B-SimpleRL-Zoo using MATH500; the original accuracy is 75.2%.
  • Intervention evidence: Ablating earlier-layer value neurons impairs subsequent dopamine-neuron TD-error prediction, while random-neuron ablation has minimal effect, supporting subsystem coherence.The intervention summary frames this as evidence for necessity, specificity, and functional coherence around reward signals.
  • Transferability: IoU for value neurons exceeds the random baseline across datasets and increases as pruning approaches 1, indicating stability of the most critical neurons across tasks.The cross-dataset comparison uses GSM8K, MATH500, and ARC in layer 3 of Qwen-2.5-14B-SimpleRL-Zoo.
  • Transferability: IoU also remains above the random baseline between two RLVR-derived Qwen models, suggesting transferability of value-neuron positions across related models.The comparison uses layers 2–4 and value probes trained on MATH500.

4 Sparse Dopamine Neurons in Large Language Models

Dopamine neurons are a sparse subset whose activations encode step-level TD errors, with pruning and visualization evidence showing reward-prediction-error signals during reasoning.

  • Identifying Dopamine Neurons: Dopamine neurons are defined as a sparse subset whose activations encode temporal-difference reward prediction errors.They predict the TD error at each reasoning step.
  • Training a Dopamine Probe: The TD error measures unexpected reward change: positive values indicate a reasoning step improves expected correctness, while negative values indicate a detrimental step.TD errors are computed from discounted value estimates at successive paragraph boundaries.
  • Training a Dopamine Probe: The dopamine probe predicts TD errors from normalized hidden states, retaining only steps with |δt| > 0.3 and identifying neurons through L1-norm pruning.Pruned probes are evaluated by Spearman correlation between predicted and Monte Carlo-estimated TD errors.
  • Empirical Evidence: Spearman correlation curves remain largely invariant across pruning ratios for Qwen-2.5-7B/14B-SimpleRL-Zoo on MATH500, indicating reward-prediction-error information is concentrated in sparse neurons.The experiments use Qwen-2.5-7B/14B-SimpleRL-Zoo models and MATH500.
  • Visualization Results: Dopamine neurons peak when the model makes critical progress and dip when it encounters a logical flaw, tracking TD errors during inference.Positive-surprise reasoning produces peaks, whereas negative-surprise reasoning produces a trough at the error moment.

5 Applications

The reward subsystem supports two applications: value neurons estimate model confidence before generation, while dopamine neurons score reasoning steps to guide inference-time search.

  • Applications of Value Neurons: Value neurons provide a lightweight confidence estimator because they encode value information before response tokens are generated.Pre-generation confidence can support adaptive compute allocation for long reasoning responses.
  • Applications of Value Neurons: Value neurons achieve the highest average AUC across four baselines, four models, and two benchmarks, showing that sparse neurons carry informative confidence signals.The comparison includes LCD, Verbal Confidence, Token Confidence, and Question Length on MATH500 and ARC.
  • Applications of Dopamine Neurons: Dopamine neuron predictions can serve as process rewards in guided search, where four candidate continuations are scored at each paragraph boundary and the highest-scoring candidate is selected.The procedure repeats until a final answer is produced.
  • Applications of Dopamine Neurons: Using dopamine neuron predictions as a process reward improves reasoning performance on MATH500, demonstrating their use for guiding reasoning.The evaluation compares the dopamine probe with greedy decoding, random selection, and an implicit PRM.

6 Related Work

Related work uses probes and activation-based methods to decode correctness, confidence, latent knowledge, and other internal reward signals from LLM representations.

  • Probing LLM Representations: Probing methods are widely used to investigate and interpret internal LLM characteristics, including prediction of forthcoming-answer correctness and extraction of syntactic relations.Prior approaches include linear probes and polar probes.
  • Reward Signals in Activations: Recent studies extract reward signals from LLM activations using value heads, unsupervised latent-knowledge discovery, and mechanistic-interpretability tools.These approaches target correctness prediction and related information in internal representations.

7 Conclusion

The paper argues that reward information is concentrated in a sparse reward subsystem comprising value and dopamine neurons, with transferable structure and practical uses for confidence estimation and search.

  • Conclusion: Probing and pruning identify value neurons encoding value information and dopamine neurons encoding reward prediction errors in a sparse reward subsystem.Intervention experiments specifically tie value neurons to reward signals.
  • Conclusion: Value-neuron positions transfer across datasets and models, while value neurons provide lightweight confidence estimates and dopamine neurons can serve as an intrinsic process reward model.These applications extend the subsystem beyond representation analysis.

A The Benefit of Using the TD Error Training Objective

The section compares TD-error training with training exclusively on final reward, using intervention results and a randomized-neuron baseline to assess value-neuron identification.

  • A The Benefit of Using the TD Error Training Objective: The ablation trains a reward model exclusively on the final reward signal and compares its value-neuron identification against the TD-error objective.The comparison uses intervention results from Qwen-2.5-7B-SimpleRL-Zoo on MATH500.
  • A The Benefit of Using the TD Error Training Objective: Intervention performance is measured by accuracy after zeroing out a 1% subset of neurons in a single layer.
  • A The Benefit of Using the TD Error Training Objective: The section also defines randomly sampled neuron sets of size k = (1 − p)N for comparing overlap against a random baseline.For two independently sampled sets, the intersection follows a hypergeometric distribution.

D.1 Robustness across Different Datasets

Value-neuron findings remain robust across datasets, model architectures and scales, and layers, with AUC staying largely stable under pruning.

  • D.1 Robustness across Different Datasets: Across MBPP+ coding and IFEval instruction-following data, the AUC curve remains largely invariant to pruning.This supports consistent observation of value neurons across different dataset types.
  • D.2 Robustness across Different Models: Across Qwen, Phi, and Llama models of different scales on MATH500 and ARC, sparse value neurons remain identifiable.The experiments include Qwen3.5-0.8B, Phi-3.5-mini-instruct, Llama-3.1-8B-Instruct, and Qwen-2.5-14B-SimpleRL-Zoo.
  • D.3 Robustness across Different Layers: Across layers of Qwen-2.5-7B/14B-SimpleRL-Zoo on GSM8K, AUC remains largely stable as the pruning ratio rises.

E Transferability Across Different Datasets: More IoU Curves

Value-neuron positions transfer across GSM8K, MATH500, and ARC, with pairwise overlap consistently exceeding the random baseline as pruning varies.

  • E Transferability Across Different Datasets: More IoU Curves: Pairwise IoU curves for value neurons across GSM8K, MATH500, and ARC consistently exceed the random baseline.The analysis uses layers 2 and 4 of Qwen-2.5-14B-SimpleRL-Zoo.
  • E Transferability Across Different Datasets: More IoU Curves: Many cross-dataset IoU curves rise significantly as the pruning ratio approaches 1.
  • F Detailed Procedures for Hidden State Normalization: Hidden states are extracted from all transformer layers in a single forward pass, then normalized across response-token positions before downstream aggregation.Paragraph-level representations average normalized activations over the tokens belonging to each paragraph.

G Further Evidence for the Characteristics of Dopamine Neurons

Dopamine neurons track positive and negative prediction errors during reasoning, while value-neuron ablation substantially disrupts their activation trajectories. Additional analyses show that value neurons encode reward information robustly at response completion, with AUC mostly above 0.8 across four models.

  • Dopamine-neuron activation: Dopamine neurons show high activation when the model makes unexpected progress or acquires unexpected rewards.On Minerva Math, a critical conclusion around the 200th token produces a sharp activation spike, followed by lower activation as later reasoning becomes predictable.
  • Dopamine-neuron activation: The Minerva Math visualization consistently correlates dopamine-neuron activation with TD error during inference.The same analysis includes both positive and negative surprise cases, supporting the proposed prediction-error interpretation.
  • Interconnection with value neurons: Value-neuron ablation significantly alters dopamine-neuron trajectories, whereas random ablation has minimal impact.Zeroing the top 20% of value neurons shifts activation peaks and troughs, causing the dopamine neuron to lose characteristic prediction-error patterns.
  • Robustness across generation positions: Value-neuron performance remains mostly above 0.8 AUC at the final response position across four models and varied pruning ratios.This completed-response setting yields stronger signals than the pregeneration setting and supports reward encoding at both ends of generation.
  • Limitations: The method requires environmental reward signals and has not been tested on open-ended tasks where such signals are difficult to obtain.The authors also report that models larger than 32B were not investigated because of resource and computational constraints.
Loading 2602.00986v2…