超线性多步注意力

在本文中，我们提出了 \textbf{Superlinear attention}，一种完全可训练的多步注意力架构，它在实现长序列次二次复杂度的同时保留了 \textbf{随机上下文访问}（又称结构性非排斥）：没有任何符合条件的位置会在结构上被排除在注意力选择之外。Superlinear attention 将标准的因果自注意力重新表述为一个具有 $N$ 步的多步搜索问题，产生 $O(L^{1+\frac{1}{N}})$ 的总体复杂度。为了说明该架构，我们提出了一个基线 $N=2$ 实现，其在算法上类似于标准的跳跃搜索。在这个 $O(L^{3/2})$ 的实例中，第一步执行 $O(L^{3/2})$ 的跨度搜索以选择序列的相关跨度，第二步应用 $O(L^{3/2})$ 的跨度注意力（限制在所选跨度内的标准注意力）。在为提升鲁棒性而放大的 $O(L^{1.54})$ 配置中，我们在单块 B200 GPU 上修改后的 30B 混合 MoE 模型实现中，在 1M 上下文长度下实现了 114 tokens/sec 的平均解码吞吐量，在 10M 上下文长度下实现了 80 tokens/sec 的平均解码吞吐量。在有限训练下，我们还在高达 256K 上下文长度的 NIAH (Needle In A Haystack) 任务上取得了强劲的性能，证明了路由跨度选择是可端到端学习的。本文强调架构表述、扩展分析和系统可行性，并给出了初步验证；跨不同长上下文任务的全面质量评估留待未来工作。