提出超线性多步注意力,用更少计算处理超长序列,还能随机访问任意位置。
Superlinear Multi-Step Attention
- 将自注意力转为多步搜索,每步选关键片段,复杂度降至O(L^1.54)
- 在100万上下文长度下实现114词/秒吞吐,1000万长度下80词/秒
- 适合需要长程依赖、追求高效推理的模型部署场景
本文提出超线性注意力,一种全可训练的多步注意力架构,在保持随机上下文访问能力的前提下,实现长序列的亚二次复杂度。该方法将标准因果自注意力重构为包含N步的多步搜索过程,整体复杂度为O(L^{1+1/N})。以基准情况N=2为例,其复杂度为O(L^{3/2}):第一步进行O(L^{3/2})的跨度搜索,筛选出相关序列片段;第二步在选定片段上执行O(L^{3/2})的跨度注意力(即限制在所选片段内的标准注意力)。在提升鲁棒性的O(L^{1.54})配置中,于单个B200 GPU上的30B混合专家模型上实现100万上下文长度下114词/秒的平均解码吞吐,1000万上下文长度下80词/秒。经有限训练后,该模型在长达256K的NIAH任务中表现优异,证明路由跨度选择可端到端学习。本文侧重架构设计、扩展性分析与系统可行性,提供初步验证;跨多样化长上下文任务的全面质量评估留待未来工作。
原文摘要 · Abstract (English)
In this paper, we propose \textbf{Superlinear attention}, a fully trainable multi-step attention architecture that achieves subquadratic complexity for long sequences while preserving \textbf{random context access} (a.k.a.\ structural non-exclusion): no eligible token position is structurally excluded from being selected for attention. Superlinear attention reformulates standard causal self-attention as a multi-step search problem with $N$ steps, yielding an overall complexity of $O(L^{1+\frac{1}{N}})$. To illustrate the architecture, we present a baseline $N=2$ implementation, which is algorithmically analogous to standard jump search. In this $O(L^{3/2})$ instantiation, the first step performs $O(L^{3/2})$ span-search to select relevant spans of the sequence, and the second step applies $O(L^{3/2})$ span-attention (standard attention restricted to the selected spans). In an upscaled $O(L^{1.54})$ configuration for robustness, we achieve an average decoding throughput of 114 tokens/sec at 1M context length and 80 tokens/sec at 10M context in our implementation on a modified 30B hybrid MoE model on a single B200 GPU. With limited training, we also obtain strong performance on the NIAH (Needle In A Haystack) task up to 256K context length, demonstrating that the routed span selection is learnable end-to-end. This paper emphasizes architectural formulation, scaling analysis, and systems feasibility, and presents initial validation; comprehensive quality evaluations across diverse long-context tasks are left to future work.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。