reward-lens: A Mechanistic Interpretability Library for Reward Models

1. Problemstellung

Reinforcement Learning from Human Feedback (RLHF) verlässt sich auf Reward Models (RMs), um das Verhalten von Large Language Models (LLMs) zu steuern. Wenn ein Reward Model falsche Präferenzen kodiert, erbt und verstärkt die daraus resultierende Policy diese Fehler. Während die verhaltensbasierte Evaluation von RMs (z. B. RewardBench) ausgereift ist, ist die mechanistische Interpretierbarkeit für RMs praktisch nicht existent.

Bestehende Interpretierbarkeitswerkzeuge (z. B. TransformerLens, Logit Lens, Direct Logit Attribution) sind für generative LLMs konzipiert. Sie stützen sich darauf, intermediate Modellzustände auf die Unembedding-Matrix des Vokabulars ($W_U$) zu projizieren, um Token-Verteilungen zu analysieren. Reward Models ersetzen $W_U$ jedoch durch einen skalaren Regressionskopf ($w_r$). Dieser strukturelle Unterschied macht Standard-Primitiven unanwendbar und schafft eine „strukturelle Lücke", in der die internen Mechanismen sicherheitskritischer Reward Models eine Black Box bleiben.

2. Methodik: Die $w_r$-zentrierte Sichtweise

Das Paper stellt reward-lens vor, eine Open-Source-Bibliothek, die mechanistische Interpretierbarkeits-Primitiven an Reward Models anpasst. Die zentrale theoretische Erkenntnis ist, dass der Gewichtsvektor des Reward-Kopfes, $w_r$, die natürliche Achse für alle Fragen der Interpretierbarkeit darstellt.

Die Bibliothek basiert auf der Beobachtung, dass die finale Belohnung $r$ ein lineares Funktional des finalen Residualstreams $h_{final}$ ist:
$$r = w_r^\top h_{final} + b_r$$
Da sich der Residualstream additiv über die Schichten zerlegen lässt ($h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$), lässt sich die Belohnung exakt zerlegen als:
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

Basierend darauf implementiert die Bibliothek folgende Primitiven, indem sie $W_U$ durch $w_r$ ersetzt:

• Reward Lens: Projiziert intermediate Residualstreams an Schicht $\ell$ auf $w_r$, um die „Kristallisation" von Präferenzsignalen über die Schichten hinweg zu verfolgen.
• Component Attribution: Berechnet den Beitrag jeder Subschicht (Attention/MLP) zum finalen skalaren Reward via $w_r^\top \cdot \text{contribution}$.
• Contrastive Activation Patching: Tauscht Aktivierungen zwischen bevorzugten und nicht-bevorzugten Completions aus und misst die Änderung in $w_r^\top h_{final}$. Dies ist für RMs aufgrund der inhärenten kontrastiven Struktur von Präferenzdaten nativ gegeben.
• SAE Feature Attribution: Nutzt Sparse Autoencoder (SAEs), um Residualstreams in Features zu zerlegen und die Feature-Ausrichtung als $w_r^\top d_i$ zu berechnen (wobei $d_i$ ein Wörterbuchvektor ist).

Architektur-Unabhängigkeit: Die Bibliothek verwendet einen HuggingFace-nativen Ansatz mit einem 10-Methoden-Adapter-Protokoll. Dies ermöglicht das Einpacken (Wrapping) standardmäßiger Modelle (Llama, Mistral, Gemma-2) und komplexer multi-objektiver Modelle (wie ArmoRM mit MoE-Gating) ohne Gewichts-Chirurgie.

Theoriebasierte Erweiterungen: Die Bibliothek operationalisiert neuere Ergebnisse der Alignment-Theorie in Code:

• Distortion Index: Sagt voraus, welche Qualitätsdimensionen von Evaluationsproben nicht abgedeckt sind (Risiko von Reward Hacking).
• Divergence-Aware Patching: Erkennt, wenn Patching Aktivierungen außerhalb der Verteilung (out-of-distribution) induziert.
• Misalignment Cascade Detector: Identifiziert korrelierte Fehlermodi über verschiedene Misalignment-Dimensionen hinweg.
• Reward-Term Conflict Analyzer: Prüft, ob verschiedene Reward-Ziele ausgerichtet, orthogonal oder in Konflikt stehen.
• Concept-Vector Analysis: Extrahiert lineare Konzeptrichtungen und misst deren kausale Dosis-Wirkungs-Beziehung auf den Reward.

3. Hauptbeiträge

1. Die Bibliothek (reward-lens): Das erste installierbare, architekturunabhängige Toolkit für die mechanistische Interpretierbarkeit von Reward Models. Es vereint beobachtende (Attribution) und kausale (Patching) Perspektiven unter einem einzigen algebraischen Rahmen.
2. Vereinheitlichender Rahmen: Es etabliert, dass jedes nützliche Interpretierbarkeits-Primitiv für ein RM eine Projektion auf oder eine Zerlegung entlang von $w_r$ ist.
3. Empirische Validierung: Systematische Anwendung auf zwei Produktionsmodelle: Skywork-Reward-Llama-3.1-8B-v0.2 (einzelner skalärer Kopf) und ArmoRM-Llama3-8B-v0.1 (multi-objektiver Kopf).
4. Negatives empirisches Ergebnis: Das Paper berichtet explizit, dass lineare Attribution die kausale Wichtigkeit in Reward Models nicht vorhersagt.

4. Hauptergebnisse & Erkenntnisse

Die Autoren evaluierten die Bibliothek an $\sim$695 RewardBench-Präferenzpaaren pro Modell.

• Kristallisationstiefe:
  • Skywork: Präferenzen kristallisieren einheitlich spät aus (fraktionale Schichttiefe $\approx 0,90$), was auf einen vereinheitlichten Mechanismus hindeutet, der nahe am Output zusammengesetzt wird.
  • ArmoRM: Präferenzen kristallisieren früher und mit höherer Varianz aus, was seine multi-objektive, MoE-gated Architektur widerspiegelt.
• Attribution vs. Kausalität (Die Faithfulness-Lücke):
  • Beobachtende Attribution: Die finalen MLP-Schichten (insbesondere mlp_L31) dominieren das Rewardsignal in beiden Modellen.
  • Kausales Patching: Komponenten früher Schichten (z. B. mlp_L0, attn_L0) sind oft am kritischsten für den finalen Reward.
  • Korrelation: Die Spearman-Rangkorrelation zwischen Attribution-Magnitude und Patching-Effekt ist negativ oder nahe null (Skywork: $\rho = -0,256$; ArmoRM: $\rho = -0,027$).
  • Interpretation: Dies ist kein Fehler, sondern eine Eigenschaft redundanter Residualstreams. Späte Schichten können ablatiert werden, da parallele Pfade kompensieren, während frühe Schichten „tragend" sind.
• Reward-Hacking-Profile:
  • Skywork bestraft Länge, Konfidenz und Formatierungs-Bias (im Gegensatz zu typischen RLHF-Modellen, die diese belohnen).
  • ArmoRM belohnt Konfidenz und Formatierung, bestraft aber Schmeichelei (Sycophancy).
• Konzept-Ausrichtung:
  • Skywork zeigt eine starke lineare Ausrichtung mit Konzepten wie „Zustimmung" (Steigung $\approx 1,0$).
  • ArmoRM zeigt eine schwache Ausrichtung ($\approx 1\%$ Magnitude) mit linearen Konzepten, was auf eine komplexere, nicht-lineare Reward-Geometrie hindeutet.

5. Bedeutung

• Paradigmenwechsel: Das Paper argumentiert, dass Forscher aufhören müssen, beobachtende Attribution als Proxy für kausale Wichtigkeit in Reward Models zu behandeln. Die Lücke zwischen beiden ist strukturell und muss gemessen, nicht ignoriert werden.
• Sicherheitswerkzeuge: Durch die Bereitstellung von Werkzeugen zur Erkennung von „Verzerrung" (nicht abgedeckte Dimensionen) und „Misalignment-Kaskaden" ermöglicht die Bibliothek ein Monitoring während des Trainings. Dies verlagert die Auditing von Reward Models von einer nachgelagerten Aktivität zu einem Echtzeit-Schutzmechanismus während der RLHF-Trainings-Schleife.
• Bearbeitbarkeit: Das Framework identifiziert spezifische Features und Konzepte, die Rewardsignale antreiben, und ebnet den Weg für interpretierbarkeitsgesteuerte Bearbeitung von Reward Models (z. B. Unterdrückung eines „Schmeichelei"-Konzepts aus dem Reward-Kopf), bevor das downstream Policy-Training beginnt.
• Open Science: Durch die Veröffentlichung von reward-lens demokratisieren die Autoren den Zugang zu den internen Mechanismen von Reward Models, ein kritischer Schritt zur Sicherstellung der Sicherheit von Frontier-AI-Systemen.

Zusammenfassend füllt reward-lens eine kritische Lücke in der KI-Sicherheit, indem es mechanistische Interpretierbarkeit an die skalare-output-Natur von Reward Models anpasst, offenbart, dass ihre interne Logik oft kontraintuitiv ist (späte Attribution $\neq$ frühe Kausalität), und die notwendigen Werkzeuge bereitstellt, um Alignment-Fehler zu diagnostizieren und zu beheben, bevor sie sich auf Policy-Modelle fortpflanzen.