Training Large Language Models To Reason In Parallel With Global Forking Tokens

Il paper propone un nuovo approccio di addestramento, denominato SSFT e GFPO, che utilizza token di diramazione globale per preservare la diversità dei percorsi di ragionamento durante la generazione parallela, superando i compromessi tra accuratezza e varietà e ottenendo prestazioni superiori su compiti di matematica e generazione di codice.

L'essenziale

Immagina di avere un gruppo di esperti (i modelli di intelligenza artificiale) a cui chiedi di risolvere un problema matematico molto difficile.

Se chiedi a un solo esperto di pensarci a lungo, potrebbe impazzire, girare in tondo e finire per dare una risposta sbagliata perché si è "sovraccaricato" di pensieri (il famoso "overthinking").
Se invece chiedi a dieci esperti di lavorare tutti insieme, ognuno con un approccio diverso, e poi prendi la risposta che vince più voti, le probabilità di successo aumentano. Questo è il concetto di "ragionamento parallelo".

Il problema è: come fai a fargli trovare dieci approcci diversi e tutti corretti? Spesso, se lasci che l'IA "speri" di trovare un approccio diverso (aumentando un po' il caos o la temperatura), finisce per generare dieci risposte tutte uguali e sbagliate, oppure dieci risposte diverse ma tutte piene di errori.

Ecco la soluzione proposta in questo paper, spiegata come una storia:

1. Il Problema: La "Forking" (La Biforcazione)

Immagina che il ragionamento dell'IA sia come un albero. All'inizio c'è il tronco (la domanda). Poi, l'albero si dirama.
Per avere risposte diverse, l'albero deve diramarsi in punti specifici chiamati "token di forking" (punti di biforcazione).
Il problema è che questi punti di diramazione corretti sono spesso nascosti molto in profondità nell'albero. È come cercare un ago in un pagliaio: se lasci che l'IA scelga da sola dove diramarsi, spesso sbaglia strada o si ripete.

2. La Soluzione: SSFT (Il "Menu" dei Pensieri)

Gli autori hanno inventato un metodo chiamato SSFT (Set Supervised Fine-Tuning).
Immagina di avere un menu con 6 opzioni speciali (chiamate <think 1>, <think 2>, ecc., come se fossero 6 chef diversi).
Invece di far scegliere all'IA quale strada prendere a caso, gli dai un "menu" di 4 soluzioni corrette (scritte da altri esperti umani o AI potenti) e gli dici:
"Ehi, devi imparare a collegare ogni soluzione del menu a uno chef specifico del tuo gruppo. Se la soluzione A è difficile, assegnala allo Chef 3. Se la soluzione B è creativa, assegnala allo Chef 5."

Il trucco è che il computer fa un abbinamento perfetto (come un puzzle matematico chiamato "matching bipartito"):

1. Guarda tutte le possibili combinazioni tra i 6 Chef e le 4 soluzioni.
2. Trova la combinazione che riduce al minimo gli errori.
3. Impara che lo Chef 3 è specializzato per il tipo di problema della Soluzione A, e lo Chef 5 per la Soluzione B.

In questo modo, l'IA non deve più "scommettere" su quale strada prendere. Impara che se le dai il comando <think 3>, deve attivare un modo di pensare specifico e diverso da quello che attiva con <think 5>.

3. Il Risultato: Non più "Tutti uguali"

Se addestri un'IA normale a vedere 4 soluzioni diverse, tende a "collassare": tutte le soluzioni diventano uguali, come se tutti gli chef avessero copiato lo stesso piatto.
Con il metodo SSFT, invece, gli chef mantengono le loro personalità.

• Se chiedi <think 1>, l'IA usa un approccio logico e rigido.
• Se chiedi <think 4>, usa un approccio creativo e visivo.
• Se chiedi <think 6>, fa un ragionamento passo-passo molto lento e attento.

Ogni "pulsante" (token) apre una porta diversa verso un ragionamento corretto.

4. Il Tocco Finale: GFPO (Il "Manager" Intelligente)

Una volta che l'IA ha imparato a distinguere i suoi 6 chef, c'è un ultimo passo. A volte, non sai quale chef sia il migliore per un problema specifico.
Gli autori aggiungono un piccolo trucco di intelligenza artificiale (chiamato GFPO) che funziona come un manager esperto.
Il manager guarda il problema e dice: "Per questo indovinello matematico, chiamiamo lo Chef 2, perché è quello che ha più probabilità di risolvere il caso!".
Invece di far scegliere a caso l'IA, il manager impara a premere il pulsante giusto per ottenere la risposta migliore.

In Sintesi

Questo paper dice: "Non lasciate che l'IA indovini come pensare in modo diverso. Insegnatele che esistono 'pulsanti' specifici per attivare modi di pensare diversi, e imparate a collegare ogni pulsante al tipo di ragionamento corretto."

Il risultato? L'IA diventa molto più brava a risolvere problemi di matematica e a scrivere codice, perché invece di pensare a lungo e in modo confuso, sa esattamente quale "strumento mentale" usare per ogni situazione. È come passare da un'auto che guida a caso a un'auto con 6 marce diverse, ognuna perfetta per una specifica strada.

Sommario tecnico

1. Il Problema

I Large Language Models (LLM) hanno dimostrato miglioramenti nelle capacità di ragionamento aumentando la potenza di calcolo al momento del test (test-time compute), generando più token prima di fornire una risposta. Tuttavia, l'estensione sequenziale di questo processo può portare al fenomeno dell'"overthinking", dove le prestazioni diminuiscono oltre una certa lunghezza della sequenza.

Un approccio alternativo è il ragionamento parallelo, che genera molteplici percorsi di ragionamento (tracce) in parallelo e li aggrega. La sfida principale risiede nella necessità di generare percorsi che siano sia diversi (per coprire diverse strategie di soluzione) che accurati.

• Il collo di bottiglia: I token che innescano percorsi di ragionamento diversi e corretti (chiamati forking tokens) si trovano tipicamente molto profondi nell'albero di campionamento.
• Il compromesso Diversità-Accuratezza: Le strategie comuni per aumentare la diversità, come l'aumento della temperatura di campionamento, spesso degradano l'accuratezza. Inoltre, aumentare la temperatura da sola non garantisce una maggiore diversità se il modello non è stato esplicitamente addestrato per garantire una copertura (coverage) delle diverse modalità di ragionamento.

2. Metodologia Proposta: SSFT e GFPO

Gli autori propongono un nuovo framework che tratta il ragionamento parallelo come un problema di predizione di un insieme di token successivi (Set of Next-Token Prediction).

A. Token di Forking Globali (Global Forking Tokens)

Il modello introduce un set di token speciali riservati (es. <think 1>, <think 2>, ..., <think N>) che agiscono come "interruttori" globali. Quando un prompt di input è condizionato da uno di questi token, il modello dovrebbe generare una specifica traccia di ragionamento distinta e corretta.

B. Set Supervised Fine-Tuning (SSFT)

Il cuore della metodologia è l'addestramento supervisionato (SFT) con una funzione di perdita basata su insiemi (Set-based loss).

1. Setup: Per ogni domanda $x$, il modello riceve un insieme di $M$ tracce di ragionamento corrette (ground-truth) e un insieme di $N$ token di forking globali ($N \ge M$).
2. Matching Bipartito Ottimale: Invece di addestrare il modello a predire una traccia specifica per un token specifico in modo rigido, SSFT calcola la perdita di predizione del token successivo (NTP loss) per tutte le combinazioni possibili tra token di forking e tracce.
3. Algoritmo di Assegnazione: Utilizzando l'Algoritmo Ungherese (Hungarian Algorithm), il sistema trova la configurazione di matching bipartito a costo minimo che associa ogni traccia corretta a un token di forking unico.
4. Funzione di Perdita: La perdita totale è la somma delle perdite NTP per le coppie (token, traccia) selezionate dal matching ottimale. Questo approccio:
   • È invariante alla permutazione (non importa quale traccia è associata a quale token, purché l'associazione sia unica).
   • Impedisce il collasso delle modalità (mode collapse), forzando il modello a imparare a distinguere le tracce in base ai token di controllo.

C. Global Forking Policy Optimization (GFPO)

Dopo l'addestramento SSFT, viene applicato un passo di ottimizzazione della politica (RL) leggero.

• Obiettivo: Ottimizzare la distribuzione di probabilità dei token di forking globali ($g(i)$) in base alla domanda di input $x$.
• Efficienza: A differenza del RL standard che ottimizza l'intera sequenza, GFPO aggiorna solo i gradienti relativi alla selezione del token di forking iniziale. Questo permette al modello di imparare a scegliere dinamicamente il token di forking più adatto per un dato problema, massimizzando la diversità e l'accuratezza.

3. Contributi Chiave

1. Introduzione dei Token di Forking Globali: Un meccanismo che permette di guidare globalmente le modalità di ragionamento da un singolo token, riducendo la dipendenza dal campionamento casuale di token di fork durante la generazione.
2. SSFT e Matching Ottimale: Dimostrazione che l'addestramento SFT standard su tracce multiple porta al collasso delle modalità (il modello ignora i token di controllo). Al contrario, SSFT preserva le modalità distinte e genera token di forking emergenti che attivano strategie di ragionamento uniche.
3. GFPO: Un metodo efficiente per incentivare il ragionamento complesso sfruttando i token appresi, senza collassare le modalità.
4. Risultati Sperimentali: I modelli addestrati con SSFT e GFPO superano costantemente le controparti addestrate con SFT standard e GRPO (Group Relative Policy Optimization) su benchmark matematici e di generazione di codice.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli Qwen2.5-32B e su dataset come AIME24/25, MATH-500, GPQA-Diamond e LiveCodeBench (codice).

• Prestazioni Pass@1: SSFT-32B ha ottenuto un 64.06% su AIME24 e 58.13% su AIME25, superando i baselines SFT misti di circa 6-8 punti percentuali.
• Copertura e Cons@k: Sotto calcolo parallelo al test-time (Cons@6 e Cons@32), SSFT mostra una copertura significativamente superiore rispetto ai metodi SFT standard. Ad esempio, su AIME25, SSFT raggiunge un Cons@6 del 73.94% e un Cons@32 dell'86.67%.
• Diversità delle Tracce: L'analisi visiva e quantitativa (Figura 4 e 5 del paper) mostra che i diversi token di forking (<think 1> vs <think 6>) attivano effettivamente strategie di ragionamento diverse con lunghezze e approcci differenti, mentre i modelli SFT standard o SSFT con matching casuale non mostrano questa diversità.
• Generalizzazione: I benefici si estendono anche a compiti out-of-distribution (come la generazione di codice su LiveCodeBench) e a modelli di dimensioni diverse (4B, 8B), sebbene i guadagni siano più marcati sui modelli più grandi.

5. Significato e Impatto

Questo lavoro rappresenta un passo avanti significativo nel campo del test-time compute e del ragionamento degli LLM:

• Superamento del compromesso Diversità-Accuratezza: Fornisce un metodo per ottenere ragionamenti paralleli diversificati senza sacrificare l'accuratezza, risolvendo un limite fondamentale delle tecniche di scaling attuali.
• Efficienza nell'Addestramento: Utilizza una perdita basata su insiemi e matching bipartito per apprendere in modo robusto le correlazioni tra prompt e strategie di ragionamento, anche con dati distillati a basso costo.
• Nuovo Paradigma di Inferenza: Introduce la possibilità di controllare esplicitamente la modalità di ragionamento di un modello tramite token specifici, aprendo la strada a sistemi più controllabili e affidabili per compiti complessi (agentic AI).
• Scalabilità: La metodologia è applicabile a diversi modelli e domini (matematica, codice), dimostrando che l'apprendimento di "token di controllo" globali è una strategia generale per migliorare le capacità di ragionamento.

In sintesi, il paper dimostra che trasformare il ragionamento parallelo in un problema di predizione di insiemi con matching ottimizzato permette di "sbloccare" e stabilizzare modalità di ragionamento complesse e distinte, portando a prestazioni superiori sia in termini di accuratezza singola che di copertura delle soluzioni.