Technische Samenvatting: Kaczmarz Lineaire Attention (KLA)

Probleemstelling

Modellering van taal met lange context wordt gehinderd door de kwadratische computatie- en geheugenkosten van standaard Softmax-attention in Transformers. Hoewel recurrente modellen met lineaire tijd dit aanpakken door context te comprimeren tot een statische van vaste grootte, wordt het ontwerp van de regel voor stat-update de kritieke bottleneck. Bestaande benaderingen, zoals Gated DeltaNet (GDN), maken gebruik van een "delta-regel" om het residu tussen een doelwaarde en de huidige statvoorspelling te schrijven, waardoor bestaande associaties worden bewerkt in plaats van alleen opgehoopt.

GDN vertrouwt echter op een geleerde scalair coëfficiënt om vergeten en updaten in evenwicht te brengen. De auteurs stellen dat deze coëfficiënt een empirische ontwerpkeuze is in plaats van een theoretisch afgeleide stapgrootte. Een ongepaste coëfficiënt kan leiden tot suboptimale update-magnitudes: te sterke correctie voor keys met een grote norm en te zwakke correctie voor keys met een kleine norm, wat resulteert in statcorruptie en suboptimale ophaalprestaties.

Methodologie: Kaczmarz Lineaire Attention

Het artikel stelt Kaczmarz Lineaire Attention (KLA) voor, een wijziging van GDN die de geleerde scalair vervangt door een theoretisch afgeleide stapgrootte gebaseerd op de Kaczmarz-projectiemethode.

Theoretische Afleiding

De auteurs formuleren de schrijfbewerking van key-value als een constrained optimalisatieprobleem. Voor een token $t$ is het doel om de stat $S_t$ te vinden die het dichtst bij de vervallen vorige stat $\tilde{S}_{t-1}$ ligt en die voldoet aan de lineaire constraint $S_t^T k_t = v_t$ (waarbij $k_t$ de key is en $v_t$ de value).

• Optimalisatiedoel: $\min_S \frac{1}{2}\|S - \tilde{S}_{t-1}\|_F^2$ onder de voorwaarde $S^T k_t = v_t$.
• Oplossing: De optimale oplossing is de orthogonale projectie van $\tilde{S}_{t-1}$ op het constraint-hypervlak. Dit levert een stapgrootte (coëfficiënt) op van $\beta_t = 1/\|k_t\|_2^2$.

De KLA Update-regel

KLA generaliseert deze exacte projectie naar een praktische update-regel die een leerbare relaxatiepoort $\eta_t$ en een numerieke stabilisator $\epsilon$ omvat:
$$\beta_t = \frac{\eta_t}{\|k_t\|_2^2 + \epsilon}$$
De stat-update wordt dan:
$$S_t = \tilde{S}_{t-1} + \beta_t k_t e_t^T$$
waarbij $e_t = v_t - \tilde{S}_{t-1}^T k_t$ het residufout is.

Belangrijkste Eigenschappen

Het artikel stelt vast dat deze coëfficiënt drie wenselijke eigenschappen gelijktijdig voldoet (in het niet-gerelaxte geval waar $\eta_t=1, \epsilon=0$):

1. Orthogonale Projectie: Het voert de verandering met de kleinste norm uit aan de stat om aan de constraint te voldoen.
2. Exacte Lijnzoektocht: Het is de exacte minimalizer van het ver-token verlies langs de update-richting.
3. Genormaliseerde Gradientstap: Het fungeert als een genormaliseerde Stochastic Gradient Descent (SGD)-stap.

Implementatie

KLA is ontworpen als een drop-in vervanging voor GDN. Het behoudt:

• De statvorm ($S \in \mathbb{R}^{d_k \times d_v}$).
• De poortmechanismen (decay- en update-poorten).
• De lineaire recurrente structuur.
• Het parallelle trainingsalgoritme per chunk.

De enige structurele wijziging is de berekening van de diagonale coëfficiëntmatrix $B = \text{Diag}(\beta_1, \dots, \beta_C)$ binnen de chunkwise solver. Er zijn geen nieuwe hardware-kernels of statopstellingen vereist.

Belangrijkste Bijdragen

1. Zelf-genormaliseerde Delta-regel: Afleiding van de KLA-coëfficiënt uit de Kaczmarz-projectie, waarmee wordt bewezen dat het een orthogonale projectie, exacte lijnzoek-minimalizer en genormaliseerde gradientstap is voor het ver-token verlies.
2. Drop-in Recurrent Algoritme: Aantonen dat KLA de GDN-recurrentie en chunkwise solver behoudt, waarbij alleen een wijziging in de diagonale coëfficiëntmatrix vereist is.
3. Empirische Winsten: Validatie dat KLA een lagere perplexiteit en betere stat-tracking bereikt dan GDN en andere lineaire tijd-baselines zonder architecturale uitbreiding.
4. Behouden Efficiëntie: Bevestiging dat KLA $O(n)$ tijd/geheugencomplexiteit behoudt en de prefill-latentie van GDN evenaart, terwijl het de decode-doorvoer verbetert.

Experimentele Resultaten

Experimenten werden uitgevoerd op de schaal van 0,4B parameters met een budget van 1B tokens op de SlimPajama-672B-dataset.

Taalmodellering

• Validatie Perplexiteit (PPL): KLA behaalde 8,09 PPL, beter dan GDN (8,50), DeltaNet (8,26), GLA (8,42), Longhorn (8,79) en Mamba2 (9,41).
• Extrapolatie van Lange Context: Zonder fine-tuning voor lange context behield KLA de laagste perplexiteit over contextlengtes van 1K tot 65K tokens, met een consistente PPL-voorsprong van 3,7–4,1 ten opzichte van GDN.

Synthetische Taken

• Single-Needle-in-a-Haystack (S-NIAH): KLA bereikte 100% nauwkeurigheid over alle geteste contextlengtes (1K–8K), terwijl GDN licht daalde en Mamba2 significant faalde.
• Multi-Query Associative Recall (MQAR): KLA verbeterde de nauwkeurigheid bij 8× lengte-extrapolatie met 7,03 punten ten opzichte van GDN (73,84% versus 66,81%).
• Stack & Palindrome: KLA toonde concurrerende of superieure prestaties in stat-tracking en omkeringstaken.

Efficiëntie

• Prefill Latentie: KLA evenaarde de prefill-latentie van GDN (bijv. ~1220ms bij 131K context), wat bevestigt dat de berekening van de Kaczmarz-coëfficiënt geen invloed heeft op de dominante kosten van de chunkwise kernel.
• Decode Doorvoer: Bij 32K context behaalde KLA 3,54 tokens/s, wat 2,1× hoger is dan GDN (1,65 tokens/s) en 1,7× hoger dan Mamba2.

Ablatiestudies

De auteurs ablateerden de normalisatiestrategie, het poortmechanisme en statuitbreiding. Resultaten gaven aan dat key-norm normalisatie cruciaal is voor stabiliteit; het verwijderen ervan of het gebruik van alleen een geleerde scalair leidde tot significante degradatie in pretrainingsstabiliteit (PPL) en prestaties op synthetische taken.

Betekenis en Claims

Het artikel beweert dat de key-norm-genormaliseerde Kaczmarz-coëfficiënt een "ontwerpas van de eerste orde" is voor delta-regel sequentiemodellen. Door de schrijfbewerking te behandelen als een constrained projectie, lost KLA het probleem van suboptimale update-magnitudes op dat inherent is aan empirisch geleerde scalar.

De betekenis ligt in het feit dat deze verbeteringen worden bereikt zonder de recurrente stat, hardware-kernel of architecturale complexiteit te wijzigen. De auteurs stellen dat de coëfficiënt zelf, in plaats van de statcapaciteit of poortstructuur, de primaire factor is die bepaalt of een delta-regel schrijfbewerking succesvol zijn lokale regressieprobleem oplost. Dit suggereert dat toekomstige lineaire tijd-modellen aanzienlijke winsten in nauwkeurigheid, extrapolatie en decode-efficiëntie kunnen behalen door simpelweg deze theoretisch onderbouwde stapgrootte toe te passen.