Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing

Technische Samenvatting: Lineaire Aandachtsarchitecturen: Mechanismen, Trade-offs en Cross-Layer Routing

Probleemstelling
Standaard softmax self-attention stelt elk token in staat om informatie op te halen uit de volledige context, een mechanisme dat centraal staat voor de expressiviteit van Transformer-taalmodellen. Echter, deze token-naar-token uitwisseling brengt een kwadratische kostenpost met zich mee in de sequentielengte ($O(T^2)$) voor zowel training als inferentie, wat een dominante bottleneck vormt naarmate de contextvensters groter worden. Hoewel lineaire aandachtsarchitecturen een oplossing bieden door aandacht te herformuleren als een recurrente update over een constante-grootte geheugenmatrix (wat resulteert in lineaire trainingstijd en constante inferentietijd per token), leden vroege varianten aanzienlijk onder nauwkeurigheidsverlies door interferentie tussen opgeslagen associaties en een gebrek aan mechanismen voor vergeten of selectief bewerken. Recente ontwikkelingen, waaronder DeltaNet en zijn opvolgers, hebben deze kloof verkleind door foutcorrigerende delta-regels en fijnmazige controle-gates te introduceren, maar de specifieke trade-offs tussen deze mechanismen, hun schaalingsgedrag en het potentieel voor lichtgewicht cross-layer informatie-uitwisseling moeten nog systematisch in kaart worden gebracht.

Methodologie en Raamwerk
De auteurs presenteren een vergelijkende studie van softmax attention en vier recente recurrente lineaire aandachtsarchitecturen: DeltaNet, Gated DeltaNet, Kimi Delta Attention en Gated DeltaNet-2.

1. Uniforme Notatie: Het artikel stelt een gemeenschappelijke recurrente-geheugennotatie op om deze mechanismen uit te drukken. Dit raamwerk isoleert de rollen van de recurrente geheugenmatrix $W$, de delta-regel residu $r$ (de fout tussen de huidige waarde en de voorspelling van het geheugen), en de diverse decay-, erase- en write-controlemechanismen.

   • DeltaNet: Vervangt naïeve additieve updates door een foutcorrigerende delta-regel, waarbij alleen het residu $r = v - \bar{v}$ naar het geheugen wordt geschreven.
   • Gated DeltaNet: Voegt een geleerde scalaire decay-gate $\alpha$ toe aan de geheugenstatus vóór het berekenen van het residu, wat een grove vergetensmechanisme introduceert.
   • Kimi Delta Attention: Verfijnt het vergeten door een kanaal-specifieke decay-vector $\alpha$ te gebruiken, waardoor verschillende feature-dimensies met verschillende snelheden kunnen vervagen.
   • Gated DeltaNet-2: Verfijnt de update verder door aparte kanaal-specifieke erase- en write-gates te introduceren, waardoor onafhankelijke controle mogelijk is over welke key-kanalen worden verwijderd en welke value-kanalen worden opgeslagen.

2. Experimentele Opzet: De studie maakt gebruik van 350M-parameter decoder-only modellen getraind op 15B tokens (FineWeb-Edu). De experimenten vergelijken validatieverlies, trainingsdoorvoersnelheid (throughput), optimizer-gevoeligheid (AdamW versus Muon) en sequentielengte-schaling. De auteurs voeren ook grotere runs uit (1.3B en 3B parameters) en evalueren downstream taken (HellaSwag, PIQA, WinoGrande).

3. Cross-Layer Routing: Om informatieverwatering in diepe stacks aan te pakken, introduceert en evalueert het artikel lichtgewicht routingmechanismen. Het onderzoekt twee assen: wat te routeren (write error versus write value) en waar het te injecteren (het value target van de volgende laag versus de gedeelde hidden stream).

   • CLER (Cross-Layer Error Residuals): Stuurt het delta-regel write error naar de value target van de volgende laag.
   • CLER-H: Routeert de write error naar de gedeelde hidden stream.
   • CLVR (Cross-Layer Value Routing): Routeert de layer's write value (het doel vóór aftrekking) naar de gedeelde hidden stream via een met nul geïnitialiseerde projectie.

Belangrijkste Resultaten

• Architectuur- en Optimizer Trade-offs: Binnen de 350M-parameter, 15B-token sweep bereikt Kimi Delta Attention, getraind met de Muon optimizer in een hybrid stack (afgewisseld met softmax lagen), het laagste uiteindelijke validatieverlies (2.273). Echter, pure Gated DeltaNet stacks getraind met AdamW vertonen de hoogste genormaliseerde trainingsdoorvoersnelheid (gesteld op 100% baseline). Hybride stacks verbeteren de loss over het algemeen ten koste van enige throughput, terwijl pure stacks de sterkste iteratietijd-schaling voor lange contexten behouden.
• Optimizer Gevoeligheid: De keuze van de optimizer is onlosmakelijk verbonden met de leercurve (learning rate) afstemming. Muon geeft de voorkeur aan lagere leercurves ($3 \times 10^{-4}$) voor lineaire aandachtsvarianten, terwijl AdamW-varianten hogere rates ($10^{-3}$) verkiezen. Muon verlaagt consequent het uiteindelijke validatieverlies ten opzichte van AdamW bij gematchte architectuurinstellingen.
• Schaalingsgedrag: Pure recurrente stacks vertonen superieur schalingsgedrag met de sequentielengte. Bij 32k tokens vertoont een pure Gated DeltaNet een minimale toename in iteratietijd (~1.1x van 4k), terwijl softmax attention met ~2.9x toeneemt. Hybride stacks zitten daar tussenin, waarbij de toename in iteratietijd ruwweg evenredig is aan de dichtheid van hun softmax lagen.
• Bevindingen over Cross-Layer Routing:
  • CLER (Error naar Value Target): Het doorsturen van de delta-regel write error naar de value target van de volgende laag leidt niet tot verbeterde prestaties ten opzichte van de gematchte baselines. De auteurs schrijven dit toe aan een "ruimte-mismatch", waarbij het gerouteerde residu concurreert met de eigen geleerde value-geometrie van de ontvanger.
  • CLVR (Value naar Hidden Stream): Het routeren van de write value naar de gedeelde hidden stream levert een bescheiden maar consistente verbetering op in het uiteindelijke validatieverlies voor zowel DeltaNet als Gated DeltaNet. Dit suggereert dat het nuttige cross-layer signaal de write target van de laag is, en niet de delta-regel correctie.
  • Verminderde Meeropbrengst: De winsten van CLVR zijn klein (vermindering van de loss met ~0.01 bij 350M/1B tokens) en nemen af naarmate de model schaal of het aantal getrainde tokens toeneemt, wat suggereert dat het voordeel het meest uitgesproken is wanneer de basisgeheugen nog niet alle noodzakelijke informatie heeft vastgelegd.

Betekenis en Claims
Het artikel verklaart geen enkele "beste" architectuur, maar beoogt de ontwerpplek begrijpelijk te maken. De primaire bijdragen zijn:

1. Een Uniform Raamwerk: Het biedt een gemeenschappelijke notatie die verduidelijkt hoe recente lineaire aandachtsvarianten verschillen in expressiviteit, geheugenverval en controle-granulariteit.
2. Empirische Trade-off Mapping: Het demonstreert dat de keuze tussen lineaire en hybride stacks, en tussen optimizers, een multi-objectieve trade-off inhoudt tussen validatieverlies, throughput en sequentielengte-schaling.
3. Routing Inzichten: Het identificeert dat hoewel naïeve error-routing (CLER) faalt door basis-misalignment, het routeren van de write value naar de uitgelijnde hidden stream (CLVR) een lichtgewicht, lineaire-tijd pad biedt voor cross-layer informatie-uitwisseling dat de prestaties verbetert zonder de recurrente structuur te verbreken.

De auteurs benadrukken dat hun conclusies gebaseerd zijn op single-run vergelijkingen en specifieke token-budgetten, en adviseren voorzichtigheid bij het generaliseren van kleine verschillen in validatieverlies zonder verificatie via herhaalde seeds. Het werk dient als fundament voor het begrijpen van de specifieke mechanismen die de prestaties in lineaire aandacht drijven en wijst op openstaande vragen over de generalisatie van routingstrategieën naar complexere architecturen zoals Gated DeltaNet-2.