Technische Samenvatting: Kalman Linear Attention (KLA)

Probleemstelling

Huidige state-space modellen (SSM's) zoals Mamba en gated linear attention (GLA) bieden lineair-complexe, paralleliseerbare alternatieven voor softmax attention, waarmee ze de kwadratische schaalbaarheidsproblematiek van Transformers aanpakken. Deze modellen vertrouwen echter op lineaire of affine verborgen toestandsupdates, wat hun expressiviteit beperkt in vergelijking met de niet-lineaire interacties die door softmax-normalisatie worden geïnduceerd. Bovendien missen bestaande lineaire RNN's en SSM's expliciete mechanismen voor het representeren van de onzekerheid van de toestand; zij behandelen verborgen toestanden als deterministisch. Dit artikel adresseert de kloof tussen de efficiëntie van lineaire updates en de expressiviteit/robuustheid van niet-lineaire filtering, door de vraag te stellen of exacte Kalman-filterupdates efficiënt kunnen worden geïmplementeerd als een sequentie-mixer.

Methodologie: Kalman Linear Attention (KLA)

De auteurs stellen Kalman Linear Attention (KLA) voor, een probabilistische sequentie-mixer die sequentiemodellering formuleert als exacte Bayesiaanse filtering. De kerninnovatie ligt in het herparameteriseren van de Kalman-filter naar de informatievorm (met behulp van precisies en natuurlijke parameters) om parallelle berekening mogelijk te maken.

1. Probabilistische Formulering

KLA modelleert sequentiemenging als inferentie in een lineair-Gaussisch state-space model:

• Dynamics Prior: De evolutie van de latente toestand wordt gemodelleerd als een Ornstein-Uhlenbeck (OU) proces, dat exact gediscretiseerd is. Dit introduceert een procesruis-term ($\bar{p}_t$) die de onzekerheid in toestandsovergangen vastlegt, wat verschilt van de deterministische dynamiek van standaard SSM's.
• Observatiemodel: Input-tokens worden behandeld als ruizige observaties van de latente toestand, geparametriseerd door een observatie-operator ($k_t$) en een waarde-precisie ($\Lambda^v_t$) die de zekerheid in het token-bewijs representeert.
• Readout: De output wordt afgeleid van het posterieure gemiddelde van de latente toestand, geprojecteerd via een query-operator ($q_t$).

2. Parallelisatie via Möbius-transformaties

Klassieke Kalman-filtering is inherent sequentieel. KLA overwint dit door gebruik te maken van de informatievorm van de filter:

• Precisie-recursie: De update van de posterieure precisie ($\lambda_t$) wordt aangetoond een Möbius (fractie-lineaire) transformatie te zijn van de vorige precisie ($\lambda_{t-1}$).
  $$\lambda_t = \frac{\alpha_t \lambda_{t-1} + \beta_t}{\gamma_t \lambda_{t-1} + \delta_t}$$
  Cruciaal is dat Möbius-transformaties associatief componeren. Dit maakt het mogelijk om de precisie-updates te berekenen via parallelle prefix-scans (associatieve scans) met $O(T)$ werk en $O(\log T)$ diepte, wat overeenkomt met de parallelle efficiëntie van moderne SSM's zoals Mamba.
• Gemiddelde-recursie: Zodra het precisiepad parallel is berekend, volgen de updates van het posterieure gemiddelde (informatie-gemiddelde) een affine transformatie die eveneals paralleliseerbaar is.

3. Niet-lineaire Gating vanuit Onzekerheid

In tegenstelling tot lineaire SSM's waar gating vaak lineair of token-afhankelijk maar deterministisch is, ontstaat de gating van KLA natuurlijk uit de ratio van precisies. De "forget gate" en "write gate" zijn geschiedenis-afhankelijk en niet-lineair, gedreven door de relatieve onzekerheid tussen de prior-toestand en de nieuwe observatie. Dit stelt het model in staat om informatie selectief op te nemen op basis van vertrouwen, zonder dat hiervoor expliciete, geleerde gating-parameters voor elke token nodig zijn.

4. Architectuur

KLA is ontworpen als een drop-in vervanging voor standaard SSM- of attention-lagen. Het volgt een gefuseerde MLP-architectuur vergelijkbaar met Mamba, waarbij 1D causale convoluties en SiLU-activaties worden geïntegreerd. De auteurs implementeren efficiënte CUDA- en Triton-kernels voor de Möbius-scan om te voorkomen dat uitgebreide toestanden in high-bandwidth memory (HBM) worden gematerialiseerd.

Belangrijkste Bijdragen

1. Associatieve Herparameterisering: De auteurs demonstreren dat diagonale lineair-Gaussische Kalman-filtering in de informatievorm een Möbius-structuur toelaat, wat exacte parallelle prefix-scans voor Bayesiaanse filtering mogelijk maakt.
2. Onzekerheid-gestuurde Niet-lineariteit: KLA introduceert niet-lineaire gating afgeleid van precisieratio's, wat strikt grotere expressiviteit biedt dan lineaire/affine SSM-updates, terwijl de lineaire-tijd scan-structuur behouden blijft.
3. KLA-laag: Een praktische, drop-in sequentie-mixer die een expliciete belief-state onzekerheid produceert en gestapeld kan worden in diepe netwerken.
4. Schaalbaarheid en Validatie: Het artikel presenteert de eerste gestapelde Bayesiaanse-filtering primitieven die getraind zijn op de schaal van miljarden tokens, waarbij de stabiliteit en prestaties worden gevalideerd.

Experimentele Resultaten

De auteurs evalueren KLA over synthetische taken, state-tracking benchmarks en grootschalige taalmodellering.

• Synthetische Taken (MAD-Lab): KLA presteert beter dan GDN, GLA, Mamba en mLSTM op taken die selectie onder onzekerheid vereisen, specifiek Fuzzy Recall en Selective Copy. Het behaalt state-of-the-art resultaten op Compression, wat suggereert dat de filtering in de informatievorm fungeert als een optimale context-compressor.
• State Tracking (A5 Permutatie): In de A5-permutatie-compositie taak (een benchmark voor NC1-complete state tracking) lost KLA de taak op met constante diepte (1–2 blokken). In contrast hiermee vereisen lineaire SSM's en attention-gebaseerde modellen een diepte die groeit met de sequentielengte. Dit demonstreert dat de fractie-lineaire updates van KLA de expressiviteitslimieten van lineaire recurrente modellen doorbreken terwijl ze parallel trainbaar blijven.
• Long-Context Associative Recall (MQAR): KLA bereikt bijna perfecte nauwkeurigheid (>95%) op de Multi-Query Associative Recall taak met een sequentielengte van 2048, waarbij het Mamba overtreft en GLA significant overtreft, dat de taak in deze setting niet leert.
• Grootschalige Pretraining:
  • Standalone: Bij 45M en 180M parameters evenaart of verbetert de standalone KLA GPT, Mamba en GDN op zero-shot commonsense benchmarks (bijv. HellaSwag, PIQA).
  • Hybride: Het vervangen van de laatste attention-laag van een GPT door KLA (GPT+KLA) levert de sterkste prestaties op onder hybriden, waarbij het GPT+Mamba en GPT+GDN overtreft.
• Ablatie: De auteurs tonen aan dat de Ornstein-Uhlenbeck (OU) prior cruciaal is voor stabiele diepe stapeling. Zonder de mean-reverting stochastische dynamica wordt het trainen van diepe stacks van KLA instabiel.

Betekenis en Claims

Het artikel claimt dat KLA er succesvol in slaagt de kloof te overbruggen tussen de efficiëntie van lineaire state-space modellen en de expressiviteit van niet-lineaire attention mechanismen. Door sequentiemenging te beschouwen als exacte Bayesiaanse filtering, biedt KLA:

• Expliciete Onzekerheid: In tegenstelling tot deterministische mixers, draagt KLA een expliciete belief-state onzekerheid met zich mee die de informatie-gating direct controleert.
• Strikt Hogere Expressiviteit: De Möbius-updates maken het mogelijk voor KLA om state-tracking taken (zoals A5) op te lossen die lineaire SSM's en attention niet kunnen, zonder de parallelle trainingscapaciteiten op te offeren.
• Schaalbaarheid: Het werk demonstreert dat exacte Bayesiaanse filtering primitieven gestapeld en getraind kunnen worden op de schaal van miljarden tokens, wat de opvatting uitdaagt dat probabilistische filtering inherent sequentieel of computationeel prohibitief is voor grootschalige taalmodellering.

De auteurs positioneren KLA niet louter als een theoretische curiositeit, maar als een praktische, drop-in primitief die de robuustheid en state-tracking capaciteiten van moderne taalmodellen verbetert, terwijl de computationele efficiëntie behouden blijft die vereist is voor lang-context verwerking.