Riepilogo Tecnico: Post-Training con Bilanciamento della Traiettoria per Modelli Linguistici Diffusivi

Problema: Blocco della Traiettoria nel RL Diffusivo

I Modelli Linguistici Diffusivi (dLLM) offrono un'alternativa promettente ai modelli autoregressivi per il ragionamento e la generazione di codice. Tuttavia, i metodi esistenti di post-training per i dLLM adattano in larga misura obiettivi di Apprendimento per Rinforzo (RL) volti alla massimizzazione della ricompensa (ad esempio, PPO, GRPO) progettati per modelli autoregressivi. Gli autori identificano una modalità di fallimento centrale in questo adattamento, denominata blocco della traiettoria.

Nei modelli diffusivi, la ricompensa viene assegnata solo al completamento denoizzato finale ($y$), mentre la distribuzione del modello è indotta sull'intera traiettoria di denoising ($\tau$). Poiché le ricompense terminali sono "indifferenti al percorso" — ovvero, qualsiasi traiettoria che raggiunge lo stesso $y$ corretto riceve la stessa ricompensa — gli aggiornamenti del gradiente della politica campionati tendono a rafforzare le traiettorie specifiche che vengono campionate casualmente nelle fasi iniziali dell'addestramento. Ciò crea un ciclo di feedback auto-rinforzante in cui la massa di probabilità collassa su un insieme ristretto di percorsi di denoising. Di conseguenza, anche se un modello impara a produrre una risposta corretta, perde la capacità di generare un insieme diversificato di soluzioni valide (modi) sotto campionamento ripetuto. Gli autori formalizzano questo fenomeno tramite il Teorema 1, dimostrando che la copertura dei modi terminali è limitata superiormente dalla diversità della traiettoria; pertanto, il blocco della traiettoria limita direttamente la copertura delle soluzioni del modello.

Metodologia: TraFL (Bilanciamento del Flusso di Traiettoria)

Per affrontare il blocco della traiettoria, gli autori propongono TraFL, un obiettivo di post-training fondato sul principio del bilanciamento della traiettoria delle Reti di Flusso Generativo (GFlowNets). Invece di massimizzare direttamente la ricompensa, TraFL addestra la politica a far corrispondere una distribuzione target inclinata dalla ricompensa ancorata a un modello di riferimento congelato.

L'Obiettivo

La distribuzione target $p^*(\tau | x)$ è definita come:
$$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$$
dove:

• $p_{\text{ref}}$ è un modello di riferimento congelato.
• $r(x, y)$ è la ricompensa terminale.
• $\beta$ controlla la nitidezza della ricompensa.
• $Z(x)$ è un termine di normalizzazione dipendente dal prompt (funzione di partizione).

L'obiettivo di addestramento minimizza il residuo quadratico tra il log-probabilità del modello corrente e il target:
$$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$$
$$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$$

Implementazione Pratica per la Diffusione

Poiché i dLLM non espongono log-probabilità esatte della traiettoria, gli autori introducono due componenti chiave per rendere TraFL praticabile:

1. Surrogato a Livello di Sequenza Compatibile con la Diffusione: Derivano un surrogato di log-probabilità trattabile ($\log \hat{p}_\theta(y|x)$) basato su un limite inferiore di ricostruzione mascherata. Ciò comporta il campionamento di un livello di mascheramento $l$, la corruzione di $l$ token e la media delle log-probabilità dei token ricostruiti. Viene utilizzato un fattore $1/l$ per normalizzare il punteggio indipendentemente dal numero di token mascherati.
2. Normalizzazione Appresa Dipendente dal Prompt: A differenza dei metodi che stimano $Z(x)$ tramite softmax su un batch (che possono indurre concentrazione), TraFL apprende una funzione di partizione $Z_\phi(x)$ tramite una piccola MLP. Questo termine riceve segnali di gradiente a ogni passo indipendentemente da quali completamenti vengono campionati, prevenendo il fallimento della normalizzazione "locale al batch" osservato in altri approcci di corrispondenza delle distribuzioni.

Contributi Chiave

1. Identificazione del Blocco della Traiettoria: Il documento identifica e dimostra formalmente che gli obiettivi di massimizzazione della ricompensa nei dLLM portano a un collasso della diversità delle traiettorie, che a sua volta limita la copertura dei modi distinti di soluzioni corrette.
2. Obiettivo TraFL: La proposta di un obiettivo di bilanciamento della traiettoria che ancori la politica a un modello di riferimento, separando il ruolo della ricompensa (inclinazione dei completamenti terminali) da quello del modello di riferimento (allocazione della massa di probabilità su traiettorie plausibili).
3. Surrogato Pratico e Normalizzazione: Lo sviluppo di un surrogato a livello di sequenza compatibile con la diffusione e di un termine di normalizzazione appreso che consente un addestramento stabile senza log-probabilità esatte della traiettoria.
4. Superiorità Empirica: TraFL è l'unico metodo di post-training valutato che migliora costantemente rispetto al modello di base in tutte le impostazioni di lunghezza del benchmark (GSM8K, MATH-500, HumanEval, MBPP) e nei budget di campionamento.

Risultati

Gli autori valutano TraFL su benchmark di ragionamento matematico e generazione di codice confrontandolo con il modello di base (LLaDA-8B-Instruct) e con baseline solide (ESPO, JUSTGRPO).

• Prestazioni In-Distribution: TraFL migliora il Pass@k in tutti i budget di campionamento ($k \in \{1, 2, 4, 8, 16\}$) e temperature. In particolare, migliora il Pass@1 (accuratezza del singolo campione) mentre scala meglio con l'aumento dei budget di campionamento, a differenza delle baseline che spesso sacrificano l'accuratezza del singolo campione per guadagni multi-campione.
• Generalizzazione su Dati Tenuti Fuori: I guadagni si trasferiscono a benchmark fuori distribuzione. Su Minerva Math, TraFL rimane al di sopra del modello di base. Su LiveCodeBench, TraFL è il metodo più forte in tutte le divisioni di difficoltà (Facile, Medio, Difficile), superando significativamente le baseline sui problemi medi e difficili.
• Analisi della Diversità: Utilizzando un protocollo LLM-as-a-judge, gli autori mostrano che i guadagni di TraFL sono associati a una copertura più ampia di strategie di soluzione distinte (ad esempio, diversi approcci dimostrativi o decomposizioni algoritmiche) piuttosto che semplicemente all'acutizzazione del punteggio di un singolo modo. Baseline come ESPO, pur migliorando l'accuratezza, non mostrano guadagni significativi di diversità rispetto al modello di base, coerentemente con l'ipotesi del blocco della traiettoria.

Significato e Affermazioni

Il documento afferma che TraFL fornisce una soluzione fondata al fallimento "cercante il modo" del post-training RL standard per i modelli diffusivi. Spostandosi dalla massimizzazione della ricompensa al bilanciamento della traiettoria, il metodo preserva la diversità dei percorsi di denoising necessaria per coprire molteplici soluzioni valide.

Gli autori affermano che TraFL è l'unico metodo valutato che migliora rispetto al modello di base in ogni impostazione di lunghezza del benchmark testata. Sottolineano che i miglioramenti sono robusti all'aumento del budget di campionamento e si generalizzano a dataset tenuti fuori, suggerendo che l'obiettivo di bilanciamento della traiettoria disaccoppia efficacemente il segnale di ricompensa dal meccanismo di allocazione del percorso. Il lavoro sostiene obiettivi di corrispondenza delle distribuzioni rispetto alla pura massimizzazione della ricompensa per il post-training dei dLLM, offrendo un quadro pratico che evita le modalità di fallimento della concentrazione intrinseche negli approcci standard di gradiente della politica.

Il documento rimane modesto riguardo alle limitazioni, notando che il metodo si basa su un surrogato approssimato a livello di sequenza piuttosto che su log-probabilità esatte della traiettoria e utilizza ricompense binarie grossolane (corrispondenza esatta o superamento del test), il che potrebbe limitare l'efficienza dell'apprendimento in scenari con successi rari.