Sintesi Tecnica: Quando la Distanza Distrae: Bias della Distanza di Rappresentazione nella BT-Loss per i Modelli di Ricompensa

1. Definizione del Problema

I modelli di ricompensa (RM) sono centrali nel Reinforcement Learning from Human Feedback (RLHF), fungendo da proxy per le preferenze umane che guidano l'allineamento dei Large Language Model (LLM). L'obiettivo standard per l'addestramento di questi modelli è la perdita di Bradley-Terry (BT), che apprende da dati a coppie consistenti in risposte "scelte" ($y_w$) e "scartate" ($y_l$).

Gli autori identificano un difetto critico nella dinamica di aggiornamento della standard BT loss: il Bias della Distanza di Rappresentazione. Mentre il gradiente della BT è destinato ad essere guidato dall'errore di predizione (l'entità dell'errore nel ranking), l'analisi rivela che la norma del gradiente scala con due componenti distinte:

1. Errore di Predizione: La differenza nelle ricompense predette tra le risposte scelte e quelle scartate.
2. Distanza di Rappresentazione: La distanza tra le due risposte nello spazio di output dell'ultimo strato ($\|h_w - h_l\|$).

Questo accoppiamento crea un segnale di apprendimento spurio in cui l'entità dell'aggiornamento è dettata da quanto le rappresentazioni della coppia siano distanti, piuttosto che esclusivamente da quanto il modello stia sbagliando.

• Coppie a Piccola Distanza: Le coppie che sono semanticamente o strutturalmente simili (ad esempio, un frammento di codice corretto rispetto a uno con un singolo errore logico) hanno spesso piccole distanze di rappresentazione. Sotto la standard BT loss, queste coppie ricevono aggiornamenti estremamente deboli, anche quando il modello le classifica erroneamente. Questo è problematico perché le distinzioni fine-grained sono spesso le più critiche per compiti come il ragionamento e la programmazione.
• Coppie a Grande Distanza: Le coppie che sono vastamente diverse (ad esempio, un rifiuto rispetto a una risposta a una query dannosa) presentano naturalmente grandi distanze di rappresentazione. Queste coppie ricevono aggiornamenti sproporzionatamente forti, potenzialmente oscurando i segnali di apprendimento delle coppie più sottili e fine-grained.

Di conseguenza, il modello sotto-addestra sulle distinzioni fine-grained difficili dove incontra maggiori difficoltà, mentre sovrastima le differenze superficiali.

2. Metodologia: NormBT

Per affrontare questo bias, gli autori propongono NormBT, uno schema di normalizzazione a coppie adattivo che riscala il contributo del gradiente di ogni coppia di preferenza.

Base Teorica

Gli autori decompongono la norma del gradiente BT $\|\nabla_\theta L_{BT}\|$ e dimostrano che essa è proporzionale a:
$$\|\nabla_\theta L_{BT}\| \approx |\sigma(d) - 1| \cdot \|h_w - h_l\|$$
dove $d$ è la differenza di ricompensa e $\|h_w - h_l\|$ è la distanza di rappresentazione. Per rendere la dimensione dell'aggiornamento dipendente principalmente dal termine di errore di predizione $|\sigma(d) - 1|$, gli autori introducono un peso $w_i$ che scala inversamente con la distanza di rappresentazione.

Dettagli di Implementazione

Il calcolo diretto della norma esatta del gradiente per la pesatura è computazionalmente costoso. Invezione, NormBT utilizza la distanza di rappresentazione dell'ultimo strato ($\|h_w - h_l\|$) come proxy, che si è dimostrata un proxy esatto per lo score head lineare e una forte approssimazione per l'intero modello.

La funzione di perdita proposta è:
$$\mathcal{L}_{NormBT}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \tilde{w}(y_w, y_l) \cdot \log \sigma(r_w - r_l) \right]$$

Dove il peso $\tilde{w}$ è definito come:
$$\tilde{w}(y_w, y_l) = \frac{\mu_t}{\|h_w - h_l\| + \epsilon}$$

Scelte di design chiave includono:

• Stop-Gradient: Il termine di distanza è calcolato con un operatore di stop-gradient per evitare che il peso di normalizzazione stesso venga ottimizzato.
• Stabilizzazione EMA: Per prevenire il drift di scala durante l'addestramento (poiché le magnitudo degli embedding possono cambiare), gli autori utilizzano una Media Mobile Esponenziale (EMA) della media batch della distanza di rappresentazione ($\mu_t$) per normalizzare i pesi. Ciò garantisce che i pesi rimangano stabili e adattivi ai cambiamenti distribuzionali.
• Integrazione Leggera: NormBT non richiede cambiamenti architettonici, preserva la base probabilistica della BT e comporta un overhead computazionale trascurabile.

3. Contributi Chiave

1. Analisi del Gradiente: Il paper fornisce un'analisi dettagliata, sia teorica che empirica, mostrando che le norme del gradiente BT sono determinate congiuntamente dall'errore di predizione e dalla distanza di rappresentazione, portando a segnali di apprendimento distorti.
2. Identificazione del Bias: Gli autori dimostrano che questo accoppiamento causa una ricezione di aggiornamenti deboli per le coppie a piccola distanza (spesso le più critiche per i compiti fine-grained), mentre le coppie a grande distanza dominano il segnale di apprendimento.
3. Proposta di NormBT: Introducono NormBT, una modifica "drop-in" che normalizza gli aggiornamenti rispetto alla distanza di rappresentazione, assicurando che i segnali di apprendimento si concentrino sugli errori di predizione.
4. Validazione Empirica: Estesi esperimenti su diversi backbone di LLM (Gemma-2b-it, Llama-3.2-3B) e dataset (Unified-Feedback, Skywork-Reward) mostrano miglioramenti costanti delle prestazioni.

4. Risultati Sperimentali

Gli autori hanno valutato NormBT rispetto alla baseline standard BT e a diverse varianti (BT + Margin, BT + Label Smoothing) sul benchmark RewardBench e sui compiti di selezione Best-of-N (BoN) a valle.

• Prestazioni di RewardBench: NormBT ha superato costantemente la baseline BT in tutte le impostazioni. I guadagni più significativi sono stati osservati nella categoria Reasoning, dove l'accuratezza è migliorata mediamente di oltre il 5%. Ciò si allinea all'ipotesi che i compiti di Reasoning dipendano da distinzioni fine-grained (piccole distanze di rappresentazione) che vengono soppresse dalla BT standard.
• RLHF a valle (BoN): Negli esperimenti di selezione Best-of-N, i modelli addestrati con NormBT hanno ottenuto "gold scores" (qualità delle risposte selezionate) più elevati rispetto alle baseline BT.
• Studi di Ablazione:
  • NormBT ha superato altre tecniche di addestramento come LayerNorm e Gradient Clipping, che non sono riuscite ad affrontare il problema del peso relativo tra le coppie.
  • Il metodo si è dimostrato complementare, migliorando le prestazioni anche quando applicato sopra altre varianti di BT (ad esempio, NormBT + Margin).
  • La rimozione della stabilizzazione EMA o l'uso di metriche di similarità alternative (similarità del coseno, embedding con average-pooling) ha comportato un degrado delle prestazioni, validando le specifiche scelte di design.

5. Significato e Rivendicazioni

Il paper sostiene che NormBT offra una soluzione principata e leggera a un limite strutturale della standard BT loss. Decoppiando l'entità dell'aggiornamento dalla distanza di rappresentazione, NormBT assicura che il modello di ricompensa apprenda in modo più fedele dalle distinzioni fine-grained, che sono spesso le più difficili ma critiche per un allineamento di alta qualità.

Gli autori sottolineano che il loro lavoro non affronta direttamente la qualità dei dati o il bias degli annotatori, ma migliora l'efficienza e la fedeltà dei segnali di apprendimento estratti dai dati a coppia esistenti. Posizionano NormBT come un miglioramento fondamentale per il paradigma della modellazione della ricompensa discriminativa, portando potenzialmente a segnali di ricompensa più affidabili nelle applicazioni RLHF a valle, particolarmente in domini che richiedono ragionamento preciso e rilevamento di errori sottili.