Stochastic Attention: Connectome-Inspired Randomized Routing for Expressive Linear-Time Attention

Zehao Jin; Yanan Sui

cs.CLApr 1, 2026

Stochastic Attention: Connectome-Inspired Randomized Routing for Expressive Linear-Time Attention

Authors: Zehao Jin, Yanan Sui

Organizations: Tsinghua University

Abstract

The whole-brain connectome of a fruit fly comprises over 130K neurons connected with a probability of merely 0.02%, yet achieves an average shortest path of only 4.4 hops. Despite being highly structured at the circuit level, the network's long-range connections are broadly distributed across brain regions, functioning as stochastic shortcuts that enable efficient global communication. Inspired by this observation, we propose Stochastic Attention (SA), a drop-in enhancement for sliding-window attention (SWA) that applies a random permutation to the token sequence before windowed attention and restores the original order afterward. This transforms the fixed local window into a stochastic global one within the same $O(nw)$ per-layer budget. Through depth, independently sampled permutations yield exponentially growing receptive fields, achieving full sequence coverage in $O(\log_w n)$ layers versus $O(n/w)$ for SWA. We validate SA in two settings: pre-training language models from scratch, where a gated SA + SWA combination achieves the best average zero-shot accuracy, and training-free inference on Qwen3-8B and Qwen3-30B-A3B, where SA consistently outperforms SWA and matches or exceeds Mixture of Block Attention at comparable compute budgets. These results suggest that connectome-inspired stochastic routing is a practical primitive for improving the expressivity of efficient attention, complementary to existing linear and sparse approaches.

Figures & tables

Figure 1: Overview of Stochastic Attention (SA). (a) A standard SWA Transformer layer. (b) The fruit fly whole-brain connectome: the adjacency matrix, shown after Reverse Cuthill–McKee reordering to expose block structure, lacks clear diagonal blocks, indicating that connections are broadly distributed across brain regions rather than confined to local modules. (c) An SA layer: token sequences are randomly permuted before windowed attention and restored afterward, producing stochastic long-range shortcuts analogous to the cross-regional connections in (b).

Figure 2: Left: Receptive field coverage as a function of depth ( $n{=}2048$ , $w{=}32$ ). SA achieves full sequence coverage in $O(\log_{w}n)$ layers via exponential growth, while SWA requires $O(n/w)$ layers with linear growth. Right: Computational cost scaling with sequence length ( $w{=}256$ ). Both SA and SWA maintain $O(nw)$ linear scaling, while full attention grows quadratically.

Model	Wiki. $\downarrow$	LMB-o $\downarrow$	LMB-s $\downarrow$	LMB-o	LMB-s	PIQA	Hella.	Wino.	ARC-e	Avg.
Full Attention	51.34	185.3	469.2	19.5	15.4	59.9	27.2	51.3	35.8	34.9
SWA ( $w{=}256$ )	57.05	156.1	370.6	21.3	17.0	59.6	27.6	48.9	36.0	35.1
SA ( $w{=}256$ )	75.83	260.1	785.9	20.1	14.1	59.3	26.4	51.9	34.2	34.3
SA + SWA ( $w{=}256$ )	51.98	131.7	371.6	22.8	17.6	59.6	27.5	50.5	37.5	35.9

Table 1: Zero-shot evaluation of language models trained on SlimPajama (15B tokens). All models share identical training setup, differing only in the attention mechanism. Wiki. and LMB ppl report perplexity (

\downarrow

), all others report accuracy (

\uparrow

). Best in bold , second best underlined .

Figure 3: Attention weight visualization (Layer 11, Head 0) on a 27-token sequence with window size $w{=}8$ . Gray regions are masked (structurally invisible). Blue intensity indicates attention weight. Full Attention exhibits the complete lower-triangular pattern. SWA shows a strict diagonal band with all out-of-window positions masked. Stochastic Attention introduces scattered non-zero entries beyond the diagonal band. These are distant tokens that became local neighbors after random permutation, enabling direct long-range information flow within the same $O(nw)$ budget. SA + SWA combines both patterns: the SWA path provides the dense diagonal band for local coherence, while the SA path adds stochastic long-range connections, with the learned gate adaptively balancing the two.

Figure 4: Average accuracy across 7 benchmarks as a function of effective window size for Qwen3-8B (left) and Qwen3-30B-A3B (right). Stochastic Attention (red) recovers the full-attention baseline (dashed gray) most rapidly as window size increases, consistently outpacing SWA (blue) and matching or exceeding MoBA (green) at comparable compute budgets.

Figure 5: Per-task accuracy vs. window size on Qwen3-8B for four representative benchmarks.

Figure 6: Per-task accuracy vs. window size on Qwen3-30B-A3B for four representative benchmarks.

Seq. Length	SA ( $w{=}256$ )	Full Attn	Speedup
2,048	5.4 ms	8.0 ms	1.5 $\times$
4,096	7.9 ms	27.2 ms	3.5 $\times$
8,192	15.2 ms	99.7 ms	6.6 $\times$
16,384	29.6 ms	379.9 ms	12.8 $\times$
32,768	52.8 ms	1,477 ms	28.0 $\times$

Table 2: Attention layer latency (ms, forward+backward) on A100 80GB. SA uses compiled FlexAttention ( Dong et al., 2024 ) with

w{=}256

. Full attention sets

w{=}L

. Measured with

B{=}16

H{=}16

d_{h}{=}64

, bf16.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Mechanism	Cost / Layer	Receptive Field	Coverage Depth
Full Attention	$O(n^{2})$	$n$ (layer 1)	$1$
Sliding Window	$O(nw)$	$\ell w$	$O(n/w)$
SA (ours)	$O(nw)$	$\Omega((w/4)^{\ell})$	$\bm{O(\log_{w}n)}$
SA + SWA	$2\cdot O(nw)$	$\Omega((w/4)^{\ell})$	$\bm{O(\log_{w}n)}$

Appendix

Table 3: Comparison of attention mechanisms.

n

: sequence length,

w

: window size,

\ell

: number of layers.

Figure 7: Per-task accuracy vs. effective window size for Qwen3-8B across all evaluated benchmarks.

Figure 8: Per-task accuracy vs. effective window size for Qwen3-30B-A3B across all evaluated benchmarks.

$w_{\text{eff}}$	Mode	Hella.	MMLU	ARC-C	BoolQ	LMB	ARC-E	HuEval	Avg.
$\infty$	Full	74.9	73.0	56.4	86.6	64.6	80.9	64.0	71.5
32	SWA	47.7	29.0	44.3	59.7	15.1	71.1	0.0	38.1
32	Stochastic	53.5	44.4	52.5	59.3	33.2	79.0	26.8	49.8
64	SWA	55.6	46.0	55.6	63.8	50.2	80.2	3.7	50.7
	Stochastic	72.4	63.0	56.1	77.1	62.9	80.9	52.4	66.4
	MoBA ( $k{=}2$ )	67.1	56.4	55.6	75.7	44.7	80.5	60.4	62.9

Appendix

Table 4: Training-free inference on Qwen3-8B. We report accuracy on 7 benchmarks at selected window sizes. Best result per column among efficient methods in bold , underlined denotes second best.

$w_{\text{eff}}$	Mode	Hella.	MMLU	ARC-C	BoolQ	LMB	ARC-E	HuEval	Avg.
$\infty$	Full	79.8	80.1	62.3	88.7	71.5	83.8	75.6	77.4
32	SWA	52.2	34.9	53.3	74.2	14.9	77.6	1.2	44.0
32	Stochastic	64.0	52.0	60.8	62.3	41.4	82.0	39.6	57.4
64	SWA	51.3	42.6	61.1	64.7	23.6	82.9	3.0	47.0
	Stochastic	78.0	69.6	62.7	82.5	69.7	83.9	65.9	73.2
	MoBA ( $k{=}2$ )	71.2	64.4	61.0	77.7	44.7	83.3	61.6	66.3

Appendix

Table 5: Training-free inference on Qwen3-30B-A3B. Same setup as Table 4 .

Explore similar work

Sep 30, 2026

Switching Linear Attention

Designing expressive sequence layers with efficient inference remains a central challenge in modern machine learning. Standard softmax attention achieves excellent sequence modeling performance through rich nonlinear token interactions, but it requires a key-value cache that grows linearly with sequence length, limiting its scalability. Linear attention enables efficient recurrent computation with a constant memory footprint, yet its reduced expressivity often yields inferior modeling performance. We introduce Switching Linear Attention (SwiLA), a novel sequence layer that bridges this gap by enhancing representational capacity while retaining the fixed-size recurrent state of linear attention. We derive the SwiLA recurrence from the test-time regression framework, casting the state update rule as online expectation-maximization in a mixture of linear regressions model. At test time, each output dimension dynamically selects among multiple linear attention components based on the input. Across associative recall, in-context language learning, and language modeling benchmarks, SwiLA shows strong performance and narrows the gap to softmax attention, even surpassing it in several settings.

Hyun Dong Lee, Xavier Gonzalez, Nicolas Zucchet +3

Stanford University · Unconventional AI. Work done at Stanford University.

Efficient Attention Linear Attention

Ask PDF

Mar 12, 2026

Why Attend to Everything? Focus is the Key

Standard attention scales quadratically with sequence length. Efficient attention methods reduce this O(n^2) cost, but when retrofitted into pretrained models, they often degrade perplexity, downstream accuracy, or both. We introduce Focus, a method that learns which token pairs matter. Focus adds a small set of learnable centroids--as few as 148K parameters per layer--that act as gates: only token pairs belonging to the same centroid group attend to each other over long ranges. Focus is composable: it can be added to any pretrained model by training only the centroids while keeping all original weights frozen. Experiments show that composing Focus onto pretrained models yields zero degradation on downstream benchmarks across model sizes from 124M to 70B parameters and five attention architectures. Surprisingly, sparse Focus attention outperforms full attention at 124M scale (30.3 vs. 31.4 perplexity) and matches full attention when trained from scratch at 7B scale (13.82 vs. 13.89). Focus is also fast: top-k group membership gives a 2x speedup with better quality than the original pretrained model. Using our FlashAttention decomposition, Focus achieves an 8.6x speedup at 1M tokens without custom kernels.

Hengshuai Yao, Xing Chen, Ahmed Murtadha +8

Sapient Intelligence · University of Alberta · Tsinghua University

Attention Mechanisms Efficient Attention Sparse Attention Long-Context Modeling

Ask PDF

Sep 28, 2026

SMat-Attention: Structured Long-Context Sequence Modeling

Long-context sequence models face a fundamental tradeoff: softmax attention uses flexible token-level interactions at quadratic cost, whereas linear attention obtains linear-time training and constant-time decoding by compressing history into a fixed-size state. In this work, we ask whether we can connect these regimes through a tunable notion of structure. To this end, we introduce Structured Matrix Attention (SMat-Attention) via a family of causal masks with structured long-range routing whose row supports have VC-dimension

d

. In our construction,

d=1

recovers the standard causal mask, and increasing

d

permits richer subset-routing patterns. We give chunkwise forward and backward algorithms to enable hardware-efficiency. For sequences of length

T

, the hard-routing construction takes

O(T^{2-3/d}+T)

work, despite the mask being dense, for our prescribed family. In fixed-horizon streaming, decoding after the distant prefix takes constant time per token using

O(T^{1-1/d})

cached states. SMat-Attention therefore makes VC-dimension an explicit knob governing access-pattern complexity, prefill cost, and decoding memory. Empirically, subset-routing and rule-assisted multi-key retrieval experiments illustrate the masks' routing expressiveness. Extensions to Mamba-2 and Gated DeltaNet using learned routing with top-

k

query reads retain subquadratic prefill, improve recall accuracy over the backbones in several settings, and achieve comparable small-scale language-modeling performance.

Emile Anand, Abdullah Ateyeh, Archer Wang +1

Massachusetts Institute of Technology

Long-Context Modeling Linear Attention Attention Mechanisms

Ask PDF