🤖 machine learning

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

受果蝇连接组学的启发，本文提出随机注意力机制，这是一种随机路由机制，通过增强滑动窗口注意力，在对数深度下实现全局感受野，同时保持线性时间复杂度，在预训练和无训练推理场景中均展现出优越性能。

原作者： Zehao Jin, Yanan Sui

发布于 2026-05-06

📖 1 分钟阅读☕ 轻松阅读

原作者： Zehao Jin, Yanan Sui

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是用简单语言和创意类比对该论文的解读。

核心思想：一种用于 AI 的“随机洗牌”技巧

想象你正在尝试阅读一本非常长的书，但只被允许每次查看 10 个单词的小窗口。许多当前的 AI 模型为了节省能量正是这样工作的；它们使用滑动窗口注意力（SWA）。它们查看彼此相邻的单词，这对于理解局部句子非常有效，但对于记住 50 页之前发生的事情却糟糕透顶。

这篇论文的作者通过观察果蝇的大脑来解决这个问题。尽管果蝇的大脑很小，且其神经元主要只与邻近的神经元相连，但它仍能极其快速地处理信息。这是如何做到的？因为它拥有几条“秘密捷径”，这些捷径随机地连接大脑的遥远部分。

这篇论文提出了一种名为**随机注意力（SA）**的新方法。这就像把那本长书随机打乱页码，阅读你的 10 个单词窗口，然后将页面按顺序放回。

工作原理：“洗牌与阅读”类比

1. “局部窗口”的问题
将标准的 AI 模型想象成一个戴着只允许看到当前单词及其前 10 个单词的遮眼板的人正在读书。

问题所在： 如果这个人需要将第 1 页的一个单词与第 100 页的一个单词联系起来，他做不到。他必须一页接一页地读完整本书，这需要耗费永恒的时间。

2. 果蝇的启发
果蝇的大脑是效率的杰作。它拥有约 13 万个神经元。大多数神经元仅与直接邻居相连（局部聚类）。但是，存在少数随机连接，它们跨越整个大脑（长程捷径）。

结果： 尽管大脑主要是局部的，但信号只需经过几个“跳跃”（约 4 步）就能从一端传到另一端。这是一个“小世界”网络。

3. 解决方案：随机注意力（SA）
作者意识到，他们可以在不构建复杂新大脑的情况下模仿果蝇的捷径。他们通过一个简单的三步技巧来实现：

步骤 1：洗牌。 在 AI 查看其“单词窗口”之前，它随机打乱整个句子的顺序。
步骤 2：阅读。 AI 查看其 10 个单词的小窗口。由于句子已被打乱，这 10 个单词实际上可能来自原始故事的开头、中间和结尾。AI 现在“看到”了文本的遥远部分，仿佛它们是邻居一样。
步骤 3：还原。 在 AI 处理完信息后，它将单词放回原始顺序，使句子再次变得通顺。

魔力所在： 通过在 AI 的层级中反复进行这种“洗牌与阅读”的技巧，模型可以非常快速地“触及”文本的每一个部分。这就像你想在一个巨大的体育场里见到每个人；与其逐排行走（缓慢），不如每隔几秒钟随机瞬移到几个不同的区域。你可以在几分钟内覆盖整个体育场。

“两全其美”的组合

论文还指出，仅仅打乱所有内容并不完美。你仍然需要理解单词的局部流动（语法、句子结构）。因此，他们创建了一个门控 SA + SWA系统。

将这想象成一条双车道公路：

车道 1（SWA）： 标准车道，汽车（单词）停留在其局部区域。这保持了语法和局部含义的完美。
车道 2（SA）： “瞬移”车道，汽车随机跳跃到公路的遥远部分以获取重要的上下文。
闸门： 一个智能交通控制器（一个学习到的门）决定每个单词从局部车道和瞬移车道分别获取多少信息。

实验结果

研究人员通过两种方式测试了这一点：

从头构建新的 AI： 他们使用这种方法训练了一个新的语言模型。结果如何？使用“洗牌 + 局部”组合的模型是最聪明的。它既很好地理解了局部句子，又比仅使用一种方法的模型更好地记住了长程上下文。
升级现有 AI（Qwen3）： 他们取了一个强大的预训练 AI（Qwen3），简单地将其注意力机制替换为这种新的“洗牌”方法，而无需重新训练它。
- 结果： 在查看长文本时，升级后的 AI 表现远优于原始版本。它恢复了全注意力模型（一次性查看所有内容）的“智能”，但运行速度更快，使用的内存更少。

为什么这很重要

该论文声称这是一种“即插即用”的升级。它不需要改变 AI 的大脑结构或增加数百万个新参数。它只是添加了一个简单的随机洗牌步骤。

速度： 它保持了快速“局部窗口”方法的速度。
智能： 它获得了缓慢的“查看所有内容”方法的智能。
效率： 它通过模仿自然界最高效的网络——果蝇的大脑——来实现这一目标。

简而言之，这篇论文证明，你不需要构建一个庞大且昂贵的网络来连接一切。有时，你只需要稍微摇动一下局面即可。

技术摘要：随机注意力

问题陈述

高效的注意力机制对于扩展 Transformer 模型至关重要，然而现有方法在计算成本与感受野覆盖范围之间面临权衡。**滑动窗口注意力（SWA）**通过将令牌限制在大小为 $w$ 的局部窗口内，将每层的复杂度降低至 $O(nw) $。然而，这种确定性的局部性将有效感受野限制为线性增长（$ \ell $层后为$ \ell w $），当$ w \ll n $时，序列的很大一部分将无法到达。虽然存在全局令牌、手工设计的稀疏模式或块级路由等补救措施，但它们往往引入了架构复杂性，或无法提供可证明的覆盖保证。相反，全注意力提供了全局连接性，但会产生不可接受的$ O(n^2)$ 成本。

本文的灵感来源于果蝇（Drosophila melanogaster）的神经连接组，这是一个包含约 130,000 个神经元的网络，连接概率仅为 0.02%。尽管存在这种稀疏性，该网络仍实现了约 4.4 跳的平均路径长度。这种“小世界”拓扑结构源于密集局部聚类与广泛分布的长程连接之间的相互作用，后者充当了随机捷径。作者提出，高效的全局信息流可以从结构化的局部计算中涌现，并辅以分布式、随机的长程连接，而无需密集的全对全连接。

方法论：随机注意力（SA）

作者提出了随机注意力（SA），这是一种无参数的 SWA 增强方案，模仿了连接组的路由策略。

核心机制

SA 通过以下三个阶段的过程（在标准 SWA 计算之前和之后应用），将固定的局部窗口转化为随机全局窗口：

置换（Permute）： 为每一层独立抽取一个随机置换 $\sigma \sim \text{Uniform}(S_n)$ 。令牌序列表示（ $Q, K, V$ ）根据 $\sigma$ 进行重排。
窗口注意力（Windowed Attention）： 在置换后的空间中应用标准 SWA。由于置换是随机的，大小为 $w$ 的固定窗口现在跨越了原始序列的一个随机子集。
撤销置换（Undo Permute）： 通过对输出应用 $\sigma^{-1}$ 恢复原始顺序。

在原始令牌空间中，位置 $i$ 关注的是一个期望大小为 $w$ 的随机邻域 $\tilde{N}^\sigma_w(i)$ ，该邻域在整个序列中均匀分布，与原始距离 $|i-j|$ 无关。这为任意令牌对创建了约 $w/n$ 的成对连接概率，类似于连接组中分布式的长程捷径。

理论属性

感受野增长： 虽然 SWA 的感受野呈线性增长（完全覆盖需要 $O(n/w)$ 层），但 SA 在各层中独立采样的置换产生了指数级增长的感受野。本文证明，完全序列覆盖可在 $O(\log_w n)$ 层内实现。
复杂度： 置换和逆置换操作是 $O(n)$ 的索引重排。注意力计算仍保持 $O(nw)$。因此，SA 保持了与标准 SWA 相同的每层 $O(nw)$ 预算。
因果约束： 在自回归设置中，因果掩码在置换后的空间中应用，确保令牌仅关注因果上可访问的前驱，同时仍能启用随机长程连接。

门控 SA + SWA 组合

为了重现“小世界”机制（高局部聚类 + 短路径），作者提出了一种结合 SA 和 SWA 的双路径架构：
$Y = g_{sa} \odot Y_{sa} + g_{swa} \odot Y_{swa}$
其中 $Y_{sa}$ 和 $Y_{swa}$ 分别来自随机路径和滑动窗口路径的输出。门控权重（ $g_{sa}, g_{swa}$ ）通过独立的 Sigmoid 投影计算得出，使模型能够学习平衡局部连贯性（来自 SWA）和全局覆盖（来自 SA）。这种组合增加了微不足道的参数量，并保持 $O(nw)$ 的复杂度。

主要贡献

随机注意力（SA）： 一种即插即用、无参数的 SWA 增强方案，通过应用随机置换，在相同的 $O(nw) $计算预算内实现指数级感受野增长（$ O(\log_w n)$）。
门控 SA + SWA： 一种双路径架构，在理论和实证上重现了生物连接组中观察到的小世界机制，平衡了局部聚类与随机捷径。
实证验证：
- 预训练： 在从头训练的 3.6 亿参数模型上，门控 SA + SWA 组合实现了最佳的平均零样本准确率，优于纯 SWA，在困惑度上与全注意力相当，并在下游任务中超越全注意力。
- 免训练推理： 应用于预训练的 Qwen3-8B 和 Qwen3-30B-A3B 模型，SA 始终优于 SWA，并在可比计算预算下达到或超过**块注意力混合（MoBA）**的性能。随着窗口大小的增加，SA 能显著更快地恢复全注意力的质量。

结果

预训练： SA + SWA 模型实现了 35.9 的平均零样本准确率（SWA 为 35.1，全注意力为 34.9）。它显著提高了对LAMBADA（需要广泛的语篇上下文）任务的性能，同时在WikiText上保持了低困惑度。纯 SA 表现出更高的困惑度，证实了语言建模中局部连贯性的必要性。
Qwen3 上的推理：
- 在有效窗口大小为 128 时，Qwen3-8B 上的 SA 达到了 70.9% 的平均准确率（在全注意力基线 1 个百分点以内），而 SWA 仅为 62.2%。
- 在 Qwen3-30B-A3B 上，SA 在窗口大小为 64 时达到了 73.2% 的准确率，显著优于 SWA（47.0%）和 MoBA（66.3%）。
- 即使在非常小的窗口大小（如 $w=32$ ）下，SA 在知识密集型任务（MMLU, BoolQ）上也表现出鲁棒性，而 SWA 的性能则在此时崩溃。
效率： SA 保持线性扩展 $O(nw)$。在 A100 GPU 上，SA 在序列长度为 2,048 时显示出 1.5 倍 的加速，而在 32,768 时加速达到 28 倍（与全注意力相比）。

意义与主张

本文主张，受连接组启发的随机路由是提升高效注意力表达力的实用原语。作者认为，全局信息流不需要密集的全对全连接，而是可以从结构化的局部计算与随深度累积的稀疏、随机长程捷径之间的相互作用中涌现。

这项工作的意义在于：

简洁性： SA 在注意力机制本身中不需要可学习参数，仅需 $O(n)$ 的索引开销，使其成为现有 SWA 内核（如 Mistral、Gemma 2、gpt-oss 中）的“即插即用”升级。
理论保证： 它提供了可证明的 $O(\log_w n)$ 覆盖深度，这是对 SWA 线性覆盖的重大改进。
互补性： 它被呈现为对现有线性和稀疏方法的补充，提供了一种新的设计原则，在不引入手工设计的稀疏模式或全局令牌复杂性的情况下，弥合了局部效率与全局表达力之间的差距。