Riepilogo Tecnico: Auto-distillazione Geometrica per la Generalizzazione del Ragionamento

1. Definizione del Problema

La distillazione on-policy (OPD) è emersa come una strategia pratica di post-training per i grandi modelli linguistici (LLM), fornendo una supervisione densa a livello di token sulle traiettorie generate dal modello studente. Una variante specifica ed efficace, l'auto-distillazione con contesto privilegiato (OPSD), condiziona il "insegnante" (lo stesso modello) su informazioni aggiuntive, come una traccia completa della soluzione o un suggerimento, mentre lo studente deve ragionare partendo solo dal contesto originale del problema.

Sebbene questa configurazione fornisca una supervisione abbondante, il paper identifica un fallimento critico: il drift predittivo (predictive drift). Poiché l'insegnante possiede un contesto privilegiato, può impegnarsi con fiducia in continuazioni che sono ovvie data la soluzione, ma che rimangono incerte o non supportate dalla prospettiva dello studente. Gli obiettivi di distillazione standard (ad esempio, la divergenza di Kullback–Leibler) pesano le preferenze dell'insegnante in base alla propria massa di probabilità. Di conseguenza, negli stati in cui l'insegnante è altamente fiducioso ma lo studente assegna una bassa probabilità (bassa sovrapposizione), il segnale del gradiente è forte. Nel corso di molti aggiornamenti, queste spinte intense costringono lo studente ad adottare le continuazioni sicure dell'insegnante prima che le informazioni dello studente stesso possano giustificarle. L'accumulo di questi aggiornamenti disallineati porta al drift, dove il modello sopprime l'incertezza e degrada le prestazioni di ragionamento fuori distribuzione (OOD), risultando spesso in un "accordo fiducioso su risposte errate".

Gli approcci esistenti tentano di mitigare questo fenomeno filtrando i segnali inaffidabili dell'insegnante o rimodellando i pesi della divergenza (ad esempio, usando la Reverse KL o meccanismi di filtraggio). Tuttavia, gli autori sostengono che tali metodi non riescano ad affrontare il problema centrale: quanto lontano un segnale affidabile debba spostare lo studente. La sola selezione non può regolare l'entità dell'aggiornamento, e un movimento incontrollato si accumula in un drift globale.

2. Metodologia: GEOSD

Gli autori introducono GEOSD (Geometric Self-Distillation), un obiettivo progettato per contrastare questo drift trattando il comportamento predittivo come un movimento all'interno di uno spazio geometrico specifico di distribuzioni del prossimo token. Il metodo opera in due modi complementari:

A. Attenuazione della Pressione Locale (Perdita di Hellinger)

Invece della standard divergenza KL, GEOSD impiega una perdita di Hellinger per modulare la spinta dell'insegnante in ogni stato visitato.

• Meccanismo: La distanza di Hellinger è derivata dalla geometria del simplesso di probabilità sotto un embedding radice quadrata ($p \to \sqrt{p}$), dove le distribuzioni giacciono su una ipersfera. Il gradiente della perdita di Hellinger scala ogni preferenza dell'insegnante tramite la media geometrica delle probabilità dello studente e dell'insegnante ($\sqrt{p_i q_i}$).
• Effetto: Se l'insegnante assegna un'alta probabilità a un token ma lo studente assegna una probabilità vicina allo zero, il termine di sovrapposizione svanisce e la spinta del gradiente viene attenuata. Al contrario, le preferenze condivise vengono trasferite con forza. Ciò assicura che la supervisione si rafforzi gradualmente man mano che lo studente e l'insegnante entrano in accordo, piuttosto che forzare lo studente ad adottare immediatamente le continuazioni sicure ma non supportate dell'insegnante.

B. Controllo del Drift Globale (Termine Prossimale di Fisher–Rao)

Sebbene la ponderazione della sovrapposizione limiti la spinta nei singoli stati, essa non impedisce l'accumulo di piccoli aggiornamenti ragionevoli in un grande drift globale durante l'addestramento.

• Meccanismo: GEOSD aggiunge un termine prossimale che penalizza la distanza tra la distribuzione corrente dello studente e un checkpoint recente ($\theta_{ckpt}$). Questa distanza è misurata utilizzando la distanza di Fisher–Rao (la lunghezza dell'arco sull'ipersfera), che traccia lo spostamento accumulato nella geometria delle distribuzioni predittive.
• Effetto: Questo agisce come una "regione di fiducia mobile", smorzando gli spostamenti bruschi pur permettendo un movimento costante quando il segnale dell'insegnante è coerente. A differenza dell'ancoraggio al modello iniziale (che impedirebbe l'adattamento), l'ancoraggio a un checkpoint recente permette al modello di adattarsi pur impedendogli di allontanarsi troppo dal proprio comportamento recente e stabile.

C. Aggiornamento del Gradiente Naturale

Sia la perdita di Hellinger che il termine prossimale di Fisher–Rao sono distanze nella stessa geometria delle distribuzioni del prossimo token. Per ottimizzare questo obiettivo, GEOSD utilizza un aggiornamento del gradiente naturale.

• Meccanismo: Invece di aggiornare i parametri nello spazio euclideo, l'aggiornamento è precondizionato dalla Matrice di Informazione di Fisher (approssimata tramite K-FAC per efficienza).
• Effetto: Ciò garantisce che ogni passo produca il minimo cambiamento nelle predizioni dello studente che migliori l'obiettivo, piuttosto che compiere un passo di dimensione fissa nello spazio dei parametri. Il termine prossimale agisce efficacementamente come un "richiamo al checkpoint" in questa geometria, riportando la policy verso il checkpoint in base alla salienza comportamentale delle direzioni dei parametri.

3. Contributi Chiave

1. Obiettivo GEOSD: Un framework di auto-distillazione che pesa le preferenze dell'insegnante tramite la sovrapposizione insegnante-studente (tramite perdita di Hellinger) e penalizza il drift predittivo da un checkpoint recente (tramite distanza di Fisher–Rao), ottimizzato tramite gradienti naturali in una geometria unificata.
2. Miglioramento Empirico: GEOSD preserva i guadagni in-distribution (ID) dell'auto-distillazione migliorando significativamente le prestazioni di ragionamento OOD. Su tre famiglie di modelli (Qwen3, Olmo, DeepSeek-R1) e cinque scale (da 1.7B a 32B), migliora la precisione media OOD di 5.7–8.6 punti rispetto al modello base, mentre gli obiettivi KL standard degradano le prestazioni OOD.
3. Analisi Meccanicistica: Il paper caratterizza perché il matching standard fallisce l'OOD: esso concentra la massa di probabilità sulla scelta principale dell'insegnante "drenando" le alternative in stati ad alta entropia, portando a un accordo fiducioso su risposte errate. GEOSD mantiene la dominanza della scelta principale mantenendo al contempo le continuazioni alternative "a portata di mano", prevenendo il collasso prematuro.

4. Risultati Sperimentali

Gli autori hanno valutato GEOSD su benchmark di ragionamento matematico (DAPO-Math per l'addestramento; AIME 2024/2025, AMC 2023, MATH-500 per i test OOD).

• Performance: GEOSD ha ottenuto la migliore precisione OOD in tutti i benchmark e le famiglie di modelli testati. Ad esempio, su Qwen3-8B, ha migliorato la precisione OOD di +8.6 punti rispetto al modello base, rispetto a un degrado di -8.1 punti per la Forward KL.
• Robustezza: I guadagni si sono mantenuti costanti attraverso diverse quantità di contesto privilegiato (da brevi prefissi a soluzioni complete) e attraverso diverse scale di modelli (da 1.7B a 32B).
• Ablazione: Rimuendo la ponderazione della sovrapposizione di Hellinger, il termine prossimale di Fisher–Rao o l'aggiornamento del gradiente naturale, si sono verificati cali significativi nelle prestazioni OOD, confermando che tutti e tre i componenti sono necessari.
• Falso Consenso: L'analisi ha mostrato che i metodi standard (come la Forward KL) esibiscono alti tassi di "falso consenso" (campioni che concordano fortemente su risposte errate) in stati ad alta entropia. GEOSD ha ridotto questo tasso al di sotto anche del modello base, mantenendo una coerenza interna informativa.

5. Significato e Rivendicazioni

Il paper sostiene che il degrado del ragionamento OOD nella auto-distillazione con contesto privilegiato non è semplicemente il risultato dell'uso di segnali inaffidabili, ma è una conseguenza di come tali segnali muovono il modello. Gli obiettivi standard permettono alla fiducia dell'insegnante di forzare lo studente in un modo stretto prematuramente, drenando la massa di probabilità delle alternative percorribili.

GEOSD dimostra che una distillazione del ragionamento efficace dipende non solo da quali segnali dell'insegnante fidarsi, ma da quanto lontano tali segnali sono autorizzati a spostare lo studente. Regolando questo movimento geometricamente — attenuando le spinte dove lo studente manca di supporto e ancorando la traiettoria a uno stato stabile recente — GEOSD raggiunge un compromesso superiore tra prestazioni in-distribution e generalizzazione fuori distribuzione. Gli autori concludono che il controllo del drift predittivo è essenziale per scalare l'auto-distillazione nei compiti di ragionamento.

Limitazioni Notevoli: Gli autori riconoscono che le loro scoperte sono attualmente stabilite solo per il ragionamento matematico. L'applicabilità della ponderazione della sovrapposizione e del controllo del drift di Fisher–Rao ad altri domini (ad esempio, codice, ragionamento scientifico) dove gli stati ad alta entropia possono avere strutture diverse rimane non testata. Inoltre, il metodo assume che la policy di base sia ampiamente competente; per modelli di base molto deboli, il meccanismo potrebbe sopprimere segnali utili che richiedono grandi cambiamenti comportamentali.