Self-Distillation as a Performance Recovery Mechanism for LLMs: Counteracting Compression and Catastrophic Forgetting

1. Il Problema

I Large Language Models (LLM) subiscono spesso un significativo degrado delle prestazioni in scenari reali a causa di due fattori principali:

• Dimenticanza Catastrofica (Catastrophic Forgetting): Durante il Supervised Fine-Tuning (SFT) multi-round per adattarsi a nuovi domini, i modelli tendono a perdere le conoscenze generali e le abilità apprese in precedenza.
• Compressione: Tecniche come la pruning (potatura) e la quantizzazione (es. NF4) sono essenziali per il deployment efficiente, ma alterano le distribuzioni interne dei parametri, portando a cali di accuratezza e coerenza logica.

Le soluzioni tradizionali per riparare questi modelli (come il retraining da zero o l'uso di dati storici) sono spesso computazionalmente proibitive o richiedono l'accesso a dati non disponibili. Inoltre, la distillazione della conoscenza (KD) classica richiede un "insegnante" esterno potente, che può introdurre spostamenti di distribuzione o costi elevati.

2. Metodologia: Il Framework di Recupero (SDFT)

Gli autori propongono un framework di recupero basato sulla Self-Distillation Fine-Tuning (SDFT). L'idea centrale è utilizzare lo stato storico del modello stesso come "insegnante" per recuperare le capacità perse, senza dipendere da modelli esterni.

Il framework si articola in tre scenari principali:

1. Recupero dalla Dimenticanza Catastrofica: Il modello degradato (studente) viene addestrato per imitare la distribuzione di output di una sua versione precedente (insegnante, es. un checkpoint prima dell'SFT aggressivo) utilizzando i dati del dominio originale. Questo agisce come un "ancoraggio" che riporta i parametri verso il manifold ad alte prestazioni originale.
2. Recupero dalla Compressione: Dopo aver applicato quantizzazione o pruning, il modello compresso viene recuperato tramite SDFT utilizzando il modello originale non compresso come insegnante statico.
3. Adattamento per Modelli Piccoli (Small-Scale): Per modelli di piccole dimensioni (es. 3B) dove l'Apprendimento Contestuale (ICL) è debole, viene proposta una strategia in due fasi:
   • Fase 1 (Off-policy): Un modello grande (es. 7B) funge da insegnante per attivare le capacità ICL del modello piccolo.
   • Fase 2 (On-policy): Il modello piccolo, ora con ICL attivato, subisce un processo di recupero SDFT per ripristinare le capacità generali perse durante la fase 1.

3. Contributi Teorici: Allineamento del Manifold

Un contributo fondamentale del paper è la spiegazione teorica del perché la SDFT funziona, spostando l'analisi dalle distribuzioni di output alla geometria interna.

• Ipotesi: La capacità generativa di un LLM dipende fondamentalmente dal manifold ad alta dimensione costruito dai suoi strati nascosti. Il degrado delle prestazioni corrisponde a una deviazione di questo manifold rispetto alla struttura ottimale.
• Metrica: Gli autori utilizzano l'Allineamento del Kernel Centrato (CKA - Centered Kernel Alignment) per quantificare l'allineamento tra le traiettorie di attivazione dello studente e dell'insegnante. Il CKA è scelto per la sua invarianza rispetto a trasformazioni ortogonali e scalature, rendendolo superiore alla semplice MSE per confrontare rappresentazioni neurali.
• Teorema: La SDFT non ottimizza solo le probabilità di output, ma regolarizza la struttura spaziale degli stati nascosti, allineando il manifold dello studente a quello ottimale dell'insegnante.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su modelli Qwen2.5-3B-Instruct e Qwen2.5-7B-Instruct su domini di Tool-use e Science.

• Recupero dalla Dimenticanza: In uno scenario a tre stadi (Esperto Science -> SFT Tool-use -> Recupero SDFT), il modello ha recuperato l'accuratezza su Science da 37.87% a 61.54% (recuperando quasi completamente la perdita), mantenendo al contempo alte prestazioni su Tool-use.
• Recupero dalla Compressione:
  • Quantizzazione (NF4): La SDFT non solo ha recuperato le perdite, ma ha portato le prestazioni del modello quantizzato a superare il modello originale in bf16 (guadagni di +15-22% sui task specifici).
  • Pruning (10% FFN): Nonostante la rimozione fisica di neuroni, la SDFT ha recuperato il 64% del degrado su MMLU e ha mostrato un trasferimento positivo cross-dominio (migliorando il Tool-use senza dati di training specifici).
• Modelli Piccoli: La strategia a due stadi ha permesso al modello 3B di raggiungere il 67.01% su Tool-use, un miglioramento del +16.49% rispetto alla SDFT diretta, preservando le capacità generali (MMLU) entro l'1.5% dal modello base.
• Validazione Teorica (CKA): È stata osservata una forte correlazione tra il recupero delle prestazioni e l'aumento del punteggio CKA.
  • Quando il modello dimentica, il CKA rispetto all'insegnante originale diminuisce.
  • Dopo il recupero SDFT, il CKA aumenta, indicando un riallineamento del manifold.
  • Il CKA si è dimostrato un predittore quantitativo affidabile della gravità della dimenticanza.

5. Significato e Implicazioni

Questo lavoro colma il divario tra framework pratici di recupero e la teoria della rappresentazione geometrica.

• Meccanismo Interno: Dimostra che la distillazione autonoma funziona agendo come un meccanismo di riallineamento geometrico del manifold, non solo come correzione comportamentale.
• Efficienza: Offre una soluzione efficiente e priva di dipendenze esterne per riparare modelli degradati, cruciale in contesti con risorse computazionali limitate o vincoli di privacy.
• Nuova Metrica: Introduce il CKA come strumento diagnostico per monitorare la salute del modello e la gravità della dimenticanza durante il training, potenzialmente utile per l'early stopping o il tuning degli iperparametri.
• Generalità: Il framework si applica universalmente a diversi tipi di degrado (SFT, pruning, quantizzazione), suggerendo che la struttura geometrica interna è un fattore chiave per la robustezza degli LLM.

In sintesi, il paper stabilisce che preservare o recuperare la struttura geometrica interna (il manifold) è la chiave per mantenere le capacità degli LLM, e la Self-Distillation è lo strumento efficace per raggiungere questo obiettivo.