Technische Samenvatting: Dual-Flow Transformers

Probleemstelling

De kosten voor het bedienen van Large Language Models (LLM's) worden steeds meer gedomineerd door inferentie in plaats van training. De twee primaire fasen van inferentie — prefill (verwerking van de prompt) en autoregressieve decode (genereren van tokens) — vertonen echter fundamenteel verschillende hardware-bottlenecks:

• Prefill: Parallelle verwerking van prompt-tokens is doorgaans compute-bound (rekenkracht-gebonden), beperkt door de rekenkundige doorvoer.
• Decode: Sequentiële generatie is doorgaals memory-bandwidth-bound (geheugenbandbreedte-gebonden), beperkt door de noodzaak om voor elk nieuw token herhaaldelijk de groeiende Key-Value (KV) cache en modelgewichten in te lezen.

Conventionele schaalstrategieën (het vergroten van de breedte of diepte van het model) verhogen de kosten voor beide fasen gelijktijdig, omdat elke toegevoegde laag tijdens zowel de prefill als de decode wordt geëvalueerd. Dit creëert een rigiditeit waarbij het verbeteren van de kwaliteit van de generatie (wat vaak profiteert van extra rekenkracht) onvermijdelijk de kosten voor het verwerken van de prompt verhoogt. Het artikel vraagt zich af of het mogelijk is om aanvullende geleerde berekening specifiek toe te wijzen aan continuatie-voorspelling (decode), terwijl de prompt-brede berekening en een enkele persistente KV-cache behouden blijven.

Methodologie: De Dual-Flow Transformer

De auteurs stellen de Dual-Flow Transformer voor, een architectuur die de primaire pad voor prompt-verwerking structureel ontkoppelt van een hulp-pad voor continuatie-voorspelling.

Kernarchitectuur

Het model bestaat uit twee residuele toestands-trajecten die het merendeel van de parameters van het model delen, maar verschillen in hun state-management en activatie-scope:

1. Primary Flow: Een volledig, standaard causaal Transformer-pad. Het verwerkt de volledere prompt, schrijft de persistente KV-cache en bepaalt de primaire staat. Het is zelfvoorzienend en is niet afhankelijk van de hulp-flow.
2. Auxiliary Flow: Een residueel traject dat de primaire staat en cache leest, maar nooit schrijft naar de persistente KV-cache en nooit terugkoppelt naar de primaire flow.
   • Activatie: De hulp-flow wordt weggelaten tijdens de initiële prompt-verwerking (prefill). Deze wordt pas geactiveerd op de laatste positie van de prompt (om het eerste continuatie-token te voorspellen) en op elke daaropvolgende gegenereerde positie.
   • Gewichtsdeling: Beide flows delen alle belangrijke attention-, MLP- en output-matrices.
   • Scheiding: Ze maken gebruik van aparte token-embeddings en lichtgewicht, geleerde koppelingsvectoren.

Kernmechanismen

• Asymmetrische gedeelde KV-Attention: De primaire flow voert standaard causale self-attention uit. De hulp-flow genereert een nieuwe query ($Q_2$), die via een geleerd vector ($a_\ell$) wordt gekoppeld aan de primaire query ($Q_1$). Cruciaal is dat de hulp-flow aandacht besteedt aan de Keys en Values van de primaire flow ($K_1, V_1$). Dit stelt de hulp-flow in staat om de gecachte staat te benutten zonder een eigen staat te schrijven, wat gegroepeerde executie mogelijk maakt waarbij gewichten en KV-regio's één keer worden ingeladen en beide flows bedienen.
• Lichtgewicht Koppeling: Informatie stroomt unidirectioneel van de primaire naar de hulp-stroom via geleerde vectoren ($a_\ell, b_\ell, c_\ell$) op het niveau van de query, de MLP-intermediaire laag en de MLP-output. Dit verrijkt de representatie van de hulp-stroom zonder afhankelijkheden te creëren die het primaire pad zouden veranderen.
• Mixture Objective: De uiteindelijke voorspelling is een mengeling van de primaire en hulp-distributies: $p_{mix} = \alpha_t p_{primary} + (1-\alpha_t) p_{aux}$. Het menggewicht $\alpha_t$ wordt dynamisch aangepast op basis van de betrouwbaarheid (botsingskans) van de primaire distributie, waardoor de primaire flow dominant blijft (minimaal 50% gewicht) tijdens inferentie.
• Router Replay (voor MoE): In Mixture-of-Experts (MoE) configuraties hergebruikt de hulp-flow dezelfde set geselecteerde experts als de primaire flow (bepaald door de primaire router), maar past deze toe met zijn eigen menggewichten. Dit behoudt de selectielogica van de primaire flow terwijl de uitgevoerde aritmetiek op de experts wordt vergroot.

Inferentie vs. Trainingskosten

• Inferentie: De prefill-fase brengt geen extra kosten met zich mee vergeleken met een standaard Transformer; de hulp-flow wordt volledig overgeslagen. De decode-fase brengt extra rekenkracht met zich mee (ongeveer een verdubbeling van de aritmetiek per token), maar verhoogt niet de grootte van de KV-cache of de set unieke gewichten die gelezen moeten worden, aangezien de twee flows gegroepeerd zijn.
• Training: Beide flows worden over de volledige sequentie geëvalueerd, wat resulteert in ongeveer 2x de training FLOPs van het primaire model. De toename in aantal parameters is echter minimaal (alleen aparte embeddings en koppelingsvectoren).

Belangrijkste Bijdragen

1. Architectuur: Introductie van een asymmetrisch gedeeld-KV ontwerp dat de prompt-brede berekening ontkoppelt van de continuatie-berekening. De primaire flow beheert de staat-persistentie, terwijl de hulp-flow een "residueel" traject biedt voor verbeterde voorspelling zonder de primaire staat te wijzigen.
2. Empirische Kwaliteit: Toonde consistente reducties in validatieverlies aan over drie experimentele assen:
   • Data Scaling (NanoGPT): Dual-Flow presteert beter dan standaard Transformers bij een gelijkwaardig tokenbudget, met een lagere passende verlies-ondergrens (fitted loss floor).
   • Model Size Scaling (Dense LLaMA): De winst blijft bestaan naarmate de modelgrootte toeneemt.
   • Sparse MoE: Gevalideerd in ijle (sparse) LLaMA-achtige MoE-configuraties met router replay.
3. Fase-specifieke MoE Allocatie: In tegenstelling tot eerdere parallelle stroomontwerpen die een uniforme rekenkracht-allocatie toepassen, maakt Dual-Flow het mogelijk om de primaire en hulp-expert fan-outs onafhankelijk te configureren. Dit maakt een afruilruimte mogelijk tussen prefill-kosten, decode-kosten en voorspellende kwaliteit, waardoor systemen kunnen optimaliseren voor specifieke workloads (bijv. de prefill-kosten vastzetten terwijl de decode-rekenkracht wordt verhoogd).

Resultaten

• Validatieverlies: Over gematchte token-budgetten in NanoGPT-experimenten bereikt Dual-Flow consequent een lager validatieverlies dan standaard Transformers. De passende asymptotische verlies-ondergrens daalde van 2.9416 (Transformer) naar 2.9013 (Dual-Flow).
• Reken-efficiëntie: In een vergelijking waarbij de training-rekenkracht gematcht is (waarbij een Dual-Flow model met de helft van de datastappen ongeveer dezelfde gedeelde-backbone FLOPs bereikt als een Transformer met dubbele stappen), behaalde Dual-Flow nog steeds een lager validatieverlies. Dit suggereert dat het alloceren van rekenkracht over interagerende stromen effectiever is dan simpelweg de hoeveelheid data voor een enkele stroom te vergroten.
• Ablaties: Componentanalyse toonde aan dat de minimale Dual-constructie (aparte embeddings, gedeelde KV) al beter presteert dan de dichtstbijzijnde eerdere stand van zaken (PHD-2). Het toevoegen van cross-flow koppeling en de waarschijnlijkheidsmengeling verbeterde de prestaties verder.

Betekenis en Claims

Het artikel claimt dat Dual-Flow een structurele mogelijkheid biedt voor fase-specifieke expert-allocatie. Door het pad van de prompt te ontkoppelen van het continuatie-pad, stelt de architectuur serving-systemen in staat om de berekening tijdens de generatie als een schaalbare bron te behandelen zonder proportioneel de kosten van de prompt-verwerking te verhogen.

Specifiek benadrukken de auteurs dat in MoE-modellen de primaire en hulp-flows onafhankelijke expert fan-outs kunnen hebben. Dit creëert een drieweg-afruimknop:

1. Prefill Kosten: Bepaald door het expert-gebruik van de primaire flow.
2. Decode Kosten: Bepaald door de som van het primaire en hulp-expert-gebruik.
3. Voorspellende Kwaliteit: Verbeterd door de totale rekenkracht die aan de continuatie-fase wordt toegewezen te vergroten.

Dit maakt deployment-specifieke flexibiliteit mogelijk waarbij een systeem de kosten voor de prompt-berekening kan behouden terwijl het de beschikbare decode-budget gebruikt voor meer hulp-experts, of vice versa, zonder de gedeelde parameterpool of de enkele-primaire-cache structuur te veranderen. Het werk ondersteunt de premisse dat de berekening van het volgende token een nuttige schaalbare bron is die intern binnen het model kan worden gealloceerd in plaats van extern via langere redeneer-traces of meerdere sampling-methoden.