Technical Summary: Stochastic Attention

Problem Statement

Efficient attention mechanisms are critical for scaling Transformer models, yet existing approaches face a trade-off between computational cost and receptive field coverage. Sliding-Window Attention (SWA) reduces complexity to $O(nw)$ per layer by restricting tokens to a local window of size $w$. However, this deterministic locality limits the effective receptive field to linear growth ($\ell w$ after $\ell$ layers), leaving large portions of the sequence unreachable when $w \ll n$. While remedies such as global tokens, hand-crafted sparse patterns, or block-level routing exist, they often introduce architectural complexity or fail to provide provable coverage guarantees. Conversely, full attention offers global connectivity but incurs prohibitive $O(n^2)$ costs.

The paper draws inspiration from the fruit fly (Drosophila melanogaster) connectome, a network of ~130,000 neurons with a connection probability of only 0.02%. Despite this sparsity, the network achieves a short average path length of ~4.4 hops. This "small-world" topology arises from the interplay of dense local clustering and broadly distributed long-range connections that function as stochastic shortcuts. The authors posit that efficient global information flow can emerge from structured local computation augmented by distributed, stochastic long-range connections, rather than requiring dense all-to-all connectivity.

Methodology: Stochastic Attention (SA)

The authors propose Stochastic Attention (SA), a parameter-free enhancement for SWA that mimics the connectome's routing strategy.

Core Mechanism

SA transforms a fixed local window into a stochastic global one through a three-stage process applied before and after the standard SWA computation:

1. Permute: A random permutation $\sigma \sim \text{Uniform}(S_n)$ is drawn independently for each layer. The token sequence representations ($Q, K, V$) are rearranged according to $\sigma$.
2. Windowed Attention: Standard SWA is applied in the permuted space. Because the permutation is random, the fixed window of size $w$ now spans a random subset of the original sequence.
3. Undo Permute: The original order is restored by applying $\sigma^{-1}$ to the output.

In the original token space, position $i$ attends to a random neighborhood $\tilde{N}^\sigma_w(i)$ of expected size $w$, uniformly distributed across the full sequence regardless of the original distance $|i-j|$. This creates a pairwise connection probability of approximately $w/n$ for any token pair, analogous to the distributed long-range shortcuts in the connectome.

Theoretical Properties

• Receptive Field Growth: While SWA's receptive field grows linearly ($O(n/w)$ layers for full coverage), SA's independently sampled permutations across layers yield exponentially growing receptive fields. The paper proves that full sequence coverage is achieved in $O(\log_w n)$ layers.
• Complexity: The permutation and inverse permutation operations are $O(n)$ index rearrangements. The attention computation remains $O(nw)$. Thus, SA maintains the same $O(nw)$ per-layer budget as standard SWA.
• Causal Constraints: In autoregressive settings, the causal mask is applied in the permuted space, ensuring tokens only attend to causally accessible predecessors while still enabling stochastic long-range connections.

Gated SA + SWA Combination

To reproduce the "small-world" regime (high local clustering + short paths), the authors propose a dual-path architecture combining SA and SWA:
$$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$$
where $Y_{sa}$ and $Y_{swa}$ are outputs from the stochastic and sliding-window paths, respectively. The gating weights ($g_{sa}, g_{swa}$) are computed via independent sigmoid projections, allowing the model to learn to balance local coherence (from SWA) and global coverage (from SA). This combination adds negligible parameters and maintains $O(nw)$ complexity.

Key Contributions

1. Stochastic Attention (SA): A drop-in, parameter-free enhancement for SWA that applies random permutations to achieve exponential receptive field growth ($O(\log_w n)$) within the same $O(nw)$ computational budget.
2. Gated SA + SWA: A dual-path architecture that theoretically and empirically reproduces the small-world regime observed in biological connectomes, balancing local clustering and stochastic shortcuts.
3. Empirical Validation:
   • Pre-training: On 360M-parameter models trained from scratch, the gated SA + SWA combination achieves the best average zero-shot accuracy, outperforming pure SWA and matching full attention in perplexity while surpassing it in downstream tasks.
   • Training-Free Inference: Applied to pretrained Qwen3-8B and Qwen3-30B-A3B models, SA consistently outperforms SWA and matches or exceeds Mixture of Block Attention (MoBA) at comparable compute budgets. It recovers full-attention quality significantly faster as window size increases.

Results

• Pre-training: The SA + SWA model achieved an average zero-shot accuracy of 35.9 (vs. 35.1 for SWA and 34.9 for Full Attention). It notably improved performance on LAMBADA (requiring broad discourse context) while maintaining low perplexity on WikiText. Pure SA suffered higher perplexity, confirming the necessity of local coherence for language modeling.
• Inference on Qwen3:
  • At an effective window size of 128, SA on Qwen3-8B reached 70.9% average accuracy (within 1 point of the full-attention baseline), whereas SWA lagged at 62.2%.
  • On Qwen3-30B-A3B, SA reached 73.2% accuracy at a window size of 64, significantly outperforming SWA (47.0%) and MoBA (66.3%).
  • SA demonstrated robustness on knowledge-intensive tasks (MMLU, BoolQ) even at very small window sizes (e.g., $w=32$), where SWA performance collapsed.
• Efficiency: SA maintains linear scaling $O(nw)$. On an A100 GPU, SA showed speedups of 1.5× at sequence length 2,048, increasing to 28× at 32,768 compared to full attention.

Significance and Claims

The paper claims that connectome-inspired stochastic routing is a practical primitive for improving the expressivity of efficient attention. The authors argue that global information flow does not require dense all-to-all connectivity but can emerge from the interplay of structured local computation and sparse, stochastic long-range shortcuts accumulated over depth.

The significance of this work lies in:

1. Simplicity: SA requires no learnable parameters in the attention mechanism itself and only $O(n)$ index overhead, making it a "drop-in" upgrade for existing SWA kernels (e.g., in Mistral, Gemma 2, gpt-oss).
2. Theoretical Guarantee: It provides provable $O(\log_w n)$ coverage depth, a significant improvement over the linear coverage of SWA.
3. Complementarity: It is presented as complementary to existing linear and sparse approaches, offering a new design principle that bridges the gap between local efficiency and global expressivity without the complexity of hand-crafted sparse patterns or global tokens.