Attention to Mamba: A Recipe for Cross-Architecture Distillation

Probleemstelling

De afgelopen jaren zijn State Space Models (SSMs), zoals Mamba, populair geworden als alternatief voor Transformers. SSMs bieden aanzienlijke voordelen in termen van geheugenefficiëntie en doorvoersnelheid bij generatie, omdat hun complexiteit lineair is met de sequentielengte ($L$), in tegenstelling tot de kwadratische complexiteit van de klassieke Self-Attention-mechanismen in Transformers.

Er is echter een groot probleem bij de adoptie van SSMs:

1. Prestatiekloof: Hoewel SSMs efficiënter zijn, presteren ze vaak slechter dan geavanceerde, voorgetrainde Transformer-modellen op downstream-taken, vooral bij schaalvergroting.
2. Distillatie-falen: Het proberen om kennis direct te distilleren van een Transformer (Teacher) naar een Mamba-architectuur (Student) via een naïeve aanpak is tot nu toe mislukt. De prestaties van de student vallen vaak sterk terug vergeleken met de leraar.
3. Bestaande oplossingen: Eerdere pogingen om dit op te lossen, leunden vaak op hybride modellen die zowel Attention- als SSM-blokken combineren, wat de efficiëntie van een pure SSM-architectuur ondermijnt.

Het doel van dit werk is een effectieve "recept" te vinden om een Attention-based model te distilleren naar een pure Mamba-architectuur (zonder enige Attention-blokken), waarbij de oorspronkelijke prestaties van de leraar behouden blijven.

Methodologie: Een Tweepasbenadering

De auteurs stellen een tweepas-benadering voor die gebaseerd is op wiskundige principes om de kloof tussen de twee architecturen te overbruggen. In plaats van kennis direct over te dragen tussen fundamenteel verschillende paradigmen, gebruiken ze een tussenstap.

Fase 1: Van Softmax Attention naar Lineaire Attention (Hedgehog)

• Doel: De niet-lineaire Softmax Attention van de Transformer vervangen door een lineaire variant die de werking goed benadert.
• Techniek: Ze maken gebruik van de Hedgehog-methode (Zhang et al., 2024). Volgens de stelling van Mercer kan de exponentiële operator in Softmax worden herschreven als een inproduct van kenmerkvectoren (feature vectors).
• Implementatie: In plaats van een vaste lineaire benadering te gebruiken, leren ze een leerbare feature map ($\phi$) via een enkel-laags MLP. Deze map benadert de actie van Softmax door de output van de Teacher-Attention-blokken te matchen met die van een lineaire versie, gemeten via cosine-similariteit.
• Resultaat: Een "Hedgehog"-laag die lineaire Attention gebruikt maar de expressiviteit van Softmax behoudt.

Fase 2: Van Lineaire Attention naar Mamba (HedgeMamba)

• Doel: De in Fase 1 geleerde Lineaire Attention-parameters gebruiken als een geprincipeerde initialisatie voor de Mamba-architectuur.
• Theoretische link: Er is een directe wiskundige connectie tussen Lineaire Attention en SSMs. Door de parameters van de SSM ($\Lambda, B, C$) te relateren aan de Query, Key en Value matrices van de Lineaire Attention, kan de SSM worden geïnitieerd om exact hetzelfde te doen als de Lineaire Attention.
• Architectuur (HedgeMamba):
  • De SSM-parameters worden geïnitieerd op basis van de geleerde Hedgehog-features.
  • De auteurs passen de Mamba-architectuur aan om een normalisatiefactor toe te voegen (om de scores te normaliseren zoals bij Attention).
  • Ze voegen de extra expressiviteit van Mamba toe: een korte convolutie en een gate-branch (met SiLU-niet-lineariteit).
  • Deze extra componenten worden aanvankelijk geïnitieerd als de identiteitsoperator, zodat ze de output van Fase 1 niet verstoren.
• Finetuning: Het volledige model (behalve de embeddinglagen) wordt vervolgens gefinetuned op de distillatiedata (OpenWebText) met een Cross-Entropy-verlies.

Belangrijkste Bijdragen

1. Nieuwe Cross-Architectuur Distillatie: Een tweestapsmethode die eerst Attention naar Lineaire Attention distilleert en vervolgens naar Mamba, gebaseerd op functionele analogieën tussen de architecturen.
2. Principe Initialisatie: Het tonen dat het gebruik van een wiskundig onderbouwde initialisatie (via de Hedgehog-feature map) cruciaal is om de prestaties van de leraar te behouden, in plaats van te vertrouwen op naïeve distillatie.
3. HedgeMamba: Een hybride sequence mixer die de efficiëntie van Mamba combineert met de expressiviteit van een geleerde lineaire Attention, zonder Attention-blokken te behouden.
4. Uitgebreide Evaluatie: Een grondige analyse op schaal van 1B parameters met 10 miljard tokens, inclusief ablatiestudies, schaalvergelijkingen en gevoeligheidsanalyses voor de token-begroting.

Resultaten

De auteurs testen hun methode door Pythia-1B (Teacher) te distilleren naar hun HedgeMamba-architectuur (Student) met 10 miljard tokens.

• Perplexiteit (PPL): De student bereikt een PPL van 14.11, wat zeer dicht bij de Teacher (Pythia-1B) ligt met 13.86. Dit is een aanzienlijke verbetering ten opzichte van de Hedgehog-baseline (14.89) en naïeve distillatie (die faalt met PPL > 100).
• Downstream Taken: De student behaalt sterke scores op diverse benchmarks (ARC, BoolQ, Lambada, etc.), waarbij hij de Teacher op veel taken bijna evenaart.
• Ablatie Studies:
  • Het toevoegen van de gate-branch levert de grootste verbetering op in prestaties.
  • Token-begroting: De optimale verdeling van tokens tussen de twee fasen is 10% voor Fase 1 (leer de feature map) en 90% voor Fase 2 (finetuning van de volledige Mamba-architectuur). Een 50/50-verdeling (zoals in eerdere werken) werkt minder goed.
• Schaalbaarheid: De methode werkt consistent goed voor modellen van 160M, 410M en 1B parameters.

Betekenis en Impact

Dit werk is significant omdat het een praktische route biedt om de efficiëntie van SSMs (Mamba) te combineren met de superieure prestaties van bestaande Transformers.

• Efficiëntie zonder inleveren: Het stelt onderzoekers en ontwikkelaars in staat om modellen te gebruiken die sneller infereren en minder geheugen nodig hebben, zonder dat ze een model van nul moeten trainen of inleveren op kwaliteit.
• Overbrugging van theorie en praktijk: Het demonstreert dat de wiskundige connectie tussen Lineaire Attention en SSMs (zoals beschreven door Dao & Gu, 2024) effectief kan worden gebruikt voor kennisoverdracht.
• Toekomstgericht: Het biedt een blauwdruk voor het distilleren van andere Transformer-varianten naar efficiëntere architecturen, wat essentieel is voor de schaalbaarheid van LLM's in de toekomst.

Kortom, de paper bewijst dat met de juiste initialisatie en een gestructureerde tweestapsdistillatie, een pure Mamba-architectuur de prestaties van een geavanceerde Transformer kan evenaren, terwijl het de voordelen van lineaire complexiteit behoudt.