Technische Samenvatting: Stochastische Attention

Probleemstelling

Efficiënte attention-mechanismen zijn cruciaal voor het schalen van Transformer-modellen, maar bestaande benaderingen staan voor een afweging tussen rekenkosten en het bestrijken van het receptieve veld. Sliding-Window Attention (SWA) reduceert de complexiteit tot $O(nw)$ per laag door tokens te beperken tot een lokaal venster van grootte $w$. Deze deterministische localiteit beperkt het effectieve receptieve veld echter tot lineaire groei ($\ell w$ na $\ell$ lagen), waardoor grote delen van de sequentie onbereikbaar blijven wanneer $w \ll n$. Hoewel er remedies bestaan zoals globale tokens, handgemaakte sparse patronen of block-level routing, introduceren deze vaak architecturale complexiteit of falen ze om bewezen dekkingst garanties te bieden. Daarentegen biedt volledige attention globale connectiviteit, maar met verbodsbepalende $O(n^2)$ kosten.

Het artikel put inspiratie uit het connectoom van de fruitvlieg (Drosophila melanogaster), een netwerk van ongeveer 130.000 neuronen met een connectieprobabiliteit van slechts 0,02%. Ondanks deze sparsiteit bereikt het netwerk een gemiddelde padlengte van ongeveer 4,4 sprongen. Deze "small-world" topologie ontstaat uit het samenspel van dichte lokale clustering en breed verdeelde langeafstandsverbindingen die fungeren als stochastische shortcuts. De auteurs stellen dat efficiënte globale informatieflow kan ontstaan uit gestructureerde lokale berekening, aangevuld met gedistribueerde, stochastische langeafstandsverbindingen, in plaats van dat dichte all-to-all connectiviteit vereist is.

Methodologie: Stochastische Attention (SA)

De auteurs stellen Stochastische Attention (SA) voor, een parameterloze uitbreiding voor SWA die de routingstrategie van het connectoom nabootst.

Kernmechanisme

SA transformeert een vast lokaal venster in een stochastisch globaal venster via een drie-trapsproces dat wordt toegepast voor en na de standaard SWA-berekening:

1. Permuteren: Een willekeurige permutatie $\sigma \sim \text{Uniform}(S_n)$ wordt onafhankelijk voor elke laag getrokken. De tokensequentie-representaties ($Q, K, V$) worden volgens $\sigma$ herschikt.
2. Venster-attention: Standaard SWA wordt toegepast in de gepermuteerde ruimte. Omdat de permutatie willekeurig is, beslaat het vaste venster van grootte $w$ nu een willekeurige subset van de originele sequentie.
3. Permutatie ongedaan maken: De originele volgorde wordt hersteld door $\sigma^{-1}$ toe te passen op de output.

In de originele tokenruimte attendeert positie $i$ op een willekeurige buurt $\tilde{N}^\sigma_w(i)$ met een verwachte grootte $w$, uniform verdeeld over de volledige sequentie ongeacht de originele afstand $|i-j|$. Dit creëert een paarsgewijze connectieprobabiliteit van ongeveer $w/n$ voor elk tokenpaar, analoog aan de gedistribueerde langeafstandsshortcuts in het connectoom.

Theoretische Eigenschappen

• Groei van het receptieve veld: Terwijl het receptieve veld van SWA lineair groeit ($O(n/w)$ lagen voor volledige dekking), leveren de onafhankelijk bemonsterde permutaties over de lagen van SA exponentieel groeiende receptieve velden op. Het artikel bewijst dat volledige sequentiedekking wordt bereikt in $O(\log_w n)$ lagen.
• Complexiteit: De permutatie- en inverse-permutatie-operaties zijn $O(n)$ index-herschikkingen. De attention-berekening blijft $O(nw)$. SA behoudt dus hetzelfde $O(nw)$ budget per laag als standaard SWA.
• Causale beperkingen: In autoregressieve settings wordt het causale masker toegepast in de gepermuteerde ruimte, zodat tokens alleen attenderen op causaal toegankelijke voorgangers, terwijl toch stochastische langeafstandsverbindingen mogelijk blijven.

Gated SA + SWA Combinatie

Om het "small-world" regime na te bootsen (hoge lokale clustering + korte paden), stellen de auteurs een dual-path architectuur voor die SA en SWA combineert:
$$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$$
waarbij $Y_{sa}$ en $Y_{swa}$ de outputs zijn van respectievelijk het stochastische en het sliding-window pad. De gating-weights ($g_{sa}, g_{swa}$) worden berekend via onafhankelijke sigmoid-projecties, waardoor het model leert om lokale coherentie (van SWA) en globale dekking (van SA) in evenwicht te brengen. Deze combinatie voegt verwaarloosbare parameters toe en behoudt $O(nw)$ complexiteit.

Belangrijkste Bijdragen

1. Stochastische Attention (SA): Een plug-and-play, parameterloze uitbreiding voor SWA die willekeurige permutaties toepast om exponentiële groei van het receptieve veld ($O(\log_w n)$) te bereiken binnen hetzelfde $O(nw)$ rekenbudget.
2. Gated SA + SWA: Een dual-path architectuur die theoretisch en empirisch het small-world regime reproduceert dat wordt waargenomen in biologische connectomen, en lokale clustering en stochastische shortcuts in evenwicht brengt.
3. Empirische Validatie:
   • Pre-training: Op modellen van 360M parameters die vanaf nul zijn getraind, bereikt de gated SA + SWA-combinatie de beste gemiddelde zero-shot nauwkeurigheid, presteert beter dan pure SWA en komt overeen met volledige attention qua perplexiteit, terwijl het deze in downstream-taken overtreft.
   • Inferentie zonder training: Toegepast op voorgeprogrammeerde Qwen3-8B en Qwen3-30B-A3B modellen, presteert SA consistent beter dan SWA en komt overeen met of overtreft Mixture of Block Attention (MoBA) bij vergelijkbare rekenbudgetten. Het herstelt de kwaliteit van volledige attention aanzienlijk sneller naarmate de venstergrootte toeneemt.

Resultaten

• Pre-training: Het SA + SWA-model bereikte een gemiddelde zero-shot nauwkeurigheid van 35,9 (tegenover 35,1 voor SWA en 34,9 voor Full Attention). Het verbeterde opmerkelijk de prestaties op LAMBADA (waarvoor een brede discourscontext vereist is) terwijl het een lage perplexiteit behield op WikiText. Pure SA leed onder een hogere perplexiteit, wat de noodzaak van lokale coherentie voor taalmodelling bevestigt.
• Inferentie op Qwen3:
  • Bij een effectieve venstergrootte van 128 bereikte SA op Qwen3-8B 70,9% gemiddelde nauwkeurigheid (binnen 1 punt van de full-attention baseline), terwijl SWA achterbleef op 62,2%.
  • Op Qwen3-30B-A3B bereikte SA 73,2% nauwkeurigheid bij een venstergrootte van 64, wat aanzienlijk beter is dan SWA (47,0%) en MoBA (66,3%).
  • SA toonde robuustheid op kennisintensieve taken (MMLU, BoolQ), zelfs bij zeer kleine venstergroottes (bijv. $w=32$), waarbij de SWA-prestaties instortten.
• Efficiëntie: SA behoudt lineaire schaling $O(nw)$. Op een A100 GPU toonde SA snelheidswinsten van 1,5× bij een sequentielengte van 2.048, wat toenam tot 28× bij 32.768 in vergelijking met volledige attention.

Betekenis en Claims

Het artikel stelt dat connectoom-geïnspireerde stochastische routing een praktische primitief is voor het verbeteren van de expressiviteit van efficiënte attention. De auteurs betogen dat globale informatieflow geen dichte all-to-all connectiviteit vereist, maar kan ontstaan uit het samenspel van gestructureerde lokale berekening en sparse, stochastische langeafstandsshortcuts die over de diepte worden opgebouwd.

De betekenis van dit werk ligt in:

1. Eenvoud: SA vereist geen leerbare parameters in het attention-mechanisme zelf en slechts $O(n)$ index-overhead, waardoor het een "plug-and-play" upgrade is voor bestaande SWA-kernels (bijv. in Mistral, Gemma 2, gpt-oss).
2. Theoretische Garantie: Het biedt een bewezen dekkingst diepte van $O(\log_w n)$, een aanzienlijke verbetering ten opzichte van de lineaire dekking van SWA.
3. Complementariteit: Het wordt gepresenteerd als complementair aan bestaande lineaire en sparse benaderingen, en biedt een nieuw ontwerpprincipe dat de kloof overbrugt tussen lokale efficiëntie en globale expressiviteit zonder de complexiteit van handgemaakte sparse patronen of globale tokens.