Reformulate LLM Reinforcement Learning for Efficient Training under Black-box Discrepancy

Sintesi Tecnica: Riformulazione dell'Apprendimento per Rinforzo degli LLM per un Addestamento Efficiente sotto Discrepanza Black-box

1. Definizione del Problema

L'Apprendimento per Rinforzo (RL) è diventato un paradigma di post-addestramento critico per potenziare le capacità di ragionamento dei Large Language Models (LLM). Tuttavia, l'addestramento RL post-training soffre frequentemente di prestazioni sub-ottimali imprevedibili o di completi collassi dell'addestramento. Indagini recenti attribuiscono questi fallimenti a una pervasiva discrepanza tra addestramento e inferenza (o mismatch).

Questa discrepanza sorge perché le esigenze industriali costringono spesso i praticanti a utilizzare motori sottostanti e formati di precisione disparati (ad esempio, diversi kernel, strategie di parallelizzazione o precisioni numeriche come BF16 vs. FP8) per gli stessi identici pesi del modello durante l'addestramento rispetto all'inferenza. Ciò crea una deviazione nascosta tra le distribuzioni di probabilità delle azioni della policy di addestramento ($\pi_{Train}$) e della policy di inferenza ($\pi_{Inf}$).

Le strategie di mitigazione esistenti presentano limitazioni significative:

• Allineamento della Precisione Uniforme: Il semplice declassamento della precisione di addestramento per farla corrispondere a quella di inferenza (ad esempio, a FP16) restringe l'intervallo di valori rappresentabili ed esacerba i rischi di esplosione del gradiente.
• Interventi Euristici: Metodi come il masking dei gradienti dai token ad alta discrepanza servono solo a sopprimere l'impatto dei dati rumorosi. Essi non affrontano la causa radice e trascurano l'obiettivo ultimo: ottimizzare la policy per il deployment di inferenza, non solo per il processo di addestramento.
• Allineamento dell'Infrastruttura: Tentare di allineare gli operatori di basso livello e gli stack hardware è laborioso e spesso impraticabile a causa della miriade di casi limite.

La sfida centrale consiste nel consentire alla policy di addestramento di percepire dinamicamente e neutralizzare attivamente la propria discrepanza con la policy di inferenza senza sacrificare l'obiettivo primario di massimizzazione della ricompensa.

2. Metodologia: Discrepancy-Constrained MDP (DCMDP)

Gli autori propongono un cambio di paradigma dal classico processo decisionale di Markov (MDP) a un Discrepancy-Constrained MDP (DCMDP). Questo framework tratta la discrepanza tra addestramento e inferenza come un problema di ottimizzazione black-box, permettendo al sistema di gestire configurazioni eterogenee di addestramento e inferenza in modo fluido.

Componenti Chiave:

• Metrica di Discrepanza:
  Gli autori valutano empiricamente tre segnali di discrepanza a livello di token: Differenza di Probabilità ($\delta_{diff}$), Rapporto di Probabilità ($\delta_{ratio}$) e Differenza di Log-Probabilità ($\delta_{log}$). Identificano che la differenza assoluta di probabilità ($\delta_{diff} = |p_{Train} - p_{Inf}|$) è la penalità più efficace. Essa offre una regolarizzazione "plastica" che corregge il drift rilevante per il deployment senza vincolare eccessivamente l'esplorazione guidata dalla ricompensa, a differenza delle penalità basate sul rapporto che possono essere troppo conservative per i token a bassa probabilità.

• La Regione di Tolleranza della Discrepanza:
  Un risultato empirico critico è l'esistenza di una regione di tolleranza della discrepanza. Penalizzare aggressivamente le discrepanze minori e trascurabili nelle fasi iniziali dell'addestramento sopprime l'esplorazione della policy e riduce l'efficienza dell'apprendimento. Al contrario, ridurre la discrepanza eccessiva al di fuori di questa regione migliora la coerenza dell'ottimizzazione.

  • Meccanismo: La formulazione DCMDP applica le penalità solo quando la discrepanza supera una soglia predefinita $c$. All'interno di questa regione di tolleranza, la policy è libera di esplorare per massimizzare le ricompense.

• Ottimizzazione a Doppio Obiettivo tramite Rilassamento Lagrangiano:
  Il problema è formulato come la massimizzazione della ricompensa attesa soggetta a un vincolo sulla discrepanza media a livello di token:
  $$\max_{\theta} J(\theta) = \mathbb{E}_{\tau \sim \pi_{\theta}} \left[ \sum r(o_t) \right]$$
  $$\text{s.t. } \frac{1}{|\tau|} \sum |\pi_{Train}(o_t) - \pi_{Inf}(o_t)| \leq c$$

Per risolverlo, gli autori impiegano il Rilassamento Lagrangiano. Un moltiplicatore apprendibile $\lambda$ regola dinamicamente il peso della penalità di discrepanza:
* Se la discrepanza viola il confine di tolleranza ($>c$), $\lambda$ aumenta, rafforzando la penalità per riportare la policy nella zona sicura.
* Se la discrepanza è entro il confine, $\lambda$ decade verso lo zero, permettendo alla policy di concentrarsi sulla massimizzazione della ricompensa.

L'obiettivo finale (Dual-GRPO) inietta questa penalità adattiva direttamente nel vantaggio normalizzato per gruppo:
$$\hat{A}^{DCMDP}_{i,t} = \hat{A}_{i,t} - \lambda_k \cdot \delta_{i,t}$$

3. Contributi Principali

1. Compensazione Autonoma della Discrepanza: Il documento dimostra che l'incorporazione di una robusta penalità black-box (differenza di probabilità assoluta a livello di token) conferisce alla policy di addestramento la capacità di eradicare autonomamente il rischio di collasso dell'addestramento indotto dalla discrepanza.
2. Cambio di Paradigma verso il DCMDP: Gli autori riformulano l'RL per LLM come un Discrepancy-Constrained MDP. Questo approccio bilancia ottimamente i doppi obiettivi di massimizzare la ricompensa del task e minimizzare la discrepanza addestramento-inferenza, andando oltre la semplice scalarizzazione degli obiettivi.
3. Paradigma di Addestramento Eterogeneo: Il metodo consente una configurazione di addestramento in cui i modelli sono ottimizzati in ambienti ad alta fedeltà e ad alta intensità di risorse (es. precisione BF16), pur essendo esplicitamente allineati per deployment di inferenza a basso costo e risorse limitate (es. FP8). Ciò trasforma l'eterogeneità dell'infrastruttura da un limite a un controllo ottimizzabile.

4. Risultati Sperimentali

Il metodo è stato valutato su modelli Qwen3-8B (dense) e Qwen3-30B-A3B (Mixture-of-Experts) su sei benchmark di ragionamento matematico (AIME24/25, AMC23, MATH-500, Minerva, OlympiadBench).

• Addestramento/Inferenza BF16: DC-GRPO ha superato costantemente il baseline vanilla GRPO in tutti i benchmark per entrambi i modelli, ottenendo punteggi medi più elevati (ad esempio, Qwen3-8B è passato da 67.3 a 71.0). Ciò conferma che l'imposizione di vincoli di discrepanza potenzia, anziché ostacolare, l'ottimizzazione della ricompensa primaria.
• Addestramento Eterogeneo (BF16) + Deployment (FP8): In questa configurazione, il vanilla GRPO ha subito un significativo degrado delle prestazioni o un collasso dovuto all'amplificazione del divario addestramento-inferenza. DC-GRPO ha mitigato con successo questo fenomeno, preservando e spesso superando le prestazioni del baseline BF16-only.
• Studi di Ablazione: Il metodo si è dimostrato robusto alle variazioni del moltiplicatore Lagrangiano iniziale ($\lambda_0$) e del tasso di apprendimento duale ($\eta_\lambda$). Il budget di tolleranza ($c$) è risultato critico; impostarlo troppo basso causava un eccesso di regolarizzazione, mentre impostarlo troppo alto faceva degenerare il metodo in un vanilla GRPO.

5. Significato e Rivendicazioni

Il documento sostiene che il DCMDP fornisca un framework algoritmico unificato per mitigare l'instabilità tra addestramento e inferenza senza richiedere ottimizzazioni esaustive e su misura dello stack hardware o software sottostante. Trattando la discrepanza come un segnale di feedback black-box, il metodo permette:

• Addestramento Stabile: Prevenire i collassi improvvisi comuni nel post-training RL.
• Esplorazione Efficiente: Permettere alla policy di esplorare liberamente entro una regione di discrepanza sicura.
• Efficienza delle Risorse: Consentire l'uso di configurazioni di addestramento ad alta fedeltà per simulare e adattarsi ad ambienti di deployment vincolati, facilitando così il deployment reale dei modelli in modo più efficiente.

Gli autori riconoscono i limiti, notando che l'ambito empirico è attualmente limitato alle precisioni BF16 e FP8 e a specifiche famiglie di modelli, con la necessità di futuri lavori per estendere il metodo ad altre combinazioni di backend e architetture di modelli.