💬 NLP

Masked Distillation: Internalizing the Chain-of-Thought in Language Models

Questo articolo introduce la "distillazione mascherata" (masked distillation), un framework che interiorizza le capacità di ragionamento dei Large Reasoning Models nei parametri di un modello studente addestrandolo a predire solo i token della soluzione pur ricevendo un feedback condizionato dalla Chain-of-Thought del docente, consentendo così la generazione diretta di risposte e riducendo la latenza di inferenza senza sacrificare l'accuratezza.

Autori originali: Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

Pubblicato 2026-07-28

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Durgesh Kalwar, Vardhan Palod, Subbarao Kambhampati

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immaginate un mondo in cui il vostro assistente IA preferito sia come uno studente brillante ma chiacchierone. Quando gli ponete una domanda difficile, non si limita a sbizzarrire la risposta; trascrive pagine di appunti, disegna diagrammi e spiega ogni singolo passaggio del suo ragionamento prima di scrivere finalmente la soluzione. Questo processo di "pensare ad alta voce", noto nel mondo tecnologico come "Chain of Thought" (Catena di Pensiero), ha reso i modelli di IA incredibilmente intelligenti nei puzzle matematici e logici. Tuttavia, c'è un problema: quelle pagine di appunti richiedono molto tempo ed energia per essere generate. È come pagare un tassista per portarti al supermercato, ma pagarlo anche per girare intorno all'intero isolato prima di capire in quale corsia hai bisogno di andare. Gli scienziati si sono spesso chiesti: possiamo insegnare all'IA di fare tutto quel pensiero dentro la sua testa? Se ci riuscissimo, l'IA salterebbe i lunghi e chiacchieroni appunti e darebbe semplicemente la risposta istantaneamente, risparmiando una quantità enorme di tempo e potenza di calcolo.

Questo articolo, intitolato "Masked Distillation" (Distillazione Mascherata), affronta proprio questa domanda. I ricercatori, lavorando con i Large Reasoning Models (Modelli di Ragionamento Grandi), si sono chiesti se potessero prendere un modello "insegnante" intelligente che ama pensare ad alta voce e addestrare un modello "studente" più piccolo per interiorizzare quel processo di pensiero. Volevano vedere se lo studente potesse imparare a saltare del tutto gli appunti e fornire solo la risposta finale, pur rimanendo altrettanto intelligente. Hanno testato questo approccio utilizzando una tecnica ingegnosa chiamata "Masked Distillation". Immaginate che l'insegnante stia risolvendo un puzzle mentre lo studente osserva. L'insegnante vede la domanda completa e il proprio processo di pensiero dettagliato, ma allo studente viene mostrata solo la domanda. L'obiettivo era vedere se lo studente potesse apprendere la logica segreta dell'insegnante e risolvere il puzzle da solo senza mai dover scrivere i passaggi, anche se l'insegnante aveva tutti gli appunti per guidare l'apprendimento.

Il team ha predisposto due esperimenti principali. Nel primo, hanno utilizzato lo stesso modello di IA sia come insegnante che come studente, chiedendogli semplicemente di pensare in modo diverso. Nel secondo, hanno usato un modello più grande e intelligente come insegnante e uno più piccolo e semplice come studente. Hanno anche introdotto un concetto di "scaffold" (impalcatura). Pensate a questo come alle rotelle di supporto su una bicicletta. A volte, lo studente potrebbe aver bisogno di vedere un po' degli appunti dell'insegnante (lo "scaffold") per arrivare alla risposta corretta, invece di provare a fare tutto da zero. Hanno testato quanto del pensiero dell'insegnante lo studente dovesse vedere per ottenere i migliori risultati.

I risultati sono stati un misto di "wow" e "non così in fretta". Su un dataset di problemi matematici della scuola elementare chiamato GSM8K, i modelli studenti sono stati straordinari. Sono riusciti ad apprendere a interiorizzare il pensiero dell'insegnante. Lo studente completamente "mascherato" (uno che non ha mai scritto appunti) ha ottenuto il 76,0% delle risposte corrette, il che è un enorme salto rispetto alle sue capacità iniziali, e lo ha fatto usando in media solo 369,6 token (unità di testo), rispetto ai 2.398,1 token dell'insegnante. Questa è una riduzione dello sforzo di ben 6,5 volte!

Tuttavia, su un puzzle numerico più complicato chiamato Countdown, la storia cambia. In questo caso, lo studente che cercava di fare tutto nella sua testa (completamente mascherato) è in realtà peggiorato, scendendo al 41,7% di accuratezza. Si è scoperto che per questo specifico tipo di puzzle, lo studente aveva bisogno di vedere alcuni degli appunti dell'insegnante per avere successo. Ma ecco la parte interessante: i ricercatori hanno trovato il punto di equilibrio ideale. Lasciando che lo studente vedesse solo una piccola parte del pensiero dell'insegnante (circa il 30% degli appunti, o un "alfa" di 0,3), l'accuratezza dello studente è schizzata all'86,2%. Era quasi pari all'insegnante (87,3%), ma utilizzava comunque circa 1,3 volte meno token rispetto all'insegnante. Era come dare allo studente solo le rotelle necessarie per bilanciarsi perfettamente senza aver bisogno dell'intera bicicletta.

L'articolo ha anche esaminato quanto bene questi studenti gestissero nuovi puzzle mai visti prima (compiti out-of-distribution). Sul lato matematico, gli studenti sono andati bene, ma sui puzzle di Countdown, i risultati sono stati un po' più contrastanti. Interessante è che, in alcuni casi, gli studenti addestrati con un po' di impalcatura (scaffolding) hanno generalizzato meglio dell'insegnante stesso quando i puzzle diventavano più difficili o cambiavano leggermente.

Un risultato chiave è stato che il metodo di insegnamento contava. Quando hanno provato a insegnare allo studente semplicemente mostrandogli gli appunti dell'insegnante e chiedendogli di copiarli (un metodo chiamato Supervised Fine-Tuning), non ha funzionato altrettanto bene. Il metodo "Masked Distillation", che utilizza un modo più complesso di confrontare le ipotesi dello studente con i pensieri effettivi dell'insegnante, è stato molto più efficace nel trasferire le capacità di "pensiero".

In conclusione, l'articolo suggerisce che non possiamo semplicemente cancellare magicamente tutti i passaggi di pensiero per ogni problema. Se un'IA possa "pensare silenziosamente" dipende fortemente dal tipo di problema e da quanto lo studente conosca già l'argomento. Per compiti familiari come la matematica di base, l'IA può imparare a pensare silenziosamente e risparmiare un sacco di tempo. Per puzzle più difficili e sconosciuti, potrebbe ancora aver bisogno di un piccolo aiuto — una piccola impalcatura di appunti — per svolgere il lavoro correttamente. I ricercatori concludono che questo "scaffold" è uno strumento potente che possiamo regolare, permettendoci di trovare il perfetto equilibrio tra quanto l'IA sia intelligente e quanto velocemente possa darci una risposta.

Riassunto Tecnico: Masked Distillation – Internalizzare la Chain-of-Thought nei Modelli di Linguaggio

Problema

I Large Reasoning Models (LRM) ottengono prestazioni elevate su compiti complessi (matematica, codice, pianificazione) generando lunghe, esplicite catene di token intermedi (Chain-of-Thought o CoT) prima di produrre una risposta finale. Sebbene queste tracce migliorino l'accuratezza, esse comportano costi significativi in termini di latenza di inferenza, uso della memoria (impronta del KV-cache) ed energia di servizio, che scalano linearmente con la lunghezza della traccia.

Recenti scoperte suggeriscono che la correttezza di queste tracce intermedie non è causalmente legata alla correttezza della risposta finale, né la lunghezza della traccia correla in modo affidabile con la complessità del problema. Ciò solleva una domanda critica: è possibile internalizzare il calcolo espresso in questi token intermedi nei parametri di un modello di linguaggio, permettendogli di produrre risposte corrette direttamente (o con tracce significativamente più brevi) senza generare esplicitamente i passaggi di ragionamento durante l'inferenza?

Metodologia: Masked Distillation

Gli autori propongono la Masked Distillation, un framework di distillazione della conoscenza progettato per internalizzare il processo di ragionamento di un modello "insegnante" (teacher) in un modello "non pensante" (student).

Framework Core

In questo framework, il modello studente è addestrato a predire solo i token della soluzione condizionati dalla domanda di input. Il modello insegnante, tuttavia, fornisce un feedback basato sia sulla domanda che sulla propria traccia di CoT generata. L'obiettivo è minimizzare la divergenza tra la distribuzione dello studente (condizionata solo dall'input) e la distribuzione dell'insegnante (condizionata dall'input e dai token di ragionamento intermedi o ITs).

L'obiettivo di addestramento utilizza la divergenza reverse-KL on-policy:
$L_{MD} = D_{KL}(\pi_S(\cdot | x, y_{<t}) \parallel \pi_T(\cdot | x, \text{ITs}, y_{<t}))$
Dove $\pi_S$ è la policy dello studente, $\pi_T$ è la policy dell'insegnante, $x$ è l'input e $y_{<t}$ sono i token precedentemente generati e $\text{ITs}$ sono i token intermedi dell'insegnante.

Impostazioni Sperimentali

Gli autori istanziano questo framework in due scenari:

Self-Distillation: Lo stesso modello funge da insegnante (in modalità di pensiero) e da studente (in modalità non pensante).
Dual-Model Setting: Un modello di ragionamento più grande funge da insegnante, supervisionando uno studente separato, più piccolo e non pensante.

Lo Scaffold $\alpha$ -Suffix

Per esplorare il compromesso tra piena internalizzazione e costo di inferenza, gli autori introducono una variante della $\alpha$ -suffix masked distillation.

$\alpha = 0$ (Completamente Mascherata): Lo studente internalizza l'intera traccia intermedia ed emette solo la soluzione.
$\alpha = 1$ (Non Mascherata): Lo studente riproduce l'intera traccia prima della risposta (equivalente alla generazione standard di CoT).
$0 < \alpha < 1$ : Lo studente internalizza il prefisso $(1-\alpha)$ della traccia dell'insegnante ed è addestrato a emettere la frazione di suffisso $\alpha$ rimanente della traccia durante l'inferenza, insieme alla soluzione.

Ciò permette agli autori di trattare i token intermedi come uno "scaffold" che può essere parzialmente mantenuto o completamente internalizzato.

Contributi Chiave

Framework Proposto: Introduzione della Masked Distillation per internalizzare il ragionamento CoT nei parametri del modello, consentendo la generazione diretta della risposta.
Analisi dello Scaffold: Definizione e valutazione del parametro $\alpha$ -suffix, stabilendo un asse controllato per regolare il compromesso tra accuratezza e costo di inferenza.
Valutazione Comprensiva: Conduzione di estesi esperimenti in due domini:
- Matematica: GSM8K (In-Distribution), MATH-500, e AIME-25 (Out-of-Distribution).
- Countdown: Un compito di ricerca in puzzle numerici con variazioni nel range target e nella profondità di ricerca (split OOD).
Studi di Ablazione: Confronto della Masked Distillation (usando la reverse-KL) rispetto al Supervised Fine-Tuning (SFT) sulle tracce dell'insegnante e all'uso di "teaching tokens" espliciti come scaffold.

Risultati

1. Dipendenza dal Task della Piena Internalizzazione

La capacità di internalizzare completamente il ragionamento ( $\alpha=0$ ) dipende fortemente dal task e dall'esposizione pregressa del modello al dominio durante il pretraining:

GSM8K (Matematica): La piena internalizzazione ha successo. Nella self-distillation, lo studente ha raggiunto il 76,0% di accuratezza (rispetto all'87,9% dell'insegnante) riducendo i token di inferenza di 6,5×. Nella dual-model setting, l'accuratezza è migliorata dal 15,69% (base) al 45,26% con una riduzione dei token di 5×.
Countdown: La piena internalizzazione fallisce senza scaffolding. Nella self-distillation, lo studente completamente mascherato è sceso al 41,7% di accuratezza (al di sotto del 56,2% del modello base). Nella dual-model setting, è migliorato dallo 0% al 34,1%, ma è rimasto molto al di sotto dell'insegnante (87,3%) e della variante non mascherata (81,1%).

2. Efficacia dello Scaffolding $\alpha$ -Suffix

Fornire una piccola quantità di scaffolding (emettendo una frazione della traccia) recupera significativamente le prestazioni, in particolare nei task in cui la piena internalizzazione fallisce:

Countdown (Self-Distillation): Impostare $\alpha = 0,3$ ha elevato l'accuratezza dal 41,7% all'86,2%, eguagliando essenzialmente l'insegnante (87,3%) pur riducendo i token di inferenza di circa 1,3× rispetto alla variante non mascherata.
Countdown (Dual-Model): La variante $\alpha = 0,5$ ha raggiunto l'84,2% di accuratezza, superando la baseline non mascherata (81,1%) pur utilizzando meno token, stabilendo un punto Pareto-ottimale.
GSM8K: Piccoli scaffold ( $\alpha=0,3$ ) hanno recuperato la maggior parte del divario di accuratezza (83,2% vs 87,9%) con circa la metà del costo di token della variante non mascherata.

3. Generalizzazione a Out-of-Distribution (OOD)

Spostamento del Range Target: I modelli addestrati con vari paradigmi di scaffolding si sono trasferiti correttamente ai task OOD dove il range numerico target è cambiato.
Spostamento della Profondità di Ricerca: Nella self-distillation sul task Countdown, gli studenti addestrati con scaffold (specificamente $\alpha=0,5$ ) si sono generalizzati meglio della variante non mascherata sotto un aumento della profondità di ricerca (32,8% vs 25,7%).
Limitazioni del Dual-Model: Nella dual-model setting su benchmark matematici più difficili (MATH-500, AIME-25), il gap di capacità tra insegnante e studente era troppo ampio; né la piena internalizzazione né lo scaffolding hanno recuperato le prestazioni OOD, con tutte le varianti che sono crollate al livello base dello studente.

4. Obiettivi di Addestramento

Reverse-KL vs. SFT: Il paper trova che il Supervised Fine-Tuning (SFT) sulle tracce dell'insegnante non riesce a sfruttare efficacementamente lo scaffolding. In Countdown, l'SFT con scaffold ha prodotto circa il 38% di accuratezza, mentre la Masked Distillation con reverse-KL con gli stessi scaffold ha raggiunto circa il 67%.
Teaching Tokens: Fornire "teaching tokens" espliciti come scaffold ha migliorato le prestazioni rispetto alla piena mascheratura, ma è stato meno efficace dell'approccio $\alpha$ -suffix.

Significato e Affermazioni

Il paper afferma che lo scaffolding del suffisso fornisce un meccanismo controllabile per bilanciare accuratezza ed efficienza di inferenza. Il punto operativo ottimale è determinato dal task specifico e dal gap di capacità tra l'insegnante e lo studente.

L'internalizzazione è possibile ma non universale: I modelli possono internalizzare il ragionamento per task con cui hanno familiarità (come GSM8K), ma faticano con nuovi task di ricerca (come Countdown) senza mantenere una traccia esplicita.
Compromesso Costo-Prestazioni: Una piccola quantità di ragionamento esplicito (es. $\alpha=0,3$ ) può recuperare quasi tutte le prestazioni della generazione completa di CoT offrendo al contempo sostanziosi risparmi nei costi di inferenza (riduzione dei token di 1,3× a 2×).
Generalizzazione: Lo scaffolding può migliorare la generalizzazione in specifici scenari OOD (es. spostamenti della profondità di ricerca nella self-distillation), suggerendo che mantenere una traccia parziale aiuta il modello a navigare meglio gli spazi di ricerca complessi rispetto a un modello completamente internalizzato o a una traccia completa in certi contesti.

Gli autori concludono che, sebbene la piena internalizzazione sia un obiettivo attraente per l'efficienza, l'approccio $\alpha$ -suffix offre un'alternativa pratica e modulabile che mantiene alte prestazioni con una latenza ridotta.