Technische Samenvatting: Multi-Granulaire Trajectoriële Alignatie (MTA) voor LLM-Distillatie

Probleemstelling

Bestaande methoden voor Kennisdistillatie (KD) voor Large Language Models (LLM's) richten zich voornamelijk op het aligneren van output-kansverdelingen (logits) of intermediate verborgen staten op vaste, statische lagen. Hoewel deze effectief zijn voor het overdragen van lokale signalen, passen deze benaderingen vaak een uniforme, token-niveau alignatiestrategie toe over alle netwerkdieptes.

Het artikel betoogt dat deze uniforme behandeling niet in overeenstemming is met de hiërarchische aard van menselijke taal en de functionele organisatie van Transformer-modellen. Taalkundig onderzoek en interpretativiteitsstudies van Transformers geven aan dat:

1. Lagere lagen oppervlakkige, lexische en feitelijke retrieval-kenmerken benadrukken.
2. Hogere lagen progressief abstracte semantische informatie en compositief redeneren coderen.

Door alle tokens geïsoleerd te behandelen en hetzelfde alignatiedoel uniform toe te passen, slagen huidige methoden er niet in om de progressieve evolutionaire trajectorie van representaties te vangen terwijl deze transformeren van lexische eenheden naar abstracte semantische structuren. Deze beperking beperkt het vermogen van het studentmodel om te generaliseren naar downstream-taken.

Methodologie: Multi-Granulaire Trajectoriële Alignatie (MTA)

Om deze beperkingen aan te pakken, stellen de auteurs MTA voor, een modulaire framework die de representatietrajectorie van docent- en studentmodellen aligneert met behulp van een laag-adaptieve strategie. MTA opereert op de hypothese dat distillatie rekening moet houden met de diepte-afhankelijke semantische rollen van Transformer-lagen.

Kerncomponenten

1. Laag-Adaptieve Multi-Granulaire Spans:

   • Lagere Lagen: Alignaties worden uitgevoerd op woord-niveau (Woordspans) om lexische verankering en fijnmazige afhankelijkheden te behouden.
   • Hogere Lagen: Alignaties verschuiven naar frase-niveau (Zelfstandig Naamwoordfrases en Werkwoordfrases) om compositieve semantiek en hogere-orde structurele relaties te vangen.
   • Spanconstructie: Semantische eenheden worden geëxtraheerd met behulp van een syntactische parser (bijv. spaCy). Token-impactgewichten worden berekend via een bidirectioneel zelf-attentie-mechanisme (exclusief zelflussen) om token-representaties te aggregeren tot span-representaties.

2. Dynamische Structurele Alignatieverlies ($L_{DSA}$):

   • Dit doel dwingt geometrische consistentie af tussen de docent en de student.
   • In plaats van individuele tokenvectoren te matchen, minimaliseert $L_{DSA}$ de discrepantie in de paarsgewijze relationele afstanden tussen semantische spans (woorden of zinnen) binnen een specifieke laag.
   • Het prioriteert semantisch opvallende spans door de alignatie te wegen op basis van de geaggregeerde impact van constitutieve tokens.
   • Wiskundig aligneert het de topologische structuur van de representatieruimte:
     $$L_{DSA} = \frac{1}{\|L_{key}\|} \sum_{l \in L_{key}} \sum_{i,j} w_{ij} \left( d(U^S_{i,l}, U^S_{j,l}) - d(U^T_{i,\phi(l)}, U^T_{j,\phi(l)}) \right)^2$$
     waarbij $d(\cdot)$ de cosinusafstand is, en $U$ de gewogen span-representatie voorstelt.

3. Alignatieverlies voor Verborgen Representaties ($L_{Hid}$):

   • Om structurele beperkingen aan te vullen, aligneert een feature-level verlies specifieke verborgen staten.
   • Aangezien de dimensies van de student doorgaans kleiner zijn dan die van de docent, maakt een leerbaar lineair projectiemiddel studentstaten af op de ruimte van de docent.
   • Het verlies minimaliseert de gewogen cosinusafstand tussen geprojecteerde studentstaten en docentstaten, met focus op de meest informatieve tokens.

4. Optimalisatie:

   • MTA is ontworpen als een plug-in module. Het totale verlies combineert het basisdistillatieverlies ($L_{Base}$) van bestaande frameworks (bijv. DistiLLM, FDD) met de voorgestelde structurele en verborgen verliezen:
     $$L_{Total} = L_{Base} + \lambda_{DSA}L_{DSA} + \lambda_{Hid}L_{Hid}$$

Belangrijkste Bijdragen

1. Identificatie van Beperkingen: De auteurs identificeren dat uniforme intermediate alignatie de hiërarchische evolutie van representaties negeert, wat leidt tot suboptimale kennisoverdracht.
2. Laag-Adaptief Framework: Zij introduceren MTA, dat dynamisch overgaat van woord-niveau alignatie in lagere lagen naar frase-niveau structurele alignatie in hogere lagen, wat de bottom-up compositieve aard van taal weerspiegelt.
3. Gegenereerd Module: MTA wordt aangetoond als een gegeneraliseerde optimalisatielaag die compatibel is met state-of-the-art distillatiemethoden (DistiLLM, DistiLLM-2, FDD) die tokenisatie delen.
4. Empirische Validatie: Uitgebreide experimenten tonen aan dat MTA consistent de prestaties verbetert over diverse modelarchitecturen (GPT-2, Qwen1.5, OPT) en benchmarks.

Experimentele Resultaten

De auteurs evalueerden MTA op instructievolgende taken met drie docent-studentparen:

• GPT-2 1.5B $\to$ 120M
• Qwen1.5 1.8B $\to$ 0.5B
• OPT 6.7B $\to$ 1.3B

Prestatiemetingen:

• ROUGE-L: MTA presteerde consistent beter dan baselines (FDD, DistiLLM, DistiLLM-2) over alle datasets (Dolly, SelfInst, Vicuna, S-NI). Bijvoorbeeld, in de GPT-2 1.5B $\to$ 120M instelling, verbeterde het toevoegen van MTA aan DistiLLM de gemiddelde ROUGE-L van 20,21 naar 21,45.
• LLM-as-a-Judge: Met behulp van GPT-4o-mini om feitelijke nauwkeurigheid, coherentie en nuttigheid te evalueren, vertoonden MTA-versterkte modellen consistente scoreverbeteringen ten opzichte van baselines.

Ablatiestudies:

• Verliescomponenten: Zowel $L_{DSA}$ (structureel) als $L_{Hid}$ (feature) droegen positief bij, waarbij hun combinatie de beste resultaten opleverde.
• Granulariteit: De "Full-level" (laag-adaptieve) strategie presteerde beter dan statische strategieën (alleen Woord of alleen Frase), wat bevestigt dat het aligneren van lexische verankering in lagere lagen en compositieve semantiek in hogere lagen cruciaal is.
• Lageselectie: De methode is robuust voor verschillende laagplanningen, hoewel een hybride toewijzing (1 Woord-niveau laag : 2 Frase-niveau lagen) als optimaal werd bevonden voor de geteste configuraties.

Efficiëntie:

• MTA introduceert trainings overhead voornamelijk door syntactische span-extractie (verminderd tot ~0,27s/stap via optimalisatie) en geometrische matching.
• Inferentiekosten blijven onveranderd aangezien span-extractie niet tijdens de testtijd wordt gebruikt.
• Tijd-gematchde vergelijkingen bevestigden dat prestatiewinsten te wijten zijn aan het alignatiedoel, niet aan een verhoogd trainingsbudget.

Betekenis en Claims

Het artikel claimt dat MTA een semantisch getrouwe benadering biedt voor kennisdistillatie door expliciet de hiërarchische representatietrajectorie van LLM's te modelleren. Door verder te gaan dan statische, token-niveau matching naar dynamische, multi-granulaire structurele alignatie, stelt MTA studentmodellen in staat om de interne informatieverwerkingsdynamiek van de docent beter te repliceren.

De auteurs positioneren MTA niet als een vervanging voor bestaande distillatieframeworks, maar als een complementaire, plug-and-play module die de generalisatie en robuustheid van gecomprimeerde LLM's verbetert. Het werk onderstreept het belang van het respecteren van taalkundige hiërarchie en diepte-afhankelijke feature-evolutie in het ontwerp van toekomstige modelcompressietechnieken.