Technische Zusammenfassung: Stochastische Aufmerksamkeit

Problemstellung

Effiziente Aufmerksamkeitsmechanismen sind entscheidend für die Skalierung von Transformer-Modellen, doch bestehende Ansätze stehen vor einem Zielkonflikt zwischen Rechenaufwand und Abdeckung des rezeptiven Feldes. Sliding-Window Attention (SWA) reduziert die Komplexität pro Schicht auf $O(nw)$, indem sie Token auf ein lokales Fenster der Größe $w$ beschränkt. Diese deterministische Lokalität begrenzt das effektive rezeptive Feld jedoch auf ein lineares Wachstum ($\ell w$ nach $\ell$ Schichten), sodass große Teile der Sequenz unerreichbar bleiben, wenn $w \ll n$. Zwar existieren Abhilfemaßnahmen wie globale Token, handgefertigte spärliche Muster oder Routing auf Blockebene, doch führen diese oft zu architektonischer Komplexität oder bieten keine nachweisbaren Abdeckungsgarantien. Umgekehrt bietet vollständige Aufmerksamkeit globale Konnektivität, verursacht jedoch prohibitiv hohe Kosten von $O(n^2)$.

Die Arbeit zieht Inspiration aus dem Connectom der Fruchtfliege (Drosophila melanogaster), einem Netzwerk von etwa 130.000 Neuronen mit einer Verbindungswahrscheinlichkeit von lediglich 0,02 %. Trotz dieser Sparsamkeit erreicht das Netzwerk eine durchschnittliche Pfadlänge von nur etwa 4,4 Sprüngen. Diese „Small-World"-Topologie entsteht aus dem Zusammenspiel dichter lokaler Clusterbildung und breit verteilter Langzeitverbindungen, die als stochastische Abkürzungen fungieren. Die Autoren gehen davon aus, dass ein effizienter globaler Informationsfluss aus strukturierter lokaler Berechnung entstehen kann, die durch verteilte, stochastische Langzeitverbindungen ergänzt wird, anstatt eine dichte All-zu-All-Konnektivität zu erfordern.

Methodik: Stochastische Aufmerksamkeit (SA)

Die Autoren schlagen Stochastische Aufmerksamkeit (SA) vor, eine parameterfreie Erweiterung für SWA, die die Routing-Strategie des Connectoms nachahmt.

Kernmechanismus

SA verwandelt ein festes lokales Fenster in ein stochastisches globales Fenster durch einen dreistufigen Prozess, der vor und nach der Standard-SWA-Berechnung angewendet wird:

1. Permutieren: Eine zufällige Permutation $\sigma \sim \text{Uniform}(S_n)$ wird unabhängig für jede Schicht gezogen. Die Token-Sequenzdarstellungen ($Q, K, V$) werden gemäß $\sigma$ neu angeordnet.
2. Fensteraufmerksamkeit: Standard-SWA wird im permutierten Raum angewendet. Da die Permutation zufällig ist, erstreckt sich das feste Fenster der Größe $w$ nun über eine zufällige Teilmenge der ursprünglichen Sequenz.
3. Permutation rückgängig machen: Die ursprüngliche Reihenfolge wird wiederhergestellt, indem $\sigma^{-1}$ auf die Ausgabe angewendet wird.

Im ursprünglichen Token-Raum attends Position $i$ an eine zufällige Nachbarschaft $\tilde{N}^\sigma_w(i)$ der erwarteten Größe $w$, die über die gesamte Sequenz gleichmäßig verteilt ist, unabhängig von der ursprünglichen Distanz $|i-j|$. Dies erzeugt eine paarweise Verbindungswahrscheinlichkeit von ungefähr $w/n$ für jedes Tokenpaar, analog zu den verteilten Langzeit-Abkürzungen im Connectom.

Theoretische Eigenschaften

• Wachstum des rezeptiven Feldes: Während das rezeptive Feld von SWA linear wächst ($O(n/w)$ Schichten für vollständige Abdeckung), führt SA durch unabhängig über die Schichten gesampelte Permutationen zu einem exponentiell wachsenden rezeptiven Feld. Die Arbeit beweist, dass eine vollständige Sequenzabdeckung in $O(\log_w n)$ Schichten erreicht wird.
• Komplexität: Die Permutations- und Inversions-Permutations-Operationen sind $O(n)$-Index-Umsortierungen. Die Aufmerksamkeitsberechnung bleibt bei $O(nw)$. Somit behält SA dasselbe $O(nw)$-Budget pro Schicht wie die Standard-SWA bei.
• Kausale Einschränkungen: In autoregressiven Settings wird die kausale Maske im permutierten Raum angewendet, wodurch sichergestellt wird, dass Token nur kausal zugängliche Vorgänger attendieren, während gleichzeitig stochastische Langzeitverbindungen ermöglicht werden.

Gated SA + SWA Kombination

Um den „Small-World"-Regime (hohe lokale Clusterbildung + kurze Pfade) nachzubilden, schlagen die Autoren eine Dual-Pfad-Architektur vor, die SA und SWA kombiniert:
$$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$$
wobei $Y_{sa}$ und $Y_{swa}$ die Ausgaben aus dem stochastischen bzw. dem Sliding-Window-Pfad sind. Die Gate-Gewichte ($g_{sa}, g_{swa}$) werden über unabhängige Sigmoid-Projektionen berechnet, wodurch das Modell lernen kann, lokale Kohärenz (von SWA) und globale Abdeckung (von SA) auszubalancieren. Diese Kombination fügt vernachlässigbare Parameter hinzu und behält die $O(nw)$-Komplexität bei.

Hauptbeiträge

1. Stochastische Aufmerksamkeit (SA): Eine Plug-in-, parameterfreie Erweiterung für SWA, die zufällige Permutationen anwendet, um innerhalb desselben $O(nw)$-Rechenbudgets ein exponentielles Wachstum des rezeptiven Feldes ($O(\log_w n)$) zu erreichen.
2. Gated SA + SWA: Eine Dual-Pfad-Architektur, die den in biologischen Connectomen beobachteten Small-World-Regime theoretisch und empirisch nachbildet und lokale Clusterbildung mit stochastischen Abkürzungen ausbalanciert.
3. Empirische Validierung:
   • Pre-Training: Bei 360-Millionen-Parameter-Modellen, die von Grund auf trainiert wurden, erreicht die Gated SA + SWA-Kombination die beste durchschnittliche Zero-Shot-Genauigkeit und übertrifft reine SWA, wobei sie in der Perplexität mit vollständiger Aufmerksamkeit gleichzieht, sie jedoch in Downstream-Aufgaben übertrifft.
   • Training-Free Inference: Auf vortrainierten Qwen3-8B- und Qwen3-30B-A3B-Modellen angewendet, übertrifft SA SWA konsistent und erreicht oder übertrifft Mixture of Block Attention (MoBA) bei vergleichbaren Rechenbudgets. Sie stellt die Qualität der vollständigen Aufmerksamkeit bei zunehmender Fenstergröße deutlich schneller wieder her.

Ergebnisse

• Pre-Training: Das SA + SWA-Modell erreichte eine durchschnittliche Zero-Shot-Genauigkeit von 35,9 (im Vergleich zu 35,1 für SWA und 34,9 für vollständige Aufmerksamkeit). Es verbesserte die Leistung auf LAMBADA (erfordert breiten Diskurskontext) erheblich, während es eine niedrige Perplexität auf WikiText beibehielt. Reine SA litt unter höherer Perplexität, was die Notwendigkeit lokaler Kohärenz für Sprachmodellierung bestätigt.
• Inference auf Qwen3:
  • Bei einer effektiven Fenstergröße von 128 erreichte SA auf Qwen3-8B eine durchschnittliche Genauigkeit von 70,9 % (innerhalb eines Punkts der Vollständigkeits-Aufmerksamkeits-Baseline), während SWA mit 62,2 % zurücklag.
  • Auf Qwen3-30B-A3B erreichte SA bei einer Fenstergröße von 64 eine Genauigkeit von 73,2 % und übertraf SWA (47,0 %) und MoBA (66,3 %) erheblich.
  • SA zeigte Robustheit bei wissensintensiven Aufgaben (MMLU, BoolQ), selbst bei sehr kleinen Fenstergrößen (z. B. $w=32$), bei denen die SWA-Leistung kollabierte.
• Effizienz: SA behält die lineare Skalierung $O(nw)$ bei. Auf einer A100-GPU zeigte SA bei einer Sequenzlänge von 2.048 Geschwindigkeitsvorteile von 1,5×, die sich bei 32.768 auf 28× im Vergleich zur vollständigen Aufmerksamkeit steigerten.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass connectom-inspiriertes stochastisches Routing eine praktische Primitive zur Verbesserung der Ausdruckskraft effizienter Aufmerksamkeit ist. Die Autoren argumentieren, dass globaler Informationsfluss keine dichte All-zu-All-Konnektivität erfordert, sondern aus dem Zusammenspiel strukturierter lokaler Berechnung und spärlicher, stochastischer Langzeit-Abkürzungen entstehen kann, die über die Tiefe hinweg akkumuliert werden.

Die Bedeutung dieser Arbeit liegt in:

1. Einfachheit: SA erfordert keine lernbaren Parameter im Aufmerksamkeitsmechanismus selbst und nur einen $O(n)$-Index-Overhead, was sie zu einem „Plug-in"-Upgrade für bestehende SWA-Kernels (z. B. in Mistral, Gemma 2, gpt-oss) macht.
2. Theoretische Garantie: Sie bietet eine nachweisbare Abdeckungstiefe von $O(\log_w n)$, eine erhebliche Verbesserung gegenüber der linearen Abdeckung von SWA.
3. Komplementarität: Sie wird als komplementär zu bestehenden linearen und spärlichen Ansätzen präsentiert und bietet ein neues Designprinzip, das die Lücke zwischen lokaler Effizienz und globaler Ausdruckskraft schließt, ohne die Komplexität handgefertigter spärlicher Muster oder globaler Token.