Technische Samenvatting: SRA (Span Representation Alignment voor Distillatie van Grootte Taalmodellen)

Probleemstelling

Het artikel adresseert de uitdaging van Cross-Tokenizer Knowledge Distillation (CTKD), waarbij een groot "leraar"-taalmodel en een kleiner "student"-model verschillende tokenizers en vocabulaires gebruiken. Bestaande CTKD-methoden vertrouwen voornamelijk op token-niveau- of logit-niveau-afstemmingstrategieën (bijvoorbeeld met behulp van edit-afstand, optimale transport of projectielagen). De auteurs betogen dat deze benaderingen vaak broos en gevoelig zijn voor discrepanties tussen tokenizers, omdat ze proberen discrete, atomaire eenheden af te stemmen die mogelijk niet semantisch overeenkomen over verschillende vocabulaire-systemen heen. Bovendien slagen huidige methoden er vaak niet in om de onderliggende dynamiek van feature-evolutie binnen de Transformer-architectuur vast te leggen wanneer tokenizers verschillen.

Methodologie: SRA-raamwerk

Het voorgestelde SRA-raamwerk herformuleert CTKD door de fundamentele eenheid van afstemming te verschuiven van individuele tokens naar robuste, tokenizer-onafhankelijke spans. Deze aanpak is gebaseerd op de fysieke analogie van Multi-Particle Dynamical Systems (MPDS), waarbij de Transformer wordt beschouwd als een systeem van interagerende deeltjes (tokens) die evolueren door de lagen heen.

1. Spanconstructie en Mapping

• Spandefinitie: In plaats van tokens direct af te stemmen, groepeert SRA tokens in "spans" die overeenkomen met dezelfde tekstuele eenheden in de ruwe invoer.
• Afstemming via LCS: Om spans tussen leraar- en studentenmodellen te matchen, berekenen de auteurs de Longest Common Subsequence (LCS) van de karakteroffsets die door beide tokenizers worden geproduceerd. Dit identificeert overeenkomende segmenten van de invoerstekst, en definieert start- en eindindices voor afgestemde spans in beide modellen zonder dat een directe token-tot-token-mapping vereist is.

2. Spanrepresentatie via Massamiddelpunt (CoM)

Geput uit het MPDS-perspectief, wordt elke span behandeld als een cluster van deeltjes. De toestand van dit cluster wordt gerepresenteerd door zijn Massamiddelpunt (CoM):

• Gewogen Pooling: De representatie van een span wordt berekend als een door aandacht gewogen gemiddelde van de verborgen toestanden van de constitutieve tokens.
• Aandachtsgewichten: Het gewicht $w_i$ voor een token wordt afgeleid van de attentiescores die worden ontvangen van het laatste token van de sequentie over alle attentiekoppen heen. Dit nabootst het fysieke concept waarbij "zwaardere" (semantisch meer saliente) deeltjes het centrum van het cluster domineren.
• Formule: $C_i = \sum w_t H_t$, waarbij $H_t$ de verborgen toestand is en $w_t$ het door attentie afgeleide gewicht.

3. Distillatiedoelstellingen

SRA maakt gebruik van een gecombineerde verliesfunctie om kennis op span-niveau over te dragen:

• Span-niveau Verborgen Toestandsoverdracht ($L_{HS}^{Span}$):

  • Stemt de CoM-representaties van studenten- en leraars-spans af met behulp van een cosine-afstandsverlies.
  • Geometrische Regularisator ($L_{Geo}$): Om de structurele integriteit van de representatieruimte te behouden, wordt een regularisator toegevoegd om de relatieve afstanden (cosine-afstand) tussen paren spans in het studentenmodel te behouden, overeenkomend met die in het leraarmodel.
  • Gewogen Verlies: Het verlies wordt gewogen op basis van de semantische salientie van de spans (afgeleid van de som van token-gewichten binnen de span), zodat de student zich richt op de meest informatieve regio's.

• Gegelijkgerichte Span-logits Distillatie ($L_{KD}^{Span}$):

  • Adresseert vocabulaire-mismatch ($V_{tea} \neq V_{stu}$) door kennisoverdracht te beperken tot de gezamenlijke deelruimte $\tilde{V} = V_{tea} \cap V_{stu}$.
  • Span-niveau logits worden geaggregeerd uit token-logits binnen het gezamenlijke vocabulaire, en een Kullback-Leibler (KL)-divergentie wordt berekend tussen de distributies van de leraar en de student over deze gezamenlijke ruimte.

• Totale Doelstelling: Het uiteindelijke verlies combineert het standaard cross-entropy-verlies op grondwaarheid-labels met de span-niveau distillatieverliezen (verborgen toestanden en logits), gebalanceerd door een hyperparameter $\alpha$.

Belangrijkste Bijdragen

1. Nieuw Conceptueel Raamwerk: Het artikel introduceert een CTKD-raamwerk dat is gebaseerd op het Multi-Particle Dynamical System- perspectief van Transformers, waarbij spans worden gemodelleerd als deeltjesclusters en hun Massamiddelpunt wordt gebruikt voor stabiele, additieve representaties.
2. Span-niveau Distillatie: Het stelt een methode voor om kennis op span-niveau te distilleren, waarbij zowel verborgen representaties als logits worden omvat. Dit omvat een door attentie afgeleid wegingsschema om saliente spans te prioriteren en een geometrische regularisator om de relatieve geometrie van de representatieruimte te behouden.
3. Empirische Validatie: De auteurs tonen aan dat SRA significant beter presteert dan state-of-the-art CTKD-baselines (waaronder ULD, DSKD, MinED en MultilevelOT) over meerdere leraar-studentparen en benchmarks.

Experimentele Resultaten

De auteurs hebben SRA geëvalueerd op diverse cross-architectuur- en cross-tokenizerparen (bijvoorbeeld Qwen1.5/2.5, Mistral-7B, GPT-2, OPT, TinyLLaMA) met behulp van instructievolgende datasets (Dolly, Vicuna, SelfInst, S-NI, DialogSum).

• Prestatie: SRA behaalde consistent de hoogste ROUGE-L-scores en de beste gemiddelde prestatie over alle instellingen in vergelijking met sterke baselines.
• Semantische Kwaliteit: Semantische evaluaties met GPT-4o-mini als rechter bevestigden dat SRA meer semantisch trouwe en contextueel afgestemde antwoorden produceert.
• Ablatiestudies:
  • Het verwijderen van de geometrische regularisator of het gewogen pooling-mechanisme resulteerde in prestatiedegradatie, wat de noodzaak bevestigt om structurele relaties te behouden en zich te richten op saliente spans.
  • De methode toonde robuustheid zelfs in scenario's met dezelfde tokenizer, en presteerde beter dan op afstand gebaseerde en andere CTKD-methoden.
• Efficiëntie: Hoewel SRA vanwege span-matching en -aggregatie iets meer GPU-geheugen vereist dan sommige baselines, behaalde het de laagste gemiddelde trainingstijd per stap onder de vergeleken methoden.

Betekenis en Claims

Het artikel claimt dat SRA een robuste, fysiek onderbouwde aanpak biedt voor kennisdistillatie die de broosheid van token-niveau-afstemming overwint. Door de eenheid van afstemming te verschuiven naar spans en het concept van het Massamiddelpunt te benutten, ondersteunt SRA op natuurlijke wijze de overdracht van dynamische, laag-tot-laag feature-evolutie. De auteurs stellen dat hun methode het principe valideert dat distillatie zich moet richten op de dynamiek van de geometrische evolutie van het systeem in plaats van alleen op de uiteindelijke output, en biedt een schaalbare oplossing voor het implementeren van efficiënte modellen in heterogene tokenizer-omgevingen.

Het werk wordt gepresenteerd als een stap naar een stabielere en meer uitgebreide kennisoverdracht, hoewel de auteurs beperkingen erkennen, zoals het statische karakter van huidige logit-mapping en de rekenkosten van teacher-inferentie op de vlucht die vereist is voor span-afstemming.