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

Switching Linear Attention
Sep 30, 2026Hyun Dong Lee, Xavier Gonzalez, Nicolas Zucchet +3Efficient Attention Linear Attention
Stanford University · Unconventional AI. Work done at Stanford University.
Why Attend to Everything? Focus is the Key
Mar 12, 2026Hengshuai Yao, Xing Chen, Ahmed Murtadha +8Attention Mechanisms Efficient Attention
Sapient Intelligence · University of Alberta · Tsinghua University
SMat-Attention: Structured Long-Context Sequence Modeling
Sep 28, 2026Emile Anand, Abdullah Ateyeh, Archer Wang +1Long-Context Modeling Linear Attention
Massachusetts Institute of Technology

Sep 30, 2026J/K move · Enter open · S save

Switching Linear Attention

Hyun Dong Lee, Xavier Gonzalez, Nicolas Zucchet, E. Kelly Buchanan+2

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

Efficient Attention Linear Attention

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Switching Linear Attention

Why Attend to Everything? Focus is the Key

SMat-Attention: Structured Long-Context Sequence Modeling