Technische Samenvatting: Multi-Mixer Modellen (Oryx)

Probleemstelling

Softmax-attention is het dominante mechanisme in moderne grote taalmodellen (LLM's), maar brengt kwadratische rekenkosten en lineaire geheugeneisen met zich mee ten opzichte van de sequentielengte, vanwege het onderhouden van een volledige key-value (KV) cache. Hoewel lineaire recurrente modellen (zoals lineaire attention en state space-modellen zoals Mamba-2) lineaire berekening en constant geheugen bieden, blijven ze vaak achter bij attention-gebaseerde modellen bij taken die sterke retrieval of in-context learning vereisen. Bestaande hybride architecturen proberen deze afwegingen te verzachten door attention en recurrente blokken statisch te intercaleren over lagen (inter-layer) of ze te fuseren binnen een enkele laag (intra-layer). Deze benaderingen fixeren echter het rekenkundige pad per token of laag, waardoor ze niet in staat zijn het modelleringmechanisme aan te passen aan de specifieke behoeften van verschillende segmenten binnen een sequentie.

Methodologie: De Oryx Architectuur

De auteurs stellen Oryx voor, een hybride architectuur op de sequentie-as die flexibele schakeling tussen verschillende sequentiemixers (bijvoorbeeld kwadratische softmax-attention en lineaire recurrente mechanismen) mogelijk maakt gedurende een enkele tokensequentie.

Kernontwerpprincipes

1. Gedeelde Representaties via Gekoppelde Projecties:

   • Oryx verenigt attention en lineaire modellen onder een associatief geheugenperspectief, waarbij beide key-value (KV) associaties onderhouden.
   • De architectuur maakt gebruik van gekoppelde projecties voor keys ($W_K$) en values ($W_V$) over mixer-modi heen. Dit stelt een enkele forward pass in staat om zowel de attention KV-cache als de lineaire recurrente staat gelijktijdig bij te werken met behulp van gedeelde representaties.
   • Mixer-specifieke Queries: In tegenstelling tot keys en values worden query-projecties ($W_Q$) apart gehouden voor elke mixer. Empirische resultaten geven aan dat het delen van query-weights de prestaties verslechtert, wat suggereert dat attention- en recurrente mechanismen verschillende readout-vectoren nodig hebben om modusspecifieke informatie uit de gedeelde staat te halen.

2. Architectonische Componenten:

   • Het blok bevat componenten die kritiek zijn voor de prestaties van lineaire modellen, zoals korte convoluties en multiplicatieve poorten, die worden toegepast op de gedeelde keys en values.
   • Om compatibiliteit te waarborgen, dwingt Oryx een eenheidshoofdstructuur (Multi-Head Attention, MHA) af over beide mixers, hoewel lineaire modellen zoals Mamba-2 doorgaans Multi-Value (MVA) structuren gebruiken.
   • Het blok onderhoudt zowel de KV-cache als de lineaire recurrente staat op elk tijdstip en werkt deze gezamenlijk bij.

3. Chunked Mixed-Mode Training:

   • Om robuuste modusschakeling tijdens inferentie mogelijk te maken, wordt Oryx getraind met een chunked mixed-mode strategie. Trainingssequenties worden gepartitioneerd in vaste lengte chunks (bijvoorbeeld 128 tokens), waarbij elke chunk willekeurig een mixer-modus (attention of lineair) toegewezen krijgt.
   • Dit trainingsregime dwingt het model om compatibele representaties te leren die naadloze overgangen tussen modi mogelijk maken, waardoor de prestatievermindering wordt voorkomen die vaak optreedt bij het schakelen van lineaire naar attention-mechanismen in sequentieniveau mixed training.

Belangrijkste Bijdragen

• Hybridisatie op de Sequentie-as: Oryx introduceert een nieuwe as voor hybridisatie, waardoor het model mechanismen dynamisch kan schakelen over de tokensequentie, in plaats van beperkt te zijn door vaste laag- of blokarchitecturen.
• Parameter-efficiëntie: Het ontwerp koppelt ten minste 90% van de parameters (specifiek $W_K, W_V, W_G, W_O$) over mixer-modi, waardoor attention- en recurrente modi kunnen opereren over gedeelde interne representaties zonder aanzienlijke parameter-overhead.
• Flexibele Inferentie: Het model ondersteunt inferentie waarbij verschillende segmenten van een sequentie (bijvoorbeeld context prefill versus generatie) kunnen worden verwerkt door verschillende mixers. Bijvoorbeeld, een model kan lineaire recurrentie gebruiken voor efficiënte contextverwerking en overschakelen naar attention voor generatie met veel retrieval.

Empirische Resultaten

Experimenten werden uitgevoerd op modellen variërend van 130M tot 1,4B parameters, met Mamba-2 en Gated DeltaNet (GDN) als lineaire mixers.

• Taalmodellering: Bij vaste tokenbudgetten bereiken Oryx-modellen (zowel Oryx-TM als Oryx-TG varianten) vergelijkbare of betere prestaties dan single-mixer baselines (pure Transformer, pure Mamba-2, pure GDN). Op schaal van 1,4B overtreffen Oryx-varianten hun respectievelijke baselines met minimaal 0,7 procentpunten op gemiddelde taalmodelleringstaken.
• Retrieval-capaciteiten:
  • Oryx behoudt de retrieval-capaciteiten van de samenstellende mixers.
  • Cruciaal, bij gebruik van een mixed-inferentiemodus (context verwerken met lineaire recurrentie en generatie met attention), bereikt Oryx prestaties die vergelijkbaar zijn met de Transformer-baseline, terwijl minder dan 10% van de tokens in attention-modus wordt verwerkt.
  • Bij real-world retrieval-taken overtreft mixed-mode Oryx lineaire baselines met 8,6 procentpunten (Oryx-TG) en 13,5 procentpunten (Oryx-TM).
  • Bij synthetische "Needle-in-a-Haystack" (NIAH) tests overtreft Oryx lineaire baselines met 38,6 tot 40,3 procentpunten, wat aantoont dat gedeelde representaties effectief lange-context retrieval ondersteunen, zelfs wanneer de lineaire mixer het grootste deel van de sequentie verwerkt.
• Modusschakeling: Perplexiteitsanalyse toont aan dat na het schakelen tussen mixers (bijvoorbeeld van attention naar Mamba-2 of vice versa), de perplexiteit van het model snel convergeert naar de baseline van de geselecteerde modus, wat aangeeft dat de gedeelde staten compatibel zijn en de overgang naadloos verloopt.

Betekenis en Claims

Het paper beweert dat Oryx aantoont dat attention en lineaire recurrente mechanismen interne representaties kunnen delen ondanks hun verschillende methodologieën. Deze bevinding daagt de aanname uit dat deze mechanismen volledig gescheiden statenruimtes vereisen.

De betekenis van dit werk ligt in:

1. Validering van Hybridisatie op de Sequentie-as: Het stelt vast dat het dynamisch schakelen van mixers binnen een sequentie een haalbare en effectieve strategie is, wat een nieuwe ontwerpruimte biedt voor adaptieve berekening.
2. Overbrugging van de Retrieval-kloof: Het toont aan dat lineaire modellen retrieval-prestaties kunnen bereiken die vergelijkbaar zijn met Transformers als ze attention-modi kunnen benutten voor specifieke segmenten, zonder dat een volledig attention-mechanisme voor de hele sequentie nodig is.
3. Toekomstige Richtingen: De auteurs suggereren dat deze architectuur kansen opent voor geleerde routing (bijvoorbeeld met behulp van RL om mixers te selecteren op basis van FLOP-besparingen) en speculatieve decoding, waarbij lineaire modi redeneersporen kunnen opstellen ter verificatie door attention-modi.

De auteurs blijven bescheiden, met de opmerking dat het onderhouden van meerdere staten (KV-cache en lineaire staat) geheugen-overhead met zich meebrengt, hoewel dit bij lange contextlengtes wordt gedomineerd door de KV-cache. Zij merken ook op dat de trainingsdynamiek van multi-mixer modellen grotendeels onontgonnen terrein blijft en mogelijk gespecialiseerde doelstellingen vereist naarmate de schaal toeneemt.