Dual-Space Knowledge Distillation with Key-Query Matching for Large Language Models with Vocabulary Mismatch

Probleemstelling

Grote Taalmodellen (LLMs) presteren uitstekend maar zijn kostbaar om te deployen vanwege hun grootte en rekenvereisten. Kennisdistillatie (Knowledge Distillation, KD) lost dit op door kleinere 'Student'-modellen te trainen om grotere 'Teacher'-modellen na te bootsen. Een groot obstakel in de huidige LLM-landschap is echter dat verschillende modellen vaak verschillende tokenizers en vocabulaires gebruiken. Dit zorgt voor een mismatch in de lengte van sequenties en de predictieruimte, wat het vergelijken en aligneren van hun output bemoeilijkt.

Bestaande methoden zoals DSKD-CMA (Dual-Space Knowledge Distillation with Cross-Model Attention) zijn state-of-the-art voor deze cross-tokenizer KD, maar hun interne werking is een "black box". Het is onduidelijk hoe ze token-alignment precies leren en of er manieren zijn om de prestaties, vooral op data buiten de trainingsdistributie (OOD), te verbeteren.

Methodologie

De auteurs hanteren een tweeledige aanpak: eerst een systematische analyse van de bestaande DSKD-CMA-methode, en daarna de introductie van een verbeterde methode.

1. Analyse van DSKD-CMA

De auteurs onderzochten het Cross-Model Attention (CMA) mechanisme door:

• Manuele Token-Alignering: Ze ontwikkelden een probe-methode om tokens van Teacher en Student expliciet te aligneren op basis van "chunks" (tekstsegmenten die door beide tokenizers worden gegenereerd).
• Visualisatie: Ze gebruikten heatmaps om te zien hoe CMA tokens koppelt.
• Ablatiestudies: Ze vergelijken de originele CMA met twee varianten:
  • Chunk-Level Projection (CLP): Vervangt de complexe attention-mechanisme door een uniforme projectie binnen gealigneerde chunks.
  • Chunk-Level Attention (CLA): Beperkt de attention-mechanisme strikt tot binnen de gealigneerde chunks.

Resultaat van de analyse: CMA leert impliciet de juiste chunk-structuur, maar is soms "ruisig" bij herhalende tokens. CLP presteert vergelijkbaar of iets beter, maar is rekenkundig duurder. CLA degradeerde de prestaties, wat suggereert dat DSKD baat heeft bij stabiele, grove aligneringen in plaats van te gevoelige, lokale relaties.

2. Nieuwe Methode: DSKD-CMA-GA

De kern van de innovatie is het herkennen dat de Keys (K) en Queries (Q) in de attention-mechanisme uit verschillende modellen komen en dus verschillende verdelingen hebben ($p_K$ en $p_Q$). Deze mismatch verzwakt de attention.

Om dit op te lossen, introduceren de auteurs Key-Query Matching via twee technieken:

• Generative Adversarial (GA) Alignment: Een discriminator probeert te onderscheiden of een vector afkomstig is van de Teacher (Key) of de Student (Query). De Student-projector (generator) wordt getraind om de discriminator te misleiden, waardoor de verdelingen van Keys en Queries worden uitgelijnd.
• Conditional Transport (CT): Een alternatieve aanpak gebaseerd op Optimal Transport om de verdelingen te aligneren.

De totale loss functie combineert de originele distillatie-verlies met een extra term voor Key-Query matching ($L_{KQ}$).

Belangrijkste Bijdragen

1. Analyserend Framework: Een systematische analyse van DSKD-CMA die aantoont dat het mechanisme chunk-georiënteerde alignering leert, maar beperkingen heeft in lokale precisie.
2. DSKD-CMA-GA en DSKD-CMA-CT: Nieuwe frameworks die Key-Query matching integreren om de verdelingsmismatch tussen verschillende tokenizers op te lossen.
3. Prestatieverbetering: Het aantonen dat deze methoden de kloof tussen cross-tokenizer en same-tokenizer distillatie verkleinen, en in sommige gevallen zelfs same-tokenizer prestaties overtreffen.

Resultaten

De experimenten zijn uitgevoerd met GPT-2 (Student) en GPT-2 XL/Qwen1.5 (Teacher) op diverse datasets (Dolly, Self-Instruct, Vicuna, etc.).

• ROUGE-L Scores: DSKD-CMA-GA behaalde consistente verbeteringen ten opzichte van de originele DSKD-CMA.
  • Op Out-of-Distribution (OOD) data was de gemiddelde verbetering +0.37 punten in ROUGE-L.
  • De methode verkleinde de prestatiekloof tussen cross-tokenizer en same-tokenizer KD aanzienlijk.
• Vergelijking Methoden: De GA-variant (Generative Adversarial) presteerde significant beter dan de CT-variant (Conditional Transport). Dit suggereert dat adversarial learning effectiever is in het aligneren van verdelingen bij tokenizer-mismatch.
• Divergentie Functies: De keuze van de divergentie-functie (bijv. Skewed KL) is cruciaal en moet worden afgestemd op de gebruikte aligneringsmethode voor optimale resultaten.
• Efficiëntie: Hoewel CLP (Chunk-Level Projection) goed presteerde, is DSKD-CMA-GA praktischer voor grootschalig training omdat het geen expliciete, dure alignering in elke stap vereist.

Betekenis en Conclusie

Dit paper biedt een belangrijke stap voorwaarts in het veld van cross-tokenizer kennisdistillatie. Door de interne werking van bestaande state-of-the-art methoden te ontrafelen en een principieel probleem (verdelingsmismatch tussen Keys en Queries) aan te pakken, tonen de auteurs aan dat cross-tokenizer KD niet langer een tweede keus hoeft te zijn.

De voorgestelde DSKD-CMA-GA methode maakt het mogelijk om kleinere modellen te trainen die bijna net zo goed presteren als modellen met dezelfde tokenizer, zelfs bij extreme vocabulaireverschillen (zoals tussen GPT-2 en Qwen). Dit opent de deur voor efficiëntere deployments van LLMs in diverse domeinen zonder dat men gebonden is aan één specifiek tokenizer-systeem. Toekomstig werk richt zich op het testen van deze methoden op complexere redeneertaken en het dynamisch wegen van de verliesfuncties.