reward-lens: A Mechanistic Interpretability Library for Reward Models

1. Problem Statement

Reinforcement Learning from Human Feedback (RLHF) relies on Reward Models (RMs) to shape the behavior of Large Language Models (LLMs). If a reward model encodes incorrect preferences, the resulting policy inherits and amplifies those errors. While behavioral evaluation of RMs (e.g., RewardBench) is mature, mechanistic interpretability for RMs is virtually non-existent.

Existing interpretability tools (e.g., TransformerLens, Logit Lens, Direct Logit Attribution) are designed for generative LLMs. They rely on projecting intermediate model states onto the vocabulary unembedding matrix ($W_U$) to analyze token distributions. However, reward models replace $W_U$ with a scalar regression head ($w_r$). This structural difference renders standard primitives inapplicable, creating a "structural gap" where the internal mechanics of safety-critical reward models remain a black box.

2. Methodology: The $w_r$-Centred View

The paper introduces reward-lens, an open-source library that adapts mechanistic interpretability primitives for reward models. The core theoretical insight is that the reward head's weight vector, $w_r$, serves as the natural axis for all interpretability questions.

The library is built on the observation that the final reward $r$ is a linear functional of the final residual stream $h_{final}$:
$$r = w_r^\top h_{final} + b_r$$
Since the residual stream decomposes additively across layers ($h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$), the reward can be exactly decomposed as:
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

Based on this, the library implements the following primitives by replacing $W_U$ with $w_r$:

• Reward Lens: Projects intermediate residual streams at layer $\ell$ onto $w_r$ to track the "crystallization" of preference signals across layers.
• Component Attribution: Computes the contribution of each sublayer (Attention/MLP) to the final scalar reward via $w_r^\top \cdot \text{contribution}$.
• Contrastive Activation Patching: Swaps activations between preferred and dispreferred completions and measures the change in $w_r^\top h_{final}$. This is native to RMs due to the inherent contrastive structure of preference data.
• SAE Feature Attribution: Uses Sparse Autoencoders (SAEs) to decompose residual streams into features, calculating feature alignment as $w_r^\top d_i$ (where $d_i$ is a dictionary vector).

Architecture Agnosticism: The library uses a HuggingFace-native approach with a 10-method adapter protocol. This allows it to wrap standard models (Llama, Mistral, Gemma-2) and complex multi-objective models (like ArmoRM with MoE gating) without weight surgery.

Theory-Grounded Extensions: The library operationalizes recent alignment theory results into code:

• Distortion Index: Predicts which quality dimensions are under-covered by evaluation probes (risk of reward hacking).
• Divergence-Aware Patching: Detects when patching induces out-of-distribution activations.
• Misalignment Cascade Detector: Identifies correlated failure modes across different misalignment dimensions.
• Reward-Term Conflict Analyzer: Checks if different reward objectives are aligned, orthogonal, or in conflict.
• Concept-Vector Analysis: Extracts linear concept directions and measures their causal dose-response on the reward.

3. Key Contributions

1. The Library (reward-lens): The first installable, architecture-agnostic toolkit for mechanistic interpretability of reward models. It unifies observational (attribution) and causal (patching) views under a single algebraic framework.
2. Unifying Framing: It establishes that every useful interpretability primitive for an RM is a projection onto or decomposition along $w_r$.
3. Empirical Validation: Systematic application to two production models: Skywork-Reward-Llama-3.1-8B-v0.2 (single-scalar head) and ArmoRM-Llama3-8B-v0.1 (multi-objective head).
4. Negative Empirical Finding: The paper explicitly reports that linear attribution does not predict causal importance in reward models.

4. Key Results & Findings

The authors evaluated the library on $\sim$695 RewardBench preference pairs per model.

• Crystallization Depth:
  • Skywork: Preferences crystallize uniformly late (fractional layer depth $\approx 0.90$), suggesting a unified mechanism assembled near the output.
  • ArmoRM: Preferences crystallize earlier and with higher variance, reflecting its multi-objective, MoE-gated architecture.
• Attribution vs. Causality (The Faithfulness Gap):
  • Observational Attribution: The final MLP layers (specifically mlp_L31) dominate the reward signal in both models.
  • Causal Patching: Early-layer components (e.g., mlp_L0, attn_L0) are often the most critical for the final reward.
  • Correlation: The Spearman rank correlation between attribution magnitude and patching effect is negative or near zero (Skywork: $\rho = -0.256$; ArmoRM: $\rho = -0.027$).
  • Interpretation: This is not a bug but a feature of redundant residual streams. Late layers can be ablated because parallel paths compensate, whereas early layers are "load-bearing."
• Reward Hacking Profiles:
  • Skywork penalizes length, confidence, and formatting biases (unlike typical RLHF models that reward them).
  • ArmoRM rewards confidence and formatting but penalizes sycophancy.
• Concept Alignment:
  • Skywork shows strong linear alignment with concepts like "agreement" (slope $\approx 1.0$).
  • ArmoRM shows weak alignment ($\approx 1\%$ magnitude) with linear concepts, suggesting a more complex, non-linear reward geometry.

5. Significance

• Paradigm Shift: The paper argues that researchers must stop treating observational attribution as a proxy for causal importance in reward models. The gap between the two is structural and must be measured, not ignored.
• Safety Tooling: By providing tools to detect "distortion" (under-covered dimensions) and "misalignment cascades," the library enables in-training monitoring. This shifts reward model auditing from a post-hoc activity to a real-time safeguard during the RLHF training loop.
• Editability: The framework identifies specific features and concepts that drive reward signals, paving the way for interpretability-guided reward model editing (e.g., suppressing a "sycophancy" concept from the reward head) before downstream policy training.
• Open Science: By releasing reward-lens, the authors democratize access to the internal mechanics of reward models, a critical step for ensuring the safety of frontier AI systems.

In conclusion, reward-lens fills a critical void in AI safety by adapting mechanistic interpretability for the scalar-output nature of reward models, revealing that their internal logic is often counter-intuitive (late attribution $\neq$ early causality) and providing the tools necessary to diagnose and fix alignment failures before they propagate to policy models.