reward-lens: A Mechanistic Interpretability Library for Reward Models

1. Enunciato del Problema

L'Apprendimento per Rinforzo dal Feedback Umano (RLHF) si basa sui Modelli di Ricompensa (RM) per modellare il comportamento dei Modelli Linguistici di Grande Dimensione (LLM). Se un modello di ricompensa codifica preferenze errate, la politica risultante eredita e amplifica tali errori. Sebbene la valutazione comportamentale degli RM (ad es. RewardBench) sia matura, l'interpretabilità meccanicistica per gli RM è praticamente inesistente.

Gli strumenti esistenti di interpretabilità (ad es. TransformerLens, Logit Lens, Direct Logit Attribution) sono progettati per LLM generativi. Si basano sul proiettare gli stati intermedi del modello sulla matrice di disincorporazione del vocabolario ($W_U$) per analizzare le distribuzioni dei token. Tuttavia, i modelli di ricompensa sostituiscono $W_U$ con un testa di regressione scalare ($w_r$). Questa differenza strutturale rende inapplicabili le primitive standard, creando un "divario strutturale" in cui i meccanismi interni dei modelli di ricompensa critici per la sicurezza rimangono una scatola nera.

2. Metodologia: La Prospettiva Centrata su $w_r$

Il documento introduce reward-lens, una libreria open-source che adatta le primitive di interpretabilità meccanicistica per i modelli di ricompensa. La intuizione teorica fondamentale è che il vettore dei pesi della testa di ricompensa, $w_r$, funge da asse naturale per tutte le domande di interpretabilità.

La libreria si fonda sull'osservazione che la ricompensa finale $r$ è un funzionale lineare del flusso residuo finale $h_{final}$:
$$r = w_r^\top h_{final} + b_r$$
Poiché il flusso residuo si decompone additivamente attraverso i livelli ($h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$), la ricompensa può essere decomposta esattamente come:
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

Sulla base di ciò, la libreria implementa le seguenti primitive sostituendo $W_U$ con $w_r$:

• Reward Lens: Proietta i flussi residui intermedi al livello $\ell$ su $w_r$ per tracciare la "cristallizzazione" dei segnali di preferenza attraverso i livelli.
• Attribuzione dei Componenti: Calcola il contributo di ogni sottolivello (Attenzione/MLP) alla ricompensa scalare finale tramite $w_r^\top \cdot \text{contributo}$.
• Patching Attivazionale Contrastivo: Scambia le attivazioni tra completamenti preferiti e non preferiti e misura la variazione in $w_r^\top h_{final}$. Questo è nativo degli RM grazie alla struttura contrastiva intrinseca dei dati di preferenza.
• Attribuzione delle Caratteristiche SAE: Utilizza Autoencoder Sparsi (SAE) per decomporre i flussi residui in caratteristiche, calcolando l'allineamento delle caratteristiche come $w_r^\top d_i$ (dove $d_i$ è un vettore del dizionario).

Indipendenza dall'Architettura: La libreria utilizza un approccio nativo HuggingFace con un protocollo di adattatore a 10 metodi. Ciò le consente di avvolgere modelli standard (Llama, Mistral, Gemma-2) e modelli complessi multi-obiettivo (come ArmoRM con gating MoE) senza chirurgia sui pesi.

Estensioni Basate sulla Teoria: La libreria trasforma in codice recenti risultati della teoria dell'allineamento:

• Indice di Distorsione: Prevede quali dimensioni di qualità sono sottocoperte dalle sonde di valutazione (rischio di reward hacking).
• Patching Consapevole della Divergenza: Rileva quando il patching induce attivazioni fuori distribuzione.
• Rilevatore di Cascata di Disallineamento: Identifica modalità di fallimento correlate tra diverse dimensioni di disallineamento.
• Analizzatore di Conflitto tra Termini di Ricompensa: Verifica se diversi obiettivi di ricompensa sono allineati, ortogonali o in conflitto.
• Analisi dei Vettori di Concetto: Estrae direzioni concettuali lineari e ne misura la risposta causale dose-dipendente sulla ricompensa.

3. Contributi Chiave

1. La Libreria (reward-lens): Il primo toolkit installabile e indipendente dall'architettura per l'interpretabilità meccanicistica dei modelli di ricompensa. Unifica le prospettive osservazionali (attribuzione) e causali (patching) sotto un unico quadro algebrico.
2. Quadro Unificante: Stabilisce che ogni primitiva di interpretabilità utile per un RM è una proiezione su o una decomposizione lungo $w_r$.
3. Validazione Empirica: Applicazione sistematica a due modelli di produzione: Skywork-Reward-Llama-3.1-8B-v0.2 (testa a singolo scalare) e ArmoRM-Llama3-8B-v0.1 (testa multi-obiettivo).
4. Riscontro Empirico Negativo: Il documento riporta esplicitamente che l'attribuzione lineare non predice l'importanza causale nei modelli di ricompensa.

4. Risultati e Scoperte Chiave

Gli autori hanno valutato la libreria su $\sim$695 coppie di preferenze RewardBench per modello.

• Profondità di Cristallizzazione:
  • Skywork: Le preferenze si cristallizzano uniformemente tardi (profondità frazionaria del livello $\approx 0.90$), suggerendo un meccanismo unificato assemblato vicino all'output.
  • ArmoRM: Le preferenze si cristallizzano prima e con maggiore varianza, riflettendo la sua architettura multi-obiettivo con gating MoE.
• Attribuzione vs Causalità (Il Divario di Fedeltà):
  • Attribuzione Osservazionale: Gli ultimi livelli MLP (in particolare mlp_L31) dominano il segnale di ricompensa in entrambi i modelli.
  • Patching Causale: I componenti dei livelli iniziali (ad es. mlp_L0, attn_L0) sono spesso i più critici per la ricompensa finale.
  • Correlazione: La correlazione di rango di Spearman tra l'entità dell'attribuzione e l'effetto del patching è negativa o vicina allo zero (Skywork: $\rho = -0.256$; ArmoRM: $\rho = -0.027$).
  • Interpretazione: Questo non è un bug ma una caratteristica dei flussi residui ridondanti. I livelli tardivi possono essere ablazionati perché percorsi paralleli compensano, mentre i livelli iniziali sono "portanti".
• Profili di Reward Hacking:
  • Skywork penalizza lunghezza, confidenza e bias di formattazione (a differenza dei tipici modelli RLHF che li premiano).
  • ArmoRM premia confidenza e formattazione ma penalizza la sycophancy (adulazione).
• Allineamento dei Concetti:
  • Skywork mostra un forte allineamento lineare con concetti come "accordo" (pendenza $\approx 1.0$).
  • ArmoRM mostra un allineamento debole ($\approx 1\%$ di entità) con concetti lineari, suggerendo una geometria di ricompensa più complessa e non lineare.

5. Significato

• Cambio di Paradigma: Il documento sostiene che i ricercatori devono smettere di trattare l'attribuzione osservazionale come un proxy per l'importanza causale nei modelli di ricompensa. Il divario tra i due è strutturale e deve essere misurato, non ignorato.
• Strumentazione per la Sicurezza: Fornendo strumenti per rilevare la "distorsione" (dimensioni sottocoperte) e le "cascate di disallineamento", la libreria abilita il monitoraggio durante l'addestramento. Questo sposta l'audit dei modelli di ricompensa da un'attività post-hoc a una salvaguardia in tempo reale durante il ciclo di addestramento RLHF.
• Modificabilità: Il quadro identifica caratteristiche e concetti specifici che guidano i segnali di ricompensa, aprendo la strada alla modifica guidata dall'interpretabilità dei modelli di ricompensa (ad es. sopprimere un concetto di "adulazione" dalla testa di ricompensa) prima dell'addestramento della politica a valle.
• Scienza Aperta: Rilasciando reward-lens, gli autori democratizzano l'accesso ai meccanismi interni dei modelli di ricompensa, un passo critico per garantire la sicurezza dei sistemi AI all'avanguardia.

In conclusione, reward-lens colma un vuoto critico nella sicurezza dell'IA adattando l'interpretabilità meccanicistica alla natura a output scalare dei modelli di ricompensa, rivelando che la loro logica interna è spesso controintuitiva (attribuzione tardiva $\neq$ causalità precoce) e fornendo gli strumenti necessari per diagnosticare e correggere i fallimenti di allineamento prima che si propaghino ai modelli di politica.