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

Titel

Self-Distillation als Prestatieherstelmechanisme voor LLM's: Het Tegengaan van Compressie en Catastrofaal Vergeten

1. Het Probleem

Grote Taalmodellen (LLM's) lijden vaak aan prestatieachteruitgang door twee hoofdprocessen die noodzakelijk zijn voor toepassing en efficiëntie:

• Catastrofaal Vergeten (Catastrophic Forgetting): Tijdens meervoudige rondes van Supervised Fine-Tuning (SFT) voor specifieke taken, vergeten modellen hun oorspronkelijke algemene kennis en vaardigheden.
• Compressie: Technieken zoals pruning (het verwijderen van neuronen) en quantization (het verlagen van de precisie van gewichten, bijv. naar 4-bit) leiden tot een verslechtering van de interne parameterverdeling, wat resulteert in lagere nauwkeurigheid en logische inconsistentie.

Traditionele oplossingen, zoals het opnieuw trainen van scratch of het gebruik van externe "leraren" (grotere modellen), zijn vaak computationally prohibitief, vereisen historische data die niet beschikbaar is, of introduceren distributieschiften. Er is een behoefte aan een efficiënt mechanisme om prestaties te herstellen zonder externe afhankelijkheden.

2. Methodologie

De auteurs introduceren een Self-Distillation Fine-Tuning (SDFT) herstelframework. In plaats van een extern model als leraar te gebruiken, gebruikt dit framework de eigen historische staten van het model als leraar.

• Het Framework:
  • Leraar ($\theta$): Een eerdere, hoogpresterende staat van het model (bijv. een checkpoint vóór vergeten of vóór compressie).
  • Student ($\theta_1$): Het gedegradeerde model (na SFT of compressie).
  • Proces: De student wordt getraind om de outputverdeling van de leraar na te bootsen, terwijl het zich aanpast aan de huidige taakdata. Dit fungeert als een "anker" dat de parameters terugtrekt naar de oorspronkelijke, optimale manifold.
• Uitbreiding voor Kleine Modellen: Voor kleinere modellen (bijv. 3B parameters) is de In-Context Learning (ICL) vaak te zwak voor directe self-distillation. De auteurs stellen een tweestapsproces voor:
  1. Off-policy distillatie: Een groot model (7B) bootst de ICL-vaardigheden van het kleine model op.
  2. On-policy SDFT: Het kleine model herstelt vervolgens zijn algemene vaardigheden via self-distillation.
• Theoretische Onderbouwing (Manifold Alignment):
  De auteurs postuleren dat de generatieve capaciteit van een LLM berust op een hoogdimensionale manifold die wordt geconstrueerd door de verborgen lagen. Herstel is dus niet alleen het optimaliseren van outputkansen, maar het opnieuw aligneren van de ruimtelijke structuur van de verborgen staten.
  • Meting: Ze gebruiken Centered Kernel Alignment (CKA) om de alignering tussen de activatietrajecten van de student en de leraar te kwantificeren. CKA is ideaal omdat het invariant is voor orthogonale transformaties en schaling, in tegenstelling tot metrics zoals MSE.

3. Belangrijkste Bijdragen

1. Praktisch Herstelframework: Een unified framework dat effectief prestatieverlies herstelt veroorzaakt door catastrofaal vergeten (na SFT) en compressie (pruning/quantization), zonder externe leraren of datasets.
2. Theoretisch Inzicht: Een rigoureuze theoretische uitleg dat self-distillation werkt door de hoogdimensionale manifold van de student te aligneren met die van de leraar.
3. CKA als Diagnostisch Instrument: Het aantonen dat CKA een sterke correlatie heeft met prestatieherstel en kan dienen als een voorspellende metric voor de ernst van vergeten of degradatie.
4. Extensie naar Kleine Modellen: Een strategie om het framework ook toepasbaar te maken op resource-beperkte scenario's (kleine modellen) via een hybride off/on-policy aanpak.

4. Resultaten

De experimenten zijn uitgevoerd op Qwen2.5-modellen (3B en 7B) in domeinen zoals Tool-use en Wetenschappelijke QA.

• Herstel van Catastrofaal Vergeten:
  • Na SFT op een nieuwe taak (Tool-use) daalde de prestatie op de oorspronkelijke taak (Science) van 59,96% naar 37,87%.
  • SDFT herstelde dit naar 61,54% (een winst van +23,67%), terwijl de prestatie op de nieuwe taak behouden bleef.
• Herstel na Compressie:
  • Quantization (NF4 4-bit): SDFT herstelde niet alleen de verloren prestaties, maar verbeterde ze zelfs boven het oorspronkelijke bf16-niveau (tot +22% winst op specifieke taken).
  • Pruning (10% FFN verwijdering): Ondanks de fysieke verwijdering van capaciteit, herstelde SDFT 64% van de degradatie op algemene benchmarks (MMLU) en toonde positieve cross-domein transfer.
• Kleine Modellen:
  • De tweestapsaanpak (off-policy bootstrapping + on-policy SDFT) resulteerde in een +16,49% verbetering op Tool-use voor het 3B-model ten opzichte van directe SDFT, terwijl algemene vaardigheden behouden bleven.
• Manifold Alignering (CKA):
  • Er werd een sterke correlatie gevonden: hoe groter de CKA-misalignering met de leraar, hoe ernstiger het prestatieverlies.
  • Na SDFT herstel nam de CKA-score significant toe (dichterbij 1,0), wat bevestigt dat de interne geometrische structuur is hersteld.

5. Betekenis en Conclusie

Dit paper biedt een paradigmaverschuiving in het begrijpen van self-distillation. Het beweert dat het succes van SDFT niet alleen voortkomt uit het imiteren van output, maar uit het geometrische herstel van de interne representatieruimte.

• Efficiëntie: Het elimineert de noodzaak voor dure externe leraren of het opnieuw trainen van modellen.
• Robuustheid: Het biedt een oplossing voor het "capability trade-off" probleem, waarbij specialisatie en generalisatie vaak tegenover elkaar staan.
• Toekomst: De correlatie tussen CKA en prestaties opent de deur voor nieuwe methoden om training te sturen (bijv. early stopping op basis van manifold alignering in plaats van alleen loss-waarden).

Samenvattend biedt dit werk zowel een praktische tool voor het repareren van beschadigde LLM's als een dieper theoretisch inzicht in hoe kennis is opgeslagen in de hoogdimensionale manifolds van neurale netwerken.