Patch the Distribution Mismatch: RL Rewriting Agent for Stable Off-Policy SFT

1. Il Problema: Divergenza di Distribuzione e Dimenticanza Catastrofica

Il lavoro affronta una sfida critica nell'adattamento dei Large Language Models (LLM): la dimenticanza catastrofica (catastrophic forgetting) che si verifica durante il Supervised Fine-Tuning (SFT) su dati di dominio specifico.

• Natura del problema: Quando un modello pre-addestrato viene fine-tunato su dati di downstream (es. ragionamento matematico) che presentano un forte spostamento di distribuzione rispetto ai dati di addestramento originali, l'SFT standard agisce come un apprendimento off-policy. Le dimostrazioni esperte ($y^*$) possono avere una probabilità molto bassa sotto la politica corrente del modello, portando a instabilità nell'ottimizzazione e a una deriva della politica verso la distribuzione di downstream, erodendo le capacità generali acquisite.
• Limitazioni degli approcci esistenti: Le tecniche attuali di data rewriting (riscrittura dei dati prima dell'SFT) tentano di colmare questo divario modificando le istruzioni o le soluzioni. Tuttavia, questi metodi soffrono di due difetti principali:
  1. Disallineamento distributivo: Riscrivono i dati condizionandoli su prompt fissi e template ($\pi_0(\cdot | x, y^*, x_{prompt})$), creando target che non corrispondono alla distribuzione naturale di generazione "stile Q&A" del modello ($\pi_0(\cdot | x)$).
  2. Collasso della diversità: L'uso di template fissi porta a una riduzione della varietà delle risposte (mode collapse), limitando l'esplorazione di realizzazioni valide ma diverse.

2. Metodologia: Agente di Riscrittura basato su RL

Gli autori propongono di modellare la riscrittura dei dati come un problema di apprendimento della politica (policy learning) risolto tramite Reinforcement Learning (RL). L'obiettivo è apprendere un agente di riscrittura che generi supervisione di alta qualità, allineata alla distribuzione naturale del modello e diversificata.

Architettura del Framework

Il sistema opera in due fasi principali:

1. Fase I: Addestramento dell'Agente di Riscrittura ($R_\phi$)

   • Base: Si utilizza un modello base istruito e congelato ($\pi_0$) come backbone.
   • Parametri Efficienti: La politica di riscrittura $R_\phi$ è implementata tramite adattatori LoRA (Low-Rank Adaptation) applicati su $\pi_0$. Questo permette un "patch" leggero e controllato senza riaddestrare l'intero modello, prevenendo la deriva eccessiva.
   • Ottimizzazione: Si utilizza l'algoritmo GRPO (Group Relative Policy Optimization), una variante on-policy che stima la baseline da gruppi di output campionati, eliminando la necessità di una rete di valore separata.
   • Obiettivo Unificato (Funzione di Ricompensa): L'agente è ottimizzato per massimizzare una ricompensa composta da tre componenti, soggette a un gate di consistenza del compito:
     • Consistenza del Compito ($r_{task}$): Un gate "hard" binario che verifica se la risposta riscritta è corretta (risultato finale e validità del ragionamento). Se fallisce, la ricompensa è zero.
     • Allineamento Distributivo ($r_{dist}$): Una volta superata la verifica, si premia la riscrittura se ha un'alta probabilità (likelihood) sotto la distribuzione QA standard del modello congelato $\pi_0(\cdot | x)$. Questo assicura che i dati riscritti siano "in-distribution" per il modello target.
     • Diversità ($r_{div}$): Si premia la diversità semantica tra le riscritture valide per lo stesso input, calcolando la distanza tra gli embedding delle risposte candidate. Questo previene il collasso dei modi e l'omogeneità indotta dai template.

2. Fase II: Costruzione del Dataset e SFT Downstream

   • L'agente addestrato genera riscritture per il dataset di downstream.
   • Viene applicato un pipeline Generate-Verify-Fallback: se una riscrittura supera il gate di consistenza, viene utilizzata come target di supervisione; altrimenti, si ricade sulla dimostrazione esperta originale ($y^*$) per non perdere dati.
   • Il modello target viene infine fine-tunato (SFT) su questo dataset riscritto ottimizzato.

3. Contributi Chiave

• Insight sulla Disallineamento: Identificazione del fatto che i metodi di riscrittura basati su prompt fissi non allineano realmente i dati alla distribuzione di generazione QA naturale del modello, limitando l'efficacia nella mitigazione della dimenticanza.
• Agente di Riscrittura RL: Proposta di un agente appreso che ottimizza simultaneamente allineamento distributivo e diversità, mantenendo la consistenza del compito come vincolo rigido.
• Obiettivo Unificato: Integrazione di metriche automaticamente valutabili (correttezza, likelihood, diversità) in un unico obiettivo di RL, risolvendo il compromesso tra stabilità, adattabilità e varietà.
• Efficienza: Utilizzo di LoRA su un backbone congelato, rendendo il metodo computazionalmente efficiente e controllabile.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi backbone (Llama-3.2-1B, 3B e Mistral-7B) con compiti di ragionamento matematico come dominio downstream.

• Performance Downstream: Il metodo proposto ("Rewriting-agent") raggiunge guadagni nelle prestazioni matematiche paragonabili o superiori allo SFT standard (Vanilla SFT). Ad esempio, su Mistral-7B, ottiene un miglioramento del +55.23% rispetto al baseline, simile allo SFT (+54.19%).
• Mitigazione della Dimenticanza: Il risultato più significativo è la riduzione drastica della dimenticanza catastrofica sui benchmark generali (out-of-domain).
  • Rispetto allo SFT standard, che causa un calo delle prestazioni generali del ~19%, l'agente di riscrittura riduce questo calo a circa il 7.35%.
  • In media, il metodo riduce la dimenticanza sui benchmark non-matematici del 12.34% rispetto all'SFT standard.
• Analisi dell'Addestramento: Le curve di perdita mostrano che l'SFT su dati riscritti dall'agente parte da una perdita iniziale più bassa e converge a un plateau più stabile rispetto all'SFT su dati originali o riscritti con template fissi (SDFT, Mind the Gap). Questo conferma che i dati riscritti sono più facili da apprendere per il modello nella sua distribuzione naturale.
• Ablazioni:
  • Rimuovere la ricompensa di allineamento distributivo ($r_{dist}$) causa il calo più grande nella ritenzione delle capacità generali, confermando la sua importanza critica.
  • L'uso di un gate "hard" (solo riscritture corrette ricevono ricompense aggiuntive) è superiore a una ricompensa "soft", poiché evita segnali di apprendimento rumorosi su risposte errate.

5. Significato e Implicazioni

Questo lavoro dimostra che l'intervento a livello di dati, guidato da Reinforcement Learning, è una strategia superiore rispetto all'adattamento puramente algoritmico (come i vincoli di aggiornamento conservativo) per gestire lo spostamento di distribuzione nell'SFT.

• Stabilità vs. Plasticità: Il metodo risolve il compromesso tra stabilità (non dimenticare le capacità generali) e plasticità (imparare nuovi compiti) rendendo le supervisioni "più facili" da apprendere per il modello, riducendo la necessità di aggiornamenti drastici e instabili.
• Generalizzabilità: Sebbene testato principalmente su matematica, il framework è progettato per essere applicabile ad altri domini dove la distribuzione dei dati di supervisione differisce significativamente dalla distribuzione di pre-addestramento.
• Qualità dei Dati: Sottolinea che la qualità dei dati per l'SFT non è solo una questione di correttezza del contenuto, ma anche di allineamento distributivo con la modalità di generazione naturale del modello.

In sintesi, il paper propone un approccio elegante che "patcha" il divario distributivo non modificando il modello di addestramento, ma "riparando" i dati di input affinché siano più compatibili con la natura del modello, ottenendo così un adattamento stabile senza sacrificare le capacità generali.