reward-lens: A Mechanistic Interpretability Library for Reward Models

1. Probleemstelling

Reinforcement Learning from Human Feedback (RLHF) is afhankelijk van Reward Models (RMs) om het gedrag van Large Language Models (LLMs) te sturen. Als een reward model incorrecte voorkeuren encodeert, erft het resulterende beleid die fouten over en versterkt ze. Waar de gedragsmatige evaluatie van RMs (bijv. RewardBench) volwassen is, bestaat mechanistische interpreteerbaarheid voor RMs vrijwel niet.

Bestaande interpreteerbaarheidstools (zoals TransformerLens, Logit Lens, Direct Logit Attribution) zijn ontworpen voor generatieve LLM's. Ze vertrouwen op het projecteren van tussenliggende modeltoestanden op de vocabulaire-unembeddingmatrix ($W_U$) om tokenverdelingen te analyseren. Reward models vervangen $W_U$ echter door een scalar regression head ($w_r$). Dit structurele verschil maakt standaardprimitieven onbruikbaar, wat een "structurele kloof" creëert waar de interne werking van veiligheidskritieke reward models een black box blijft.

2. Methodologie: De $w_r$-gecentreerde Visie

Het paper introduceert reward-lens, een open-source bibliotheek die mechanistische interpreteerbaarheidsprimitieven aanpast voor reward models. Het kerntheoretische inzicht is dat het gewichtsvector van de reward head, $w_r$, fungeert als de natuurlijke as voor alle interpreteerbaarheidsvragen.

De bibliotheek is gebaseerd op de observatie dat de uiteindelijke beloning $r$ een lineaire functionaal is van de uiteindelijke residual stream $h_{final}$:
$$r = w_r^\top h_{final} + b_r$$
Omdat de residual stream additief over lagen decomposeert ($h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$), kan de beloning exact worden ontbonden als:
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

Gebaseerd hierop implementeert de bibliotheek de volgende primitieven door $W_U$ te vervangen door $w_r$:

• Reward Lens: Projecteert tussenliggende residual streams op laag $\ell$ op $w_r$ om de "kristallisatie" van voorkeurssignalen over de lagen te volgen.
• Component Attribution: Berekent de bijdrage van elke sublaag (Attention/MLP) aan de uiteindelijke scalar beloning via $w_r^\top \cdot \text{bijdrage}$.
• Contrastive Activation Patching: Vervangt activaties tussen geprefereerde en niet-geprefereerde voltooiingen en meet de verandering in $w_r^\top h_{final}$. Dit is native voor RMs vanwege de inherente contrastieve structuur van voorkeursdata.
• SAE Feature Attribution: Gebruikt Sparse Autoencoders (SAEs) om residual streams te decomponeren in features, waarbij feature-alignement wordt berekend als $w_r^\top d_i$ (waarbij $d_i$ een dictionary-vector is).

Architectuuronafhankelijkheid: De bibliotheek gebruikt een HuggingFace-native aanpak met een 10-methode adapterprotocol. Hiermee kan het standaardmodellen (Llama, Mistral, Gemma-2) en complexe multi-objectieve modellen (zoals ArmoRM met MoE-gating) omhullen zonder gewichtschirurgie.

Theoretisch Onderbouwde Uitbreidingen: De bibliotheek operationaliseert recente resultaten uit alignment-theorie in code:

• Distortion Index: Voorspelt welke kwaliteitsdimensies onderbelicht zijn door evaluatieprobes (risico op reward hacking).
• Divergence-Aware Patching: Detecteert wanneer patching activaties buiten de verdeling (out-of-distribution) induceert.
• Misalignment Cascade Detector: Identificeert gecorreleerde faalmodi over verschillende misalignment-dimensies.
• Reward-Term Conflict Analyzer: Controleert of verschillende reward-doelen gealigneerd, orthogonaal of in conflict zijn.
• Concept-Vector Analysis: Extrahert lineaire conceptrichtingen en meet hun causale dosis-respons op de beloning.

3. Belangrijkste Bijdragen

1. De Bibliotheek (reward-lens): De eerste installeerbare, architectuuronafhankelijke toolkit voor mechanistische interpreteerbaarheid van reward models. Het verenigt observationele (attribution) en causale (patching) visies onder één algebraïsch raamwerk.
2. Unificerend Kader: Het stelt vast dat elke nuttige interpreteerbaarheidsprimitief voor een RM een projectie op of ontbinding langs $w_r$ is.
3. Empirische Validatie: Systematische toepassing op twee productiemodellen: Skywork-Reward-Llama-3.1-8B-v0.2 (single-scalar head) en ArmoRM-Llama3-8B-v0.1 (multi-objective head).
4. Negatief Empirisch Vondst: Het paper rapporteert expliciet dat lineaire attribution de causale belangrijkheid in reward models niet voorspelt.

4. Belangrijkste Resultaten & Vondsten

De auteurs evalueerden de bibliotheek op ongeveer 695 RewardBench-voorkeursparen per model.

• Kristallisatiediepte:
  • Skywork: Voorkeuren kristalliseren uniform laat (fractionele laagdiepte $\approx 0.90$), wat wijst op een unificerend mechanisme dat dicht bij de output wordt samengesteld.
  • ArmoRM: Voorkeuren kristalliseren eerder en met hogere variantie, wat zijn multi-objectieve, MoE-gated architectuur weerspiegelt.
• Attribution vs. Causaliteit (De Faithfulness-kloof):
  • Observational Attribution: De laatste MLP-lagen (specifiek mlp_L31) domineren het beloningssignaal in beide modellen.
  • Causal Patching: Componenten uit vroege lagen (bijv. mlp_L0, attn_L0) zijn vaak het meest kritiek voor de uiteindelijke beloning.
  • Correlatie: De Spearman-rangcorrelatie tussen attribution-magnitude en patching-effect is negatief of dicht bij nul (Skywork: $\rho = -0.256$; ArmoRM: $\rho = -0.027$).
  • Interpretatie: Dit is geen bug maar een eigenschap van redundante residual streams. Late lagen kunnen worden geablateerd omdat parallelle paden compenseren, terwijl vroege lagen "dragerend" zijn.
• Reward Hacking-profielen:
  • Skywork straft lengte, zelfvertrouwen en opmaakbias (in tegenstelling tot typische RLHF-modellen die deze belonen).
  • ArmoRM belooft zelfvertrouwen en opmaak maar straft sycophancy.
• Concept Alignment:
  • Skywork toont sterke lineaire alignement met concepten zoals "akkoord" (helling $\approx 1.0$).
  • ArmoRM toont zwakke alignement ($\approx 1\%$ magnitude) met lineaire concepten, wat wijst op een complexere, niet-lineaire reward-geometrie.

5. Betekenis

• Paradigmaverschuiving: Het paper betoogt dat onderzoekers moeten stoppen met het behandelen van observationele attribution als een proxy voor causale belangrijkheid in reward models. De kloof tussen de twee is structureel en moet worden gemeten, niet genegeerd.
• Veiligheidstooling: Door tools te bieden om "distortion" (onderbelichte dimensies) en "misalignment cascades" te detecteren, maakt de bibliotheek monitoring tijdens het trainen mogelijk. Dit verschuift reward model auditing van een post-hoc activiteit naar een real-time veiligheidsmechanisme tijdens de RLHF-traininglus.
• Bewerkbaarheid: Het raamwerk identificeert specifieke features en concepten die beloningssignalen aandrijven, wat de weg effent voor interpreteerbaarheidsgeleide reward model editing (bijv. het onderdrukken van een "sycophancy"-concept uit de reward head) voordat downstream beleidstraining plaatsvindt.
• Open Science: Door reward-lens vrij te geven, democratiseren de auteurs de toegang tot de interne werking van reward models, een kritieke stap voor het waarborgen van de veiligheid van frontier AI-systemen.

Concluderend vult reward-lens een kritieke leemte in AI-veiligheid door mechanistische interpreteerbaarheid aan te passen aan de scalar-output aard van reward models. Het onthult dat hun interne logica vaak tegen-intuïtief is (late attribution $\neq$ vroege causaliteit) en levert de noodzakelijke tools om alignment-fouten te diagnosticeren en te verhelpen voordat ze zich voortplanten naar beleidsmodellen.