Sintesi Tecnica: Allineamento Sicuro al Tempo di Inferenza tramite l'Aumento del Reward Lagrangiano (LARA)

Definizione del Problema

Il Reinforcement Learning from Human Feedback (RLHF) è il framework standard per allineare i Large Language Models (LLM) alle preferenze umane. Tuttavia, l'allineamento tradizionale si basa sul fine-tuning della policy (ad esempio, tramite RLHF o Direct Preference Optimization come Dano DPO), il che è computazionalmente costoso e richiede un nuovo addestramento quando i requisiti di sicurezza o le popolazioni di utenti cambiano.

Sebbene i metodi di allineamento al tempo di inferenza (come il reranking Best-of-N o la decodifica guidata dal reward a livello di token) siano emersi per guidare modelli congelati utilizzando segnali di reward ausiliari senza aggiornamenti dei pesi, essi mancano attualmente di un meccanismo principato per l'allineamento sicuro. I metodi esistenti ottimizzano tipicamente un singolo punteggio scalare. Quando gli obiettivi di utilità (helpfulness) e sicurezza (harmlessness) entrano in conflitto, i professionini devono regolare manualmente il compromesso tra di essi. Inoltre, questi metodi non impongono intrinsecamente vincoli espliciti di sicurezza (ad esempio, un budget sul costo atteso) durante la decodifica, portando a output o troppo insicuri o con una riduzione non necessaria dell'utilità.

Metodologia: Aumento del Reward Lagrangiano (LARA)

Gli autori propongono LARA, un framework decoder-agnostic che trasferisce la struttura di ottimizzazione vincolata del Safe RLHF alla fase di inferenza. Invece di aggiornare i pesi del modello, LARA calibra una variabile duale scalare per creare un segnale di reward aumentato.

1. Formulazione Teorica

Il metodo inizia con l'obiettivo del Safe RLHF, che massimizza l'utilità attesa ($r$) soggetto a un vincolo di costo atteso ($c \le \tau$), regolarizzato da un termine di divergenza KL rispetto a una policy di riferimento ($\pi_{ref}$):
$$\max_{\pi} \mathbb{E}[r(x,y)] - \beta D_{KL}(\pi \parallel \pi_{ref}) \quad \text{s.t.} \quad \mathbb{E}[c(x,y)] \le \tau$$

Utilizzando la dualità Lagrangiana, gli autori dualizzano il vincolo di sicurezza. Questo trasforma il problema di ottimizzazione della policy a dimensione infinita in un problema di ottimizzazione convessa a una dimensione sulla variabile duale non negativa $\lambda$ (il moltiplicatore di Lagrange).

• Reward Aumentato: La variabile duale definisce un reward aumentato: $r_\lambda(x, y) = r(x, y) - \lambda c(x, y)$.
• Obiettivo Duale: Si dimostra che la funzione duale $g(\lambda)$ è convessa. La sua derivata corrisponde alla differenza tra il budget di sicurezza $\tau$ e il costo atteso sotto la policy Gibbs-tilted indotta da $r_\lambda$.
• Ottimalità: La $\lambda^*$ ottimale soddisfa la condizione di complementarità slack, garantendo che il costo atteso sia uguale al budget $\tau$ (se il vincolo è attivo) massimizzando al contempo l'utilità.

2. Calibrazione Empirica

LARA non richiede la risoluzione analitica dell'ottimizzazione. Invece, stima $\lambda^*$ tramite calibrazione empirica su un piccolo dataset:

1. Dati: Un piccolo set di prompt viene campionato e molteplici risposte vengono generate dalla policy di riferimento congelata.
2. Stima: La variabile duale viene stimata minimizzando l'obiettivo duale empirico su un intervallo compatto $[0, \Lambda]$.
3. Algoritmo: Grazie alla convessità e alla monotonicità del gradiente duale, la $\lambda^*$ viene trovata efficientemente usando la ricerca binaria.
4. Robustezza: Per gestire il disallineamento distributivo tra calibrazione e deployment, gli autori suggeriscono di fare il bootstrap del set di calibrazione e selezionare un quantile superiore della distribuzione di $\lambda$ stimata come scelta conservativa.

3. Integrazione con la Decodifica

Una volta calibrato, $\hat{\lambda}$ viene utilizzato per calcolare il reward aumentato $r_{\hat{\lambda}}$. Questo punteggio viene passato come segnale "drop-in" ai metodi esistenti di allineamento al tempo di inferenza:

• A livello di sequenza (es. Best-of-N): Il $\lambda$ calibrato fornisce una soluzione esatta al problema del costo atteso vincolato.
• A livello di token (es. Ricerca Guidata dal Reward): Il $\lambda$ calibrato fornisce un'euristica principata derivata dall'obiettivo vincolato, sebbene non garantisca il soddisfacimento esatto del vincolo a livello di popolazione a causa della natura greedy della decodifica a livello di token.

Contributi Chiave

1. Formulazione: Il paper formula l'allineamento sicuro al tempo di inferenza come il duale dell'obiettivo Safe RLHF, riducendo il problema a un compito di calibrazione convessa a una dimensione.
2. Geometria ed Efficienza: Gli autori caratterizzano la geometria dell'obiettivo duale, provandone la convessità e la monotonicità del gradiente. Ciò consente una procedura di calibrazione tramite ricerca binaria efficiente.
3. Garanzie Teoriche: Viene fornita una garanzia a campioni finiti per lo stimatore $\hat{\lambda}_{N,K}$, mostrando che l'errore si concentra attorno al minimizzatore della popolazione $\lambda^*$, controllato dal numero di prompt ($N$) e dai campioni per prompt ($K$).
4. Framework Decoder-Agnostic: LARA è progettato come uno strato di calibrazione compatibile con qualsiasi metodo esistente di allineamento al tempo di inferenza, sostituendo il tuning manuale dello scalare con una costruzione principata basata su vincoli di sicurezza.

Risultati Sperimentali

Gli autori hanno valutato LARA su modelli Llama-3.2-3B e Qwen-2.5-3B utilizzando il dataset BeaverTails.

• Soddisfacimento del Vincolo: Gli esperimenti hanno confermato che la $\lambda^*$ calibrata si trova al confine del vincolo di sicurezza. Valori di $\lambda < \lambda^*$ violavano il budget di sicurezza, mentre $\lambda > \lambda^*$ soddisfacevano il vincolo ma riducevano inutilmente l'utilità.
• Performance vs. Baseline:
  • Tempo di inferenza: Tra i metodi al tempo di inferenza, Best-of-N (BoN) combinato con LARA ha ottenuto il miglior compromesso tra utilità e innocuità. Ha superato SafeDPO in termini di utilità pur rimanendo competitivo sull'innocuità, il tutto senza aggiornamenti dei pesi.
  • Livello di Token: I metodi a livello di token (ARGS, PARGS, CD-Q) sono stati generalmente peggiori in termini di utilità rispetto a BoN, sebbene CD-Q abbia mostrato forti miglioramenti dell'innocuità all'aumentare di $\lambda$.
  • vs. Fine-tuning: Sebbene il BoN guidato da LARA si avvicini alle performance dei baselines basati su fine-tuning (Safe-RLHF, SafeDPO), i metodi di fine-tuning hanno comunque raggiunto il miglior compromesso complessivo. Tuttavia, LARA offre un vantaggio significativo in termini di flessibilità di deployment e costi.
• Valutazione: I risultati sono stati coerenti sia nelle valutazioni basate su modelli (usando i modelli di reward/costo di BeaverTails) sia nelle valutazioni basate su GPT-4o-mini, sebbene GPT abbia mostrato meno sensibilità alle variazioni di $\lambda$ rispetto ai modelli di reward.

Significato e Rivendicazioni

Il paper sostiene che LARA fornisce un'alternativa principata al tuning manuale delle penalità per l'allineamento sicuro al tempo di inferenza. Dualizzando il vincolo di sicurezza, LARA traduce un budget di sicurezza in un obiettivo di decodifica senza modificare i pesi del modello.

• Per i campionatori a livello di sequenza: LARA offre un'interpretazione esatta del costo atteso, garantendo teoricamente il compromesso ottimale sotto l'obiettivo vincolato.
• Per i decoder a livello di token: Offre un'euristica principata che rispetta la geometria dell'obiettivo vincolato, anche se non può imporre strettamente il vincolo di popolazione.

Gli autori concludono che LARA riesce a trasferire i vincoli di sicurezza dall'allineamento al tempo di addestramento al controllo pratico durante la decodifica, migliorando il compromesso tra utilità e innocuità e avvicinando i metodi al tempo di inferenza agli approcci basati sul fine-tuning. Il lavoro sottolinea che i vincoli di sicurezza possono essere gestiti efficacemente al tempo di inferenza, evitando l'overhead computazionale di ripetuti addestramenti.