RACE 注意力机制：面向超大上下文训练的严格线性时间注意力层
机器学习
2026-04-21 v5 人工智能
摘要
Softmax 注意力的时间复杂度随序列长度呈二次增长，这在处理长序列时会成为瓶颈，即使采用高度优化的 GPU 内核亦然。例如，FlashAttention-2/3（精确的 GPU 优化 Softmax 注意力实现）在 NVIDIA GH200（96 GB）上，无法完成超过约400万个 token 的上下文环境下单个注意力层的前向-反向传播。我们提出 Repeated Arrays-of-Count Estimators (RACE) 注意力机制，这是一种严格线性于序列长度和嵌入维度的替代方案。RACE 注意力将指数核替换为锐化的角度相似度，通过高斯随机投影和软局部敏感哈希（LSH）来近似注意力输出，从而避免构建完整的注意力矩阵。在语言建模、掩码语言建模和文本/图像分类等任务上，RACE 注意力在序列长度达到64K时与强大基线相当或优于，同时显著降低了 wall-clock 时间和内存占用。此外，我们在单个注意力层上进行了受控的规模实验，展示了在 NVIDIA GH200 GPU 上处理750万个 token、在 Intel Xeon Gold 5220R CPU 上处理7500万个 token的能力，这远超当前主流注意力实现的能力。RACE 注意力为当今硬件上的长上下文训练提供了一种实用且理论上有据可循的机制。我们已在 https://github.com/sahiljoshi515/RACE_Attention 上发布代码。
引用
@article{arxiv.2510.04008,
title = {RACE Attention: A Strictly Linear-Time Attention Layer for Training on Outrageously Large Contexts},
author = {Sahil Joshi and Agniva Chowdhury and Amar Kanakamedala and Ekam Singh and Evan Tu and Anshumali Shrivastava},
journal= {arXiv preprint arXiv:2510.04008},
year = {2026}
}
备注
Accepted at ICLR 2026. 29 pages, 8 figures