Latent Reasoning with Supervised Thinking States

Riassunto Tecnico: Latent Reasoning with Supervised Thinking States

1. Il Problema: Il costo del ragionamento (CoT)

Il ragionamento tramite Chain-of-Thought (CoT) permette ai Large Language Models (LLM) di risolvere compiti complessi generando passaggi intermedi in linguaggio naturale. Tuttavia, questo approccio presenta due svantaggi principali:

• Costo di inferenza elevato: La generazione di lunghe sequenze di token di ragionamento aumenta significativamente la latenza e il consumo computazionale.
• Efficienza del contesto: I token di CoT occupano spazio nella finestra di contesto, limitando la lunghezza della sequenza che il modello può gestire.

Esistono metodi di "ragionamento latente" (dove il pensiero è rappresentato da embedding continui), ma questi sono difficili da supervisionare e richiedono spesso il Backpropagation Through Time (BPTT), rendendo l'addestramento estremamente lento e difficile da scalare.

2. Metodologia: Thinking States

Gli autori propongono Thinking States, un'architettura che permette al modello di ragionare mentre elabora l'input, senza aggiungere token di testo al contesto.

Componenti Architetturali:
Il modello integra un LLM di base con due moduli aggiuntivi:

1. Thinking Block ($T$): Un decoder Transformer leggero che genera una sequenza di "token di pensiero" in linguaggio naturale partendo dalle rappresentazioni profonde dell'LLM.
2. Compression Block ($C$): Un encoder Transformer che comprime la sequenza variabile di token di pensiero in uno stato (state) a dimensione fissa ($S$).

Il Processo Ricorrente:
L'input viene diviso in blocchi (chunks). Per ogni blocco $i$:

• Lo stato precedente $S_i$ viene iniettato nelle rappresentazioni dei token del blocco corrente a un livello superficiale (shallow layer) dell'LLM.
• L'LLM elabora il blocco e produce rappresentazioni profonde.
• Il Thinking Block estrae il ragionamento da queste rappresentazioni.
• Il Compression Block trasforma il ragionamento in un nuovo stato $S_{i+1}$ da passare al blocco successivo.

Addestramento con Supervisione (Teacher-Forcing):
A differenza dei metodi latenti continui, Thinking States utilizza la supervisione in linguaggio naturale. Gli autori allineano i passaggi di un CoT esistente (generato da un modello teacher) con posizioni specifiche dell'input. Durante l'addestramento, si utilizza il teacher-forcing: si iniettano gli stati "gold" (derivati dai ragionamenti corretti) nel modello. Questo permette un addestramento completamente parallelo, eliminando la necessità di BPTT e riducendo drasticamente i tempi di training.

3. Contributi Chiave

• Ragionamento Ricorrente senza espansione del contesto: Il modello mantiene la capacità di ragionamento sequenziale senza aumentare la lunghezza della sequenza di input.
• Efficienza nell'addestramento: Grazie all'uso di stati supervisionati e teacher-forcing, l'addestramento è parallelizzabile e non soffre della complessità computazionale del BPTT.
• Interpretabilità: Poiché i "pensieri" sono generati in linguaggio naturale prima di essere compressi, il processo di ragionamento rimane ispezionabile dall'uomo.
• Speculative Thinking (Prefill veloce): Un algoritmo di prefill che specula su stati "trivial" (senza ragionamento) per accelerare la fase iniziale di elaborazione della query.

4. Risultati Sperimentali

Il modello è stato testato su compiti di State Tracking (es. Parità, Assegnazione di variabili) e ragionamento generale (GSM8K, Multi-hop QA).

• State Tracking: Thinking States mostra una capacità di estrapolazione della lunghezza superiore al CoT. Mentre il CoT fatica quando la sequenza di operazioni supera quella vista in addestramento, Thinking States mantiene un'accuratezza vicina al 100% grazie alla sua natura ricorrente.
• Ragionamento Generale:
  • Su GSM8K, supera i metodi di ragionamento latente esistenti (come Coconut e iCoT) e si avvicina alle prestazioni del CoT, offrendo però un guadagno di velocità (speedup) di circa 2.66x.
  • Su Multi-hop QA, eguaglia le prestazioni del CoT con una latenza significativamente inferiore.
• Analisi degli Errori: Il modello è meno soggetto ad allucinazioni di passaggi intermedi rispetto al CoT, ma può soffrire di "ambiguità dello stato" se la domanda cruciale si trova alla fine della query (problema risolvibile spostando la domanda all'inizio).

5. Significato e Conclusioni

Il lavoro di Amos et al. rappresenta un passo fondamentale verso LLM più efficienti. Dimostra che è possibile "interiorizzare" il ragionamento senza sacrificare l'accuratezza o l'interpretabilità. La capacità di comprimere il ragionamento in stati latenti supervisionati apre la strada a modelli che possono ragionare in modo profondo con un costo computazionale molto più contenuto rispetto ai modelli che generano testualmente ogni singolo passaggio logico.