SuperThoughts: Reasoning Tokens in Superposition

Riepilogo Tecnico: SuperThoughts: Token di Ragionamento in Superposizione

Definizione del Problema

I Large Language Models (LLM) risolvono problemi complessi generando sequenze esplicite di Chain-of-Thought (CoT), dove ogni token rappresenta un'unità di computazione (un passaggio forward). Sebbene modelli recenti come OpenAI o1 e DeepSeek-R1 dimostrino che l'aumento del calcolo durante il test migliora significativamente le prestazioni, questa generazione sequenziale di token è computazionalmente costosa.

Gli approcci esistenti che tentano di bypassare la generazione discreta di token per ragionare in spazi latenti continui (ad esempio, COCONUT, ragionamento latente) affrontano una sfida critica: la mancanza di supervisione intermedia. L'addestramento standard della CoT si basa sulla perdita di cross-entropy a livello di token in ogni passaggio, fornendo segnali di gradiente densi. Quando il ragionamento avviene in uno spazio latente non vincolato, questa supervisione svanisce, portando a instabilità nell'addestramento, deriva rappresentazionale e scarse prestazioni su compiti a lungo orizzonte. Al contrario, i metodi che si limitano a comprimere il contesto o a ridurre il numero di token discreti tramite fine-tuning non riducono il numero di passaggi forward richiesti, limitandone l'impatto sul throughput di inferenza.

La domanda centrale affrontata da questo lavoro è: È possibile addestrare un modello a ragionare in uno spazio latente più ricco, mantenendo al contempo i segnali di supervisione intermedia della generazione di token discreti?

Metodologia

Gli autori propongono SuperThoughts, un framework che comprime coppie di token CoT consecutivi in singole rappresentazioni latenti, permettendo al modello di decodificare due token per passaggio. Questo approccio dimezza il numero di passaggi forward richiesti per una catena di ragionamento di lunghezza $L$ (riducendolo a $L/2$ passaggi) preservando la supervisione dei token discreti.

Architettura

Il modello è composto da tre componenti primari:

1. Compressore: Codifica una coppia di embedding di token $(c_{2i-1}, c_{2i})$ in un singolo vettore latente $x_i$. Gli autori esplorano due implementazioni:
   • Proiezione Lineare: Concatena gli embedding e li proietta tramite una matrice apprendibile.
   • Blocco Transformer: Elabora la coppia attraverso un piccolo strato Transformer, estraendo lo stato nascosto del secondo token.
2. Modulo Principale (Main Module): Il backbone di ragionamento primario (inizializzato da un LLM pre-addestrato). Elabora il vettore compresso $x_i$ per produrre uno stato nascosto $h_i$ e predire il token successivo con indice dispari ($c_{2i+1}$).
3. Modulo Multi-Token Prediction (MTP): Un modulo leggero che predice il token successivo con indice pari ($c_{2i+2}$). Prende come input il precedente token pari, il nuovo token dispari appena predetto e lo stato nascosto del modulo Main, proiettandoli in un vettore $x'_i$ per generare la successiva predizione. Entrambi i moduli condividono lo stesso modulo di linguaggio (language modeling head) in uscita.

Strategia di Addestramento

Per affrontare l'instabilità del ragionamento latente, viene impiegato un protocollo di addestramento in due fasi:

1. Distillazione Latente (Fase 1): Il Compressore viene addestrato per allineare lo spazio latente compresso con il manifold del linguaggio pre-addestrato. Un modello "insegnante" elabora la sequenza di token discreti, mentre un modello "studente" elabora le coppie compresse. Il Compressore è ottimizzato per minimizzare la distanza Smooth-L1 tra gli stati nascosti dello studente (dopo l'elaborazione di $x_i$) e gli stati nascosti dell'insegnante (dopo l'elaborazione della corrispondente coppia di token). Ciò assicura che le rappresentazioni compresse preservino lo stato di ragionamento della sequenza discreta originale.
2. Addestramento Congiunto End-to-End (Fase 2): Il modulo Main, il modulo MTP e il Compressore vengono addestrati congiuntamente utilizzando la perdita di cross-entropy standard per tutti i token predetti (sia per l'output del Main che per quello dell'MTP).

Inferenza Adattiva

Riconoscendo che non tutti i passaggi di ragionamento sono ugualmente difficili, gli autori introducono un meccanismo adattivo basato sulla confidenza. Ad ogni passaggio, se la confidenza della predizione del modulo MTP (probabilità massima) scende al di sotto di una soglia $\tau$, il modello rifiuta la predizione MTP. In tal caso, torna alla decodifica standard per quel passaggio, fornendo il singolo token direttamente al modulo Main. Ciò consente al modello di allocare più calcolo ai token "difficili" pur mantenendo l'efficienza su quelli "facili".

Contributi Chiave

1. Architettura SuperThoughts: Un nuovo framework che fonde coppie di token in embedding in superposizione e decodifica due token per passaggio utilizzando un modulo Main e un modulo MTP, dimezzando efficacemente i passaggi forward richiesti.
2. Protocollo di Addestramento in Due Fasi: Un metodo per stabilizzare il ragionamento latente allineando prima lo spazio compresso tramite distillazione e successivamente procedendo con l'addestramento end-to-end con supervisione discreta.
3. Meccanismo di Inferenza Adattiva: Un sistema di fallback dinamico che scambia il throughput con l'accuratezza sui passaggi di ragionamento incerti, prevenendo la degradazione delle prestazioni su compiti complessi.
4. Valutazione Empirica: Valutazione completa su benchmark di ragionamento matematico e scientifico (MATH500, AMC23, OlympiadBench, GPQA-Diamond) che dimostra una significativa riduzione della lunghezza della CoT con una perdita minima di accuratezza.

Risultati

Gli autori hanno valutato SuperThoughts su modelli Qwen2.5-Math (varianti da 1.5B, 7B e 14B) su molteplici benchmark:

• Riduzione della Lunghezza: Con il meccanismo adattivo, SuperThoughts ottiene una riduzione del 20–35% della lunghezza della CoT.
• Accuratezza: Il metodo mantiene l'accuratezza entro 1–2 punti rispetto alla baseline standard della CoT discreta nella maggior parte dei compiti. Ad esempio, sul modello da 1.5B con $\tau=0.999$, l'accuratezza su MATH500 ha raggiunto il 73.0% (rispetto al 72.4% della baseline) con una riduzione della lunghezza della CoT del 36%.
• Effetti di Scala: I modelli più grandi (7B) tollerano meglio la compressione aggressiva rispetto a quelli più piccoli (1.5B). Senza l'inferenza adattiva, il modello da 1.5B ha subito cali di accuratezza significativi (14–21 punti), mentre il modello da 7B è sceso solo di 5–12 punti.
• Confronto con HAMburger: SuperThoughts ha superato la baseline HAMburger (Liu & Zhang, 2025) a ogni livello di compressione, ottenendo un'accuratezza superiore per la stessa lunghezza di CoT o lunghezze di CoT più brevi per la stessa accuratezza.
• Accelerazione dell'Inferenza: Sebbene lo speedup teorico sia proporzionale alla riduzione della lunghezza, lo speedup effettivo del tempo di esecuzione (wall-clock time) è leggermente inferiore a causa dell'overhead dei moduli Compressor e MTP. Tuttavia, questo overhead diminuisce all'aumentare della dimensione del modello; per il modello da 14B, una riduzione della lunghezza del ~33% si è tradotta in una riduzione del ~28% del tempo di generazione.

Significato

Il paper sostiene che SuperThoughts dimostri con successo che i modelli possono ragionare in uno spazio latente in superposizione senza sacrificare la stabilità fornita dalla supervisione di token discreti. Comprimendo le coppie di token e utilizzando l'inferenza adattiva, il framework offre una via pratica per ridurre il costo computazionale dei compiti di ragionamento a lungo orizzonte. Il lavoro suggerisce che i modelli più grandi possiedano una capacità per passaggio sufficiente per gestire in modo affidabile i token in superposizione, lasciando intendere che lo scaling futuro potrebbe consentire una compressione ancora maggiore (ad esempio, superponendo tre o più token). Inoltre, gli autori osservano che, man mano che i moduli MTP nativi diventeranno standard nei futuri LLM (ad esempio, Qwen3-Next, MiMo), l'applicazione di SuperThoughts diventerà ancora più semplice ed efficace.