Reward Models Can Improve Themselves: Reward-Guided Adversarial Failure Mode Discovery for Robust Reward Modeling

1. Il Problema

La modellazione della ricompensa (Reward Modeling - RM) è fondamentale per allineare i Large Language Models (LLM) alle preferenze umane, utilizzata in compiti come il fine-tuning, il filtraggio delle risposte e il ranking. Tuttavia, i modelli di ricompensa esistenti soffrono di due problemi principali:

• Fragilità: Tendono a fallire sotto perturbazioni distribuzionali (Out-of-Distribution - OOD) o perturbazioni avversarie, assegnando punteggi errati a risposte che dovrebbero essere preferite o non preferite (fenomeno noto come reward hacking o mis-specification).
• Limitazioni delle metodologie attuali: Gli approcci esistenti per identificare queste modalità di fallimento (es. generazione di controfattuali basata su attributi) richiedono conoscenze preliminari sulle distribuzioni delle preferenze o sugli attributi rilevanti. Questo li rende poco pratici in scenari reali dove tali informazioni non sono disponibili e non generalizzano bene a diversi modelli o distribuzioni.

L'obiettivo della ricerca è rispondere alla domanda: È possibile identificare in modo trattabile le modalità di fallimento dei modelli di ricompensa senza conoscere a priori la distribuzione delle preferenze, e utilizzare queste informazioni per migliorare la robustezza del modello stesso?

2. Metodologia: Il Framework REFORM

Gli autori propongono REFORM, un framework di auto-miglioramento che utilizza il modello di ricompensa stesso per guidare la scoperta delle proprie vulnerabilità. Il processo si articola in tre fasi principali:

A. Scoperta di Modalità di Fallimento tramite Decodifica Controllata

Invece di affidarsi a prompt esterni o attributi predefiniti, REFORM utilizza una decodifica controllata guidata dalla ricompensa.

• Obiettivo: Generare risposte che appartengono chiaramente a una classe di preferenza nota (es. una risposta "preferita" o "non preferita") ma che il modello di ricompensa attuale classifica erroneamente (es. assegna un punteggio basso a una risposta preferita).
• Meccanismo:
  • Si parte da una policy allineata (per generare risposte preferite) o disallineata (per risposte non preferite).
  • Durante la generazione token per token, si seleziona il token successivo tra i top-k candidati della policy, ma scegliendo quello che minimizza (per le risposte preferite) o massimizza (per le risposte non preferite) il punteggio dato dal modello di ricompensa.
  • Viene introdotta una regolarizzazione basata sulla likelihood del modello linguistico per garantire che le risposte generate rimangano fluide e semanticamente coerenti, evitando output incoerenti.
• Risultato: Si ottengono esempi avversari "class-consistenti ma reward-inconsistenti" (es. una risposta innocua che il modello penalizza ingiustamente).

B. Augmentation dei Dati e Fine-tuning

Una volta identificati questi fallimenti:

• Vengono selezionati i punti di addestramento più influenti (definiti come quelli con il minor loss di Bradley-Terry, ovvero i casi su cui il modello è più "sicuro" ma che potrebbero nascondere vulnerabilità sistemiche).
• Si generano varianti di fallimento per il 5% di questi campioni.
• Si crea un dataset aumentato combinando i dati originali con le nuove coppie di preferenze corrette (dove la risposta preferita riceve effettivamente un punteggio più alto rispetto alla variante avversaria).
• Il modello di ricompensa viene ri-addestrato (o fine-tuned) su questo dataset combinato.

C. Valutazione della Robustezza

Il modello risultante viene testato non solo sulla sua accuratezza originale, ma anche sulla sua resistenza a perturbazioni specifiche (verbosità eccessiva, capitalizzazione di parole dannose, ripetizioni, errori di ortografia).

3. Contributi Chiave

1. Metodo Agnostico alla Distribuzione: REFORM non richiede conoscenze preliminari sugli attributi che influenzano la ricompensa. Utilizza il modello stesso come guida per la ricerca delle vulnerabilità, rendendolo applicabile a qualsiasi modello di ricompensa.
2. Auto-Miglioramento: Dimostra che un modello di ricompensa può utilizzare le proprie modalità di fallimento scoperte per migliorare la propria robustezza senza intervento umano massivo o ri-etichettatura estensiva.
3. Efficienza dei Dati: Il metodo richiede l'aggiunta di solo il 5% di dati aumentati rispetto al dataset originale per ottenere miglioramenti significativi, superando metodi precedenti che richiedevano dataset di addestramento molto più grandi (es. RRM).
4. Preservazione delle Prestazioni: Il framework migliora la robustezza senza degradare la qualità del modello di ricompensa sui dati non perturbati, né compromettere la qualità dell'allineamento nelle fasi successive.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su due dataset di preferenze ampiamente utilizzati: Anthropic Helpful-Harmless (HH) e PKU Beavertails, utilizzando modelli come Mistral 7B e Qwen 2.5 14B.

• RQ1 (Scoperta dei Fallimenti): REFORM ha dimostrato di essere significativamente più efficace dei baselines (perturbazioni casuali o basate su attributi) nell'identificare modalità di fallimento per le risposte preferite, mantenendo alta la leggibilità e l'appropriatezza semantica.
• RQ2 (Robustezza): I modelli di ricompensa addestrati con REFORM hanno mostrato una resistenza superiore alle perturbazioni OOD (verbosità, capitalizzazione, ripetizione, errori di ortografia). La caduta nel tasso di vittoria (win rate drop) è stata notevolmente inferiore rispetto ai modelli originali o addestrati con altri metodi di robustezza.
• RQ3 (Qualità e Allineamento):
  • Accuratezza: La capacità del modello di ricompensa di valutare correttamente le risposte non perturbate è rimasta stabile (con un calo minimo trascurabile).
  • Allineamento Downstream: L'uso di REFORM in strategie di allineamento come Best-of-N (BoN), PPO e DPO ha prodotto policy che mantengono alta la qualità (leggibilità, utilità) e la diversità.
  • Rimozione di Correlazioni Spurie: È stato osservato che REFORM aiuta il modello a non imparare correlazioni spurie (es. penalizzare risposte solo perché sono lunghe o contengono certe parole chiave), migliorando la qualità dell'allineamento finale.

5. Significato e Implicazioni

Questo lavoro è significativo perché sposta il paradigma dalla semplice identificazione dei fallimenti alla loro correzione attiva e automatizzata.

• Sicurezza: Migliora la sicurezza degli LLM rendendo i modelli di ricompensa meno suscettibili a manipolazioni superficiali (es. scrivere parole dannose in maiuscolo per ingannare il filtro).
• Scalabilità: Offre un metodo scalabile per migliorare la robustezza dei modelli di ricompensa senza dipendere da annotazioni umane costose o da conoscenze di dominio specifiche.
• Futuro: Suggerisce che la robustezza può essere ottenuta attraverso l'addestramento su esempi di "fallimento controllato", aprendo la strada a sistemi di allineamento più resilienti e affidabili in ambienti reali dinamici.

In sintesi, REFORM dimostra che i modelli di ricompensa possono "imparare dai propri errori" in modo autonomo, diventando più robusti e affidabili senza sacrificare le loro prestazioni fondamentali.