Complementing reinforcement learning with SFT through logit averaging in the post training of LLMs

Riepilogo Tecnico: Complementare l'Apprendimento per Rinforzo con SFT tramite Media dei Logit

1. Enunciato del Problema

Il post-addestramento dei Large Language Models (LLM) si basa tipicamente su due paradigmi: Supervised Fine-Tuning (SFT) e Reinforcement Learning (RL) con ricompense verificabili (RLVR). Mentre l'SFT fornisce una supervisione densa per token per apprendere formattazione e stile, è limitato dalla distribuzione delle dimostrazioni di addestramento. Al contrario, i metodi RL come Proximal Policy Optimization (PPO) e Group Relative Policy Optimization (GRPO) possono superare questi limiti distributivi, ma soffrono di segnali di ricompensa sparsi e rumorosi.

L'approccio standard combina questi metodi in sequenza: l'SFT viene applicato per primo, seguito da RL con un termine di regolarizzazione Kullback-Leibler (KL), $\beta \text{KL}(\pi_\theta \| \pi_{\text{sft}})$, per ancorare la politica addestrabile $\pi_\theta$ al riferimento SFT $\pi_{\text{sft}}$. Gli autori identificano una tensione fondamentale in questa ricetta "ancorata al KL": la penalità KL trascina $\pi_\theta$ verso $\pi_{\text{sft}}$, potenzialmente vincolando la politica ai limiti del checkpoint SFT (ad esempio, errori algebrici, cicli di ripetizione) e impedendole di sfruttare appieno le capacità di ragionamento sviluppate durante l'RL. I praticanti spesso mitigano questo problema scegliendo un $\beta$ piccolo, ma gli autori propongono un'alternativa strutturale alla regolarizzazione.

2. Metodologia

Il paper introduce un metodo che sostituisce la regolarizzazione KL con la media dei logit. Invece di penalizzare la divergenza, la politica addestrabile $\pi_\theta$ e una politica di riferimento congelata $\pi_{\text{ref}}$ (ad esempio SFT) sono accoppiate mediando i loro logit per formare una politica mista $\pi^{\text{mix}}_{\theta, \alpha}$.

Meccanismo Principale:
Dato un peso $\alpha \in [0, 1]$ e uno stato $s$, i logit misti $z^{\text{mix}}_{\theta, \alpha}(s)$ sono definiti come:
$$z^{\text{mix}}_{\theta, \alpha}(s) = (1 - \alpha) z_\theta(s) + \alpha z_{\text{ref}}(s)$$
La politica mista è quindi $\pi^{\text{mix}}_{\theta, \alpha}(a | s) = \text{softmax}(z^{\text{mix}}_{\theta, \alpha}(s))$.

Obiettivo di Addestramento:
Gli autori adattano l'obiettivo GRPO a questa politica mista:

• Rollout: Le traiettorie sono campionate dalla politica mista $\pi^{\text{mix}}_{\theta, \alpha}$ piuttosto che dalla politica addestrabile pura $\pi_\theta$.
• Stima del Vantaggio: I vantaggi relativi al gruppo sono calcolati basandosi sulle ricompense di questi rollout misti.
• Ottimizzazione: I parametri addestrabili $\theta$ sono aggiornati per massimizzare la ricompensa sotto la politica mista. I logit di riferimento $z_{\text{ref}}$ rimangono congelati per tutta la durata.
• Schemi di Pesatura:
  1. Pesi Fissi: $\alpha$ è mantenuto costante (ad esempio, $\alpha=0.5$) durante tutto l'addestramento.
  2. Pesi Adattivi: $\alpha$ viene aggiornato dinamicamente basandosi su un set di convalida tenuto da parte. La regola di aggiornamento utilizza una metrica "net-gain" che confronta la correttezza della politica pura rispetto alla politica mista, spingendo $\alpha$ verso 1 se la miscelazione produce un guadagno netto e verso 0 altrimenti.

Confronto con RLVR:
A differenza dei metodi RLVR standard (PPO, GRPO, DAPO) che utilizzano un critic e regolarizzazione KL, questo approccio non ha termine KL né critic. La politica di riferimento partecipa all'esplorazione tramite media dei logit, e i rapporti di importanza nella funzione di perdita sono calcolati tra le politiche miste corrente e precedente.

3. Contributi Chiave e Risultati

Competenze Complementari (Pattern 1 & 2):
L'analisi empirica sui dataset MATH, cn-k12 e MMLU rivela che l'SFT e la politica addestrata con RL possiedono competenze complementari:

• SFT eccelle nella formattazione (aderendo ai template di output) ma può fallire su ragionamenti complessi.
• Politica RL ($\pi_\theta$) eccelle nel ragionamento (derivando risposte numeriche corrette) ma spesso perde la disciplina nella formattazione.
• Complementarità: Una frazione significativa di problemi (22%–38%) è risolta correttamente da esattamente una delle due politiche. La media dei logit agisce come un "Prodotto di Esperti", combinando la forza nella formattazione dell'SFT con la forza nel ragionamento della politica RL.

Risultati sulle Prestazioni:
Valutato sui modelli Qwen2.5-Instruct (1.5B, 3B, 7B) su tre dataset:

• Accuratezza: Il metodo di media dei logit raggiunge un'accuratezza delle risposte più alta o comparabile rispetto al GRPO regolarizzato con KL canonico in tutte e nove le configurazioni (dataset $\times$ dimensione del modello).
• Ablazione: Sostituire l'ancora SFT con un modello base Instruct risulta in prestazioni significativamente inferiori, confermando che i guadagni derivano dalla miscelazione con un riferimento consapevole della formattazione.
• Media dei Logit vs. Media delle Probabilità: La media dei logit (media geometrica delle distribuzioni) supera la media delle probabilità (media aritmetica). La media dei logit raggiunge tassi di risoluzione più alti con pesi di miscelazione più piccoli e ha un picco a $\alpha=0.5$, mentre la media delle probabilità richiede pesi più alti e ha un picco a $\alpha=0.6$.

Connessione Teorica:
Il paper stabilisce che la media dei logit è matematicamente equivalente a un Prodotto di Esperti (POE). La politica risultante minimizza una somma pesata di divergenze KL verso le politiche RL e di riferimento. Questa struttura permette alla politica di concentrare la probabilità sui token favoriti da entrambi gli esperti, correggendo efficacemente gli errori algebrici dell'SFT pur mantenendo i suoi vincoli di formattazione.

4. Significato e Affermazioni

Gli autori affermano che questo metodo offre un vantaggio distinto rispetto alla ricetta standard regolarizzata con KL, disaccoppiando l'"ancora di formattazione" dal "vincolo del gradiente".

• Superamento dei Limiti: Mentre la regolarizzazione KL forza $\pi_\theta$ a rimanere vicino a $\pi_{\text{sft}}$ (potenzialmente ereditandone i difetti di ragionamento), la media dei logit permette a $\pi_\theta$ di superare i limiti del riferimento mantenendo i vantaggi di formattazione attraverso l'accoppiamento strutturale della politica mista.
• Regione di Fiducia Implicita: La miscelazione agisce come una regione di fiducia implicita, prevenendo la deriva distributiva lontano dagli stili di formattazione che spesso affligge l'RL puro su larga scala.
• Priors Densi: A differenza dell'RL che propaga segnali solo dal token finale di ricompensa, la media dei logit incorpora un prior denso a livello di token dall'ancora SFT a ogni passo di decodifica.
• Scalabilità: L'approccio si estende naturalmente a ensemble di Small Language Models (SLM), permettendo a un sistema di sfruttare esperti specializzati (ad esempio, per codice, sicurezza o terminologia di dominio) senza il costo di addestrare un'unica politica massiva.

Il paper conclude che la media dei logit fornisce un'alternativa robusta alla regolarizzazione KL, sfruttando con successo la complementarità tra SFT e RL per migliorare l'accuratezza del ragionamento senza sacrificare la qualità della formattazione.