Beyond Token-Level Policy Gradients for Complex Reasoning with Large Language Models

1. Il Problema: La Disallineamento Granulare nel Ragionamento

I metodi esistenti di ottimizzazione basati sul gradiente di politica (come PPO, GRPO e DAPO) per i Large Language Models (LLM) trattano la generazione di ogni singolo token come un'azione indipendente e sequenziale.

• Limitazione: Sebbene efficace per molti compiti di generazione, questo approccio non cattura pienamente la struttura dei compiti di ragionamento complesso (es. matematica, programmazione). In questi scenari, una singola decisione semantica (come definire una variabile, comporre un'equazione o scrivere una funzione) è spesso realizzata attraverso un blocco di più token consecutivi.
• Conseguenza: C'è un disallineamento fondamentale tra l'ottimizzazione a livello di singolo token e la natura intrinsecamente "a blocchi" del ragionamento. Questo frammenta le azioni semantiche coerenti, portando a una perdita di coerenza nelle traiettorie di ragionamento e a una stima del vantaggio (advantage) meno efficace.

2. Metodologia: Multi-token Policy Gradient Optimization (MPO)

Per colmare questo divario, gli autori propongono MPO, un framework che tratta sequenze di $K$ token consecutivi come azioni semantiche unificate (blocchi), piuttosto che come azioni isolate.

Componenti Chiave:

• Azione a Blocchi: Invece di calcolare il rapporto di importanza (importance sampling ratio) per ogni singolo token, MPO aggrega $K$ token contigui. Il rapporto di importanza viene calcolato sulla probabilità congiunta di generare l'intero blocco di $K$ token dato il contesto precedente.
• Moduli Multi-Token Prediction (MTP): Per implementare ciò, il framework utilizza moduli MTP (simili a quelli usati in DeepSeek-V3) che permettono al modello di prevedere fino a $K$ token futuri in un'unica passata in avanti, mantenendo la coerenza causale.
• Ottimizzazione del Rapporto di Importanza:
  • Una formulazione diretta come prodotto di rapporti ($\prod \pi/\pi_{old}$) aumenterebbe eccessivamente la varianza.
  • MPO adotta una formulazione alternativa basata su una somma logaritmica pesata (ispirata al metodo Log-COP-TD):
    $$\tilde{R}^{(K)}_{i,t}(\theta) = \exp \left( \sum_{n=1}^{K} \beta_n \log r_{i,t+n}(\theta) \right)$$
    dove $\beta_n$ sono pesi che decadono ($\beta_k = \beta_2 \times \lambda^{k-2}$), dando più peso ai token più vicini e meno a quelli più distanti per controllare la varianza.
• Funzione Obiettivo: L'obiettivo surrogate di MPO è analogo a quello di PPO, ma utilizza questo rapporto aggregato a blocchi moltiplicato per la stima del vantaggio, incoraggiando il modello a ottimizzare pattern multi-token coerenti.
• Strategia di Addestramento:
  1. Warm-up: I moduli MTP vengono inizializzati e "riscaldati" (warm-up) tramite un obiettivo di previsione multi-token supervisionato prima dell'addestramento RL, per garantire la qualità delle previsioni future.
  2. Fine-tuning RL: Durante la fase di ottimizzazione della politica, i gradienti vengono calcolati sui blocchi di token, mentre il backbone del modello rimane congelato (in alcune configurazioni) o aggiornato con cautela.

3. Contributi Principali

1. Proposta di MPO: Introduzione di un nuovo algoritmo di ottimizzazione della politica che opera su unità di ragionamento strutturate (blocchi di token) invece che su token isolati.
2. Integrazione nel Post-Training: Dimostrazione che l'integrazione di ottimizzazioni strutturate multi-token nella fase di post-training (RLHF) migliora significativamente la coerenza del ragionamento.
3. Performance Superiori: Validazione empirica che MPO supera sistematicamente i baseline a livello di token (PPO, GRPO, DAPO) su benchmark di ragionamento matematico e generazione di codice.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli di diverse dimensioni (da 1B a 7B parametri) e su tre benchmark principali: GSM8K (ragionamento matematico di livello universitario), MATH (competizioni matematiche) e HumanEval (generazione di codice).

• Ragionamento Matematico: MPO ha mostrato miglioramenti costanti rispetto ai baseline. Ad esempio, su GSM8K con il modello DeepSeek-Qwen2.5-1.5B, MPO ha raggiunto un'accuratezza del 88.2%, superando GRPO (85.0%) e PPO (86.6%). Su MATH, i guadagni sono stati ancora più marcati, indicando una migliore capacità di gestire problemi complessi.
• Generazione di Codice: Su HumanEval, MPO ha ottenuto risultati superiori (es. 64.0% su DeepSeek-Qwen2.5-1.5B contro il 59.8% di PPO), dimostrando che la dipendenza semantica multi-step è cruciale anche per la sintesi di programmi.
• Stabilità dell'Addestramento: L'analisi ha rivelato che MPO riduce la varianza dei rapporti di importanza e la frazione di "clipping" (taglio dei gradienti) durante l'addestramento, rendendo il processo più stabile rispetto ai metodi token-level.
• Trade-off Bias-Varianza: È stato identificato un punto ottimale (circa $K=5$ e un peso cumulativo dei moduli MTP del 10-20%) che bilancia la riduzione del bias (migliore visione a lungo termine) con l'aumento del rumore di previsione.

5. Significato e Implicazioni

Questo lavoro segna un cambiamento metodologico nell'ottimizzazione dei LLM per compiti di ragionamento:

• Oltre il Token: Dimostra che l'ottimizzazione a livello di singolo token è insufficiente per compiti dove il significato emerge da strutture più ampie (equazioni, funzioni).
• Efficienza: Nonostante l'aggiunta di moduli MTP, l'overhead computazionale è gestibile (circa il 30% più veloce di GRPO in termini di iterazioni per secondo, con un consumo di memoria simile), rendendo l'approccio scalabile.
• Futuro della Ricerca: MPO apre la strada a future ricerche che esplorano la granularità delle decisioni nelle politiche di apprendimento, spingendo verso processi decisionali strutturati semanticamente piuttosto che puramente sequenziali a livello di token.

In sintesi, MPO risolve il problema della frammentazione semantica nel ragionamento complesso aggregando le azioni di ottimizzazione in blocchi significativi, migliorando sia la coerenza logica che le prestazioni finali dei modelli.