RACE Attention: A Strictly Linear-Time Attention for Long-Sequence Training¶

会议: ICLR 2026
arXiv: 2510.04008
代码: https://github.com/sahiljoshi515/RACE_Attention
领域: LLM效率 / 注意力机制
关键词: 线性注意力, LSH, 角核, 长序列训练, 注意力近似

一句话总结¶

提出 RACE Attention——用幂次角核替代 softmax 并通过可微 LSH 草图近似注意力输出，实现严格线性时间复杂度，支持单 GPU 处理 1200 万 token、单 CPU 处理 7500 万 token，在多种任务上匹配或超越 softmax 精度。

领域现状：Softmax 注意力的 \(O(N^2 d)\) 复杂度是长上下文训练的根本瓶颈。即使 FlashAttention-2/3 优化，GH200 上单层也无法处理超过 ~400 万 token。
现有痛点：线性注意力（Linear Attention、Performer）精度下降；低秩近似（Linformer）不支持自回归；YOSO 用硬 LSH 但无理论保证且不支持因果 LM。
核心矛盾：现有近似方法缺乏严格数学框架刻画效率-精度权衡，设计决策 ad hoc 且跨任务不稳定。
本文要解决什么？：设计有理论保证的严格线性时间注意力，支持因果和非因果，可处理数千万 token。
切入角度：角核的 LSH 碰撞概率恰好等于角相似度，RACE 草图可线性时间无偏估计核密度和。
核心idea一句话：用幂次角核替代 softmax + 可微 RACE 草图实现 \(O(N)\) 注意力。

用角核 \((1 - \frac{\arccos(\hat{q}_i \cdot \hat{k}_j)}{\pi})^\gamma\) 替代 softmax。不构造注意力矩阵，将 key-value 哈希到 \(S = L \times R\) 个桶，查询时聚合同桶统计量。

幂次角核: 角相似度的 \(\gamma\) 次幂替代 softmax，\(\gamma\) 越大越尖锐。LSH 碰撞概率 = 角相似度，可直接用 RACE 草图理论。
可微 RACE 草图: 用 sigmoid 软分配替代硬 SimHash，保持近似质量同时支持梯度训练。\(L\) 个独立哈希表取平均降低方差。
因果 RACE: 用前缀和流式维护因果桶计数器，支持自回归 LM。

即插即用替换 Softmax Attention，标准交叉熵训练。\(L\)（哈希表数）和 \(R\)（桶数）控制方差-精度权衡。

方法	复杂度	64K 支持	精度
Softmax (FA2)	\(O(N^2)\)	OOM	基线
Linear Attn	\(O(N)\)	✓	差
Performer	\(O(Nd^2)\)	部分	差
RACE	\(O(N)\)	✓	≈基线

硬件	Softmax 最大	RACE 最大
GH200 (96GB)	~4M	12M
CPU	N/A	75M