Riepilogo Tecnico: Addestramento Congiunto della Predizione Multi-Token nel Reinforcement Learning tramite Calibrazione Ottimale dei Coefficienti

Enunciato del Problema

Il Reinforcement Learning da Ricompense Verificabili (RLVR) è diventato il paradigma standard per potenziare le capacità di ragionamento dei Large Language Models (LLM). Parallelamente, la Predizione Multi-Token (MTP), un obiettivo ausiliario addestrato durante il pre-addestramento per prevedere più token futuri, ha dimostrato benefici significativi nell'efficienza del campione e nel fungere da modello di bozza per la decodifica speculativa.

Intuitivamente, l'ottimizzazione congiunta di MTP e degli obiettivi RL durante il post-addestramento dovrebbe permettere ai modelli di beneficiare della supervisione multi-token mentre acquisiscono capacità di ragionamento. Tuttavia, le evidenze empiriche mostrano che l'addestramento congiunto MTP-RL spesso porta a un degrado significativo delle prestazioni del modello principale. Di conseguenza, le pratiche attuali (ad esempio, in framework come veRL e slime) tipicamente staccano (detach) i gradienti MTP dal modello principale durante l'addestramento RL. Questo distacco costringe i parametri MTP ad adattarsi solo tramite un fine-tuning post-hoc indipendente, sacrificando i potenziali guadagni derivanti dalla supervisione multi-token simultanea. Il problema centrale affrontato è: Perché l'addestramento congiunto MTP-RL fallisce e può essere progettata una strategia fondata per abilitare un addestramento congiunto efficace?

Metodologia

Analisi Teorica

Gli autori affrontano il problema da una prospettiva di ottimizzazione, analizzando l'effetto per passo di MTP sull'obiettivo RL $J(\theta)$. Assumendo una funzione obiettivo $L$-liscia, derivano un limite inferiore per il miglioramento per passo. Questa decomposizione rivela che l'effetto di MTP è governato da due termini in competizione:

1. Termine di Correlazione del Primo Ordine: Determinato dall'accordo direzionale (prodotto scalare) tra il gradiente RL ($g_{RL}$) e il gradiente MTP ($g_{MTP}$).
2. Penalità di Perturbazione del Secondo Ordine: Un termine non positivo che riflette la perturbazione introdotta dalla magnitudine del gradiente MTP.

Il miglioramento delle prestazioni si verifica solo se il termine di correlazione supera il termine di penalità.

Analisi degli Regimi Esistenti

Utilizzando questa decomposizione, il lavoro unifica e spiega tre regimi di addestramento esistenti:

• Stacco (Detach): $g_{MTP} = 0$ rispetto al modello principale. $\Delta_{MTP} = 0$. Questo è sicuro ma rinuncia ai potenziali guadagni.
• Perdita Cross-Entropy (CE): Il termine di correlazione è vicino a zero in media perché l'RL aumenta il peso dei campioni ad alta ricompensa mentre CE tratta tutti i campioni allo stesso modo. Di conseguenza, la negativa penalità del secondo ordine domina, causando degrado.
• Perdita della Politica (Policy Loss): Utilizza la stessa formulazione dell'obiettivo RL per MTP. Inizialmente, i gradienti sono allineati e il termine di correlazione domina, producendo guadagni. Tuttavia, man mano che l'addestramento procede e il modello entra in regioni piatte del paesaggio della perdita, la correlazione decade mentre la norma del gradiente (e quindi la penalità) persiste. Ciò porta a una curva di prestazioni "crescita-poi-declino" in cui un coefficiente fisso non può tracciare la transizione di fase.

Soluzione Proposta: Calibrazione Ottimale dei Coefficienti (OCC)

Per affrontare la transizione di fase, gli autori propongono la Calibrazione Ottimale dei Coefficienti (OCC), uno schema adattivo che traccia online il coefficiente MTP ottimale ($\lambda$).

1. Ottimo in Forma Chiusa: L'analisi produce un'espressione in forma chiusa per l'ottimo $\lambda^*$ a ogni passo, che è proporzionale al rapporto tra la correlazione dei gradienti ($c$) e il quadrato della norma del gradiente ($v^2$).
2. Proxy della Log-Probabilità: Il calcolo dei gradienti del modello completo per $c$ e $v^2$ è proibitivamente costoso nei sistemi distribuiti su larga scala. OCC utilizza un proxy della log-probabilità derivato dalla variazione della log-probabilità sotto aggiornamenti a piccoli passi ($\delta \approx \log \pi_\theta - \log \pi_{old}$).
   • Proxy di Allineamento ($\hat{c}_t$): Stima il prodotto scalare.
   • Proxy di Varianza ($\hat{v}^2_t$): Stima il quadrato della norma.
3. Aggiornamento Adattivo: Il coefficiente viene aggiornato come $\lambda_t = \lambda_+ \cdot \frac{\hat{c}_t}{\hat{v}^2_t + \epsilon}$. Ciò permette a OCC di regolare dinamicamente il peso di MTP, mantenendo guadagni positivi durante tutto l'addestramento senza calcolare gradienti completi.

Contributi Chiave

1. Unificazione Teorica: Il lavoro fornisce un'analisi teorica che unifica i regimi Stacco, Cross-Entropy e Perdita della Politica, spiegando il loro successo o fallimento in base all'interazione tra correlazione del primo ordine e perturbazione del secondo ordine.
2. Algoritmo OCC: Propone la Calibrazione Ottimale dei Coefficienti, una procedura adattiva che calibra il coefficiente MTP online utilizzando un proxy della log-probabilità. Questo metodo richiede un sovraccarico computazionale trascurabile ed evita il calcolo dei gradienti del modello completo.
3. Validazione Empirica: Esperimenti estesi su più modelli (MiMo-7B, GLM-4.5-Air) e algoritmi (DAPO, GSPO) su benchmark di ragionamento matematico di livello competitivo dimostrano che OCC corrisponde o supera costantemente la linea di base Stacco.

Risultati Sperimentali

Gli autori hanno valutato OCC su sei benchmark: AIME24, AIME25, AMC, MATH, Minerva e OlympiadBench.

• Prestazioni: Su tutti i benchmark e le scale dei modelli (7B denso e 106B MoE), OCC ha costantemente superato la linea di base Stacco e la Perdita della Politica a coefficiente fisso.
  • Su MiMo-7B con DAPO, OCC ha raggiunto una precisione media del 61,7%, superando Stacco (58,9%) e Perdita della Politica (57,4%).
  • Su GLM-4.5-Air, OCC ha raggiunto una precisione media del 67,6% contro il 65,9% di Stacco.
• Dinamiche di Addestramento: Mentre la Perdita della Politica ha mostrato una tendenza di crescita-poi-declino (peggioramento nelle fasi successive), OCC ha mantenuto miglioramenti stabili durante tutto l'addestramento. La Perdita CE ha costantemente sottoperformato rispetto a Stacco.
• Efficienza: OCC ha introdotto un sovraccarico trascurabile. Il tempo di addestramento reale per passo per OCC (8,07s) era comparabile a Stacco (8,11s), mentre la retropropagazione del modello completo richiedeva 45,06s (un rallentamento di 5,6×).
• Ablazione: L'analisi di sensibilità ha mostrato che il fattore di scala predefinito $\lambda_+$ è robusto su un intervallo di valori, con $\lambda_+ = 1,0$ che produce le migliori prestazioni medie.

Significato e Affermazioni

Il lavoro afferma che il fallimento dell'addestramento congiunto MTP-RL non è intrinseco ma deriva da una discrepanza tra coefficienti statici e la deriva dinamica dell'allineamento dei gradienti durante l'addestramento. Fornendo una decomposizione teorica del paesaggio di ottimizzazione e proponendo OCC, gli autori dimostrano che MTP può essere sicuramente reintrodotto nel post-addestramento RL.

Il significato risiede nel permettere ai modelli di acquisire simultaneamente capacità di ragionamento tramite RL e sfruttare rappresentazioni multi-token senza sacrificare le prestazioni o incorrere in costi computazionali proibitivi. Il lavoro suggerisce che la calibrazione adattiva è essenziale per l'addestramento congiunto nei paesaggi di deep learning non convessi, offrendo un'alternativa fondata allo standard industriale attuale di distacco dei gradienti.

Limitazioni

Gli autori notano che la loro analisi teorica si basa su assunzioni di $L$-liscietà che potrebbero non valere esattamente nei paesaggi non convessi, fornendo indicazioni qualitative piuttosto che quantitative. Il proxy della log-probabilità è un'approssimazione del primo ordine che potrebbe perdere fedeltà sotto tassi di apprendimento elevati. Inoltre, gli esperimenti sono focalizzati sul ragionamento matematico con ricompense verificabili; la generalizzazione a contesti come RLHF con modelli di ricompensa appresi o generazione aperta rimane da validare.