Technische Samenvatting: N-vium – Mixture-of-Exits Transformer voor Versnelde Exacte Generatie

1. Probleemstelling

Standaard autoregressieve Large Language Models (LLMs) verwerken tokens sequentieel door een stapel decoderlagen. Deze forward pass dient twee gekoppelde functies: (1) het bouwen van Key-Value (KV) representaties op elke laag voor toekomstige attention, en (2) het projecteren van de laatste verborgen toestand om het volgende token te voorspellen. Bijgevolg vereist elk token een volledige forward pass door alle lagen, ongeacht of de voorspellingsverdeling al is geconvergeerd.

Bestaande benaderingen om de inferentie-efficiëntie te verbeteren, zoals Early Exit-methodes, proberen deze functies te ontkoppelen door voorspellingskoppen aan te brengen op tussenliggende lagen. Deze methodes behandelen echter doorgaans tussenliggende voorspellingen als benaderingen van de verdeling van de laatste laag. Dit leidt tot verschillende nadelen:

• Benaderingsfouten: Tussenliggende voorspellingen zijn vaak zwakker dan die met volledige diepte, wat een train-test mismatch creëert.
• Heuristische Routing: Beslissingen over wanneer te exiten, berusten vaak op handgemaakte betrouwbaarheidsdrempels die kalibratie vereisen.
• Onvolledige KV Caches: Om rekenkracht te besparen, overslaan early-exit methodes vaak de berekening van KV caches voor overgeslagen lagen, wat vereist dat deze bij benadering worden gekopieerd of overgeslagen, wat de kwaliteit verslechtert.
• Kwaliteitsdegradatie: Snelheid wordt vaak bereikt ten koste van perplexiteit en generatiekwaliteit.

Het artikel betoogt dat, hoewel empirische observaties (bijvoorbeeld via de "logit lens") tonen dat voorspellingsverdelingen vaak stabiliseren voordat de laatste laag wordt bereikt, huidige methodes dit niet benutten voor exacte generatie zonder kwaliteitsverlies.

2. Methodologie: N-vium

De auteurs introduceren N-vium, een mixture-of-exits framework dat de standaard transformer-architectuur strikt generaliseert. In plaats van vroege exits te behandelen als benaderingen, herdefinieert N-vium de output van het model als een geleerde mixverdeling over meerdere exits.

Architectuur

• Mixture of Exits: Het model partitioneert de $L$ lagen van de backbone in $N_V$ blokken. Aan het einde van elk blok (knooppunt $k$) bepaalt een router de waarschijnlijkheid om het volgende token te bemonsteren uit de verdeling $\pi_k$ van die laag.
• Token-Adaptieve Routing: Een router-MLP verwerkt de verborgen toestand $h^{(\ell_k)}_t$ op knooppunt $k$ om een waarschijnlijkheid $w_k$ te outputten. De waarschijnlijkheid $p_k$ om te exiten bij knooppunt $k$ wordt berekend als:
  $$p_k(x_t) = w_k(h^{(\ell_k)}_t) \prod_{j=1}^{k-1} (1 - w_j(h^{(\ell_j)}_t))$$
  waarbij $w_{N_V} = 1$ bij conventie.
• Mixverdeling: De uiteindelijke verdeling voor het volgende token is een gewogen som van alle exit-verdelingen:
  $$\pi_{mix}(x_{t+1} | x_{1:t}) = \sum_{k=1}^{N_V} p_k(x_t) \pi_k(x_{t+1} | x_{1:t})$$
• Gedeelde Kop: Alle exits delen de uiteindelijke Language Modeling (LM) projectiekop $W_{lm}$, met lichtgewicht adapter-MLP's die tussenliggende verborgen toestanden transformeren om te matchen met de verwachte geometrie van de kop.

Training

• Doel: Het model wordt getraind om de cross-entropy loss van de gezamenlijke mixverdeling $\pi_{mix}$ te minimaliseren, in plaats van individuele koppen geïsoleerd te optimaliseren om de laatste laag te matchen.
• Rekenstraf: Een hulploss-term $\mathcal{L}_{compute}$ wordt toegevoegd om de router aan te moedigen waarschijnlijkheidsmassa toe te wijzen aan eerdere exits, waardoor het model wordt gestuurd naar efficiëntie:
  $$\mathcal{L}_{total} = \mathcal{L}_{mix} + \beta \mathcal{L}_{compute}$$
  waarbij $\mathcal{L}_{compute}$ de verwachte genormaliseerde laagdiepte is.
• Warm-up: Een router warm-up-fase wordt toegepast om "router collapse" te voorkomen (waarbij alle tokens naar de laatste laag routeren) door afwijkingen van een uniforme exit-verdeling aanvankelijk te bestraffen.

Inferentie: Piggybacking

Een cruciale innovatie is de Piggybacking-techniek om exacte KV-cacheherstel te garanderen zonder berekening over te slaan:

1. Uitgestelde Berekening: Wanneer een token exit bij knooppunt $k$, wordt de berekening voor lagen $k+1$ tot $L$ uitgesteld.
2. Parallel Batching: Naarmate toekomstige tokens worden verwerkt, bereiken hun forward passes de uitgestelde lagen. Op deze diepten worden de verborgen toestanden van de uitgestelde tokens samengevoegd met de toestand van de huidige token en parallel verwerkt.
3. Resultaat: Dit vult de ontbrekende KV-toestanden voor alle uitgestelde tokens in één parallelle pass in, waardoor de KV-cache compleet, exact en consistent is met het bemonsterde token. Er wordt geen berekening overgeslagen; deze wordt slechts herschikt van sequentiële naar parallelle gebatchte uitvoering.

3. Belangrijkste Bijdragen

1. Exacte Bemonsteringsmix: N-vium is de eerste methode die token-adaptieve diepterouting mogelijk maakt met exacte bemonstering uit de getrainde verdeling, waardoor de benaderingsfouten worden vermeden die inherent zijn aan standaard early-exit methodes.
2. Volledig KV-herstel: Door berekeningen van bovenste lagen uit te stellen en te batchen, herstelt N-vium volledige en exacte KV-caches, een functie die vaak verloren gaat bij andere adaptieve berekeningsmethodes.
3. Trainingsframework: De auteurs bieden een praktische trainingspipeline die de mixdoelstelling direct optimaliseert, inclusief strategieën om modellen naar efficiëntie te sturen via een leerbare rekenstraf.
4. Schaalbaarheid: De methode is gevalideerd tot schalen van 1,5 miljard parameters, wat aantoont dat de aanpak niet beperkt is tot kleine modellen.

4. Experimentele Resultaten

De auteurs hebben N-vium-modellen (specifiek "Quadrivium" met 4 exits) voorgetraind op het C4-dataset met de LLaMA 2-architectuur, en deze vergeleken met parameter- en datagematchte dichte baselines.

• Snelheidswinst vs. Perplexiteit:
  • Het grootste model (1,5 miljard parameters, 48 lagen, 1536 verborgen breedte) behaalde een 57,9% snelheidswinst in wandkloktijd ten opzichte van de dichte baseline met geen verslechtering van de perplexiteit (0% verschil).
  • Kleinere modellen toonden consistente snelheidswinsten variërend van 23% tot 41% met lichte verbeteringen in perplexiteit (bijvoorbeeld toonde een model met 36 lagen een snelheidswinst van 36,4% met een verbetering van 3,7% in perplexiteit).
• Pareto-grens: In kwaliteits-snelheidsplots bezet N-vium het "rechts-boven" kwadrant (sneller en gelijke of betere kwaliteit), terwijl methodes zoals CALM (hergeïmplementeerd als CALM*) vallen in het "rechts-onder" (sneller maar slechtere kwaliteit) en LayerSkip ligt op een lijn met constante perplexiteit onder de baseline.
• Downstream Prestaties: Na supervised fine-tuning (SFT) op de Tulu 3-mix, matchten of overtroffen N-vium-modellen dichte baselines op downstream benchmarks (MMLU, ARC, HellaSwag, etc.), wat bevestigt dat de modellen bestuurbaar blijven en niet overfitten op de cross-entropy doelstelling ten koste van taakprestaties.
• Hardware Robuustheid: Snelheidswinstmetingen op A100, RTX A6000 en RTX 3090 GPU's leverden consistente resultaten op (~57%), terwijl oudere GPU's met Turing-architectuur lagere winsten toonden door gebrek aan native bfloat16-ondersteuning.

5. Betekenis en Claims

Het artikel stelt dat N-vium een paradigma-verschuiving in inferentie-efficiëntie vertegenwoordigt. In tegenstelling tot de meeste methodes die tokens per seconde verbeteren door FLOPs per token te verminderen (vaak via benadering), verbetert N-vium de efficiëntie door FLOPs per seconde te verhogen. Dit bereikt het door sequentiële berekening te herschikken in parallelle gebatchte passes, waardoor hardware-parallelisme beter wordt benut zonder berekening te verwerpen.

De auteurs benadrukken dat N-vium een strikt generalisatie is van de standaard transformer: wanneer de routinggewichten voor tussenliggende exits nul zijn, gedraagt het model zich exact als een standaard transformer. De methode voldoet gelijktijdig aan vier kritieke criteria, wat de auteurs beweren dat geen enkele andere bestaande methode doet:

1. Exacte Inferentie: Bemonsteren uit de getrainde verdeling zonder benadering.
2. Volledig KV-cacheherstel: Geen overgeslagen lagen of benaderende cache-invulling.
3. Geleerde Routing: Berekening per token wordt toegewezen door een getrainde router, niet door heuristieken.
4. Decode-tijd snelheidswinst: Vermindert wandkloktijd latentie tijdens autoregressieve generatie.

Het artikel concludeert dat, hoewel het trainen van N-vium een hogere FLOP-kost per token met zich meebrengt (door het evalueren van koppen op meerdere knooppunten), deze overhead gerechtvaardigd is voor productie-implementatie waar inferentiekosten domineren, en dat de relatieve overhead afneemt naarmate de modelgrootte schaalt.