Do LLMs Encode Functional Importance of Reasoning Tokens?

1. Il Problema

I Large Language Models (LLM) risolvono compiti complessi (matematica, scienze, ragionamento decisionale) generando lunghe catene di ragionamento (Chain-of-Thought). Sebbene queste catene migliorino l'accuratezza, comportano costi computazionali elevati (latenza, memoria) e rendono difficile isolare le parti del ragionamento funzionalmente rilevanti per la risposta finale.

Le ricerche precedenti sulla "ragionamento compatto" (compact reasoning) si sono concentrate sulla generazione di catene più corte tramite campionamento probabilistico, euristiche o supervisione da modelli all'avanguardia (frontier models). Tuttavia, questi approcci offrono scarse intuizioni su se i modelli LLM codifichino internamente un'importanza funzionale a livello di token per la generazione della risposta. Il lavoro si pone la domanda fondamentale: I LLM codificano internamente l'importanza funzionale dei token di ragionamento?

2. Metodologia: Greedy Pruning

Gli autori introducono un metodo diagnostico chiamato Greedy Pruning (Potatura Greedy), una procedura di cancellazione che preserva la verosimiglianza (likelihood-preserving).

• Concetto Base: A differenza del greedy decoding che aggiunge token per massimizzare la verosimiglianza, il greedy pruning parte da una catena di ragionamento completa e rimuove iterativamente i token la cui cancellazione degrada minimamente la verosimiglianza del modello rispetto a un obiettivo specifico.
• Framework Teacher-Pruner-Student:
  1. Teacher: Genera una catena di ragionamento completa.
  2. Pruner: Valida la catena rimuovendo token uno alla volta. Per ogni token candidato, calcola la verosimiglianza post-cancellazione ($L_{del}$). Il token con la diminuzione di verosimiglianza minima viene rimosso. Questo processo genera un ranking globale dei token basato sulla loro importanza funzionale.
  3. Student: Viene addestrato (via SFT - Supervised Fine-Tuning) sulle catene di ragionamento potate.
• Obiettivi di Potatura:
  • L_ANS: Preserva la verosimiglianza della sola risposta finale.
  • L_JOINT: Preserva la verosimiglianza congiunta del ragionamento e della risposta (più restrittivo, favorisce catene coerenti con la traiettoria di ragionamento del modello).

3. Contributi Chiave

1. Nuova Prospettiva Diagnostica: Invece di proporre un nuovo metodo di compressione basato su esplorazione esterna, lo studio tratta la compressione come un problema diagnostico per rivelare la struttura funzionale interna codificata dal modello.
2. Greedy Pruning: Un algoritmo semplice ma efficace che genera un ranking di importanza dei token basato sulla sensibilità della verosimiglianza alla rimozione, senza richiedere supervisione esterna (come GPT-4).
3. Analisi Strutturale: Dimostrazione che i modelli LLM possiedono una struttura funzionale non banale: i token critici per il calcolo simbolico sono preservati, mentre il "scaffolding" linguistico (riferimenti, descrizioni) viene rimosso per primo.
4. Predicibilità tramite Attenzione: Evidenza che i punteggi di attenzione del modello possono predire con alta accuratezza l'ordine di potatura, suggerendo che i segnali di importanza sono accessibili internamente durante la generazione.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su benchmark come GSM8K, MATH e MMLU-Pro, utilizzando modelli come Llama-3.1, Qwen-2.5, Llama-2 e Mistral.

• Performance di Distillazione: Gli studenti addestrati su catene potate con il metodo Greedy superano significativamente le baseline (incluso TokenSkip, che usa supervisione da modelli frontier, e metodi basati su Surprisal o H2O) a parità di lunghezza di ragionamento. Questo conferma che il pruning preserva i token funzionalmente importanti.
• Struttura Funzionale:
  • I token di tipo SYMBMATH (calcoli simbolici, equazioni) vengono preservati in modo sproporzionato anche con potature aggressive.
  • I token COREF (riferimenti, pronomi) e GRAMMAR vengono rimossi precocemente.
  • L'uso dell'obiettivo JOINT produce una struttura funzionale più nitida rispetto all'obiettivo ANS (solo risposta).
  • La forza del modello "Pruner" influenza la qualità della struttura: un pruner più debole mantiene il calcolo simbolico ma disturba l'equilibrio delle strutture non simboliche, riducendo le prestazioni dello studente.
• Dinamicità: L'importanza dei token non è statica; viene rivalutata dinamicamente man mano che il contesto si contrae. Un ranking "congelato" (calcolato all'inizio) è meno efficace di un ranking dinamico locale.
• Predizione tramite Attenzione: Un semplice MLP addestrato sugli score di attenzione del modello teacher raggiunge una correlazione di Pearson di 0.88 con i veri punteggi di potatura, indicando che l'attenzione codifica informazioni sufficienti per prevedere l'importanza funzionale.
• Qualità e Leggibilità: Anche con riduzioni del 30% dei token, la degradazione è principalmente superficiale (mancanza di punteggiatura o articoli), mentre la coerenza semantica e lo stato matematico rimangono intatti.

5. Significato e Implicazioni

Questo lavoro stabilisce che i LLM non generano solo sequenze di testo casuali, ma codificano una struttura funzionale gerarchica all'interno delle loro catene di ragionamento.

• Compressione Efficiente: Il Greedy Pruning offre un metodo controllato e principiato per comprimere il ragionamento, superando i metodi basati su euristiche o supervisione esterna.
• Strumento Diagnostico: Fornisce un nuovo strumento per "probing" (sondaggio) della cognizione interna dei modelli, rivelando quali parti del ragionamento sono essenziali per la predittività.
• Futuro: I risultati suggeriscono che i segnali di importanza potrebbero essere utilizzati per ottimizzare l'inferenza in tempo reale o per creare curricula di addestramento più efficienti, spostando il focus dalla semplice compressione alla comprensione della struttura interna del ragionamento dei modelli.

In sintesi, il paper dimostra che è possibile "potare" il ragionamento di un LLM in modo intelligente preservando la sua capacità di risolvere problemi, e che questa capacità di distinguere l'importante dal ridondante è intrinseca al modello stesso, riflettibile attraverso la sua verosimiglianza e i suoi meccanismi di attenzione.