arXiv:2510.04008cs.LGcs.AI2025-10中稿 · ICLR被引 4

提出线性时间注意力机制,让超长文本训练成为可能。

RACE Attention: A Strictly Linear-Time Attention Layer for Training on Outrageously Large Contexts

  • 用角度相似度替代指数核,避免构建完整注意力矩阵。
  • 单次前向反向传播可处理1200万词元(GPU)和7500万词元(CPU)。
  • 适合需要超长上下文的大模型训练,如文档理解、代码生成。

Softmax注意力在序列长度上具有二次时间复杂度,导致长序列计算成本过高。例如,即使使用高度优化的FlashAttention-2/3,在NVIDIA GH200(96GB)上当上下文超过约400万词元时,也无法完成单层注意力的前向反向传播。本文提出受核启发的重复数组计数估计器(RACE)注意力,其时间复杂度严格线性于序列长度和嵌入维度。RACE通过锐化角度相似度替代指数核,并利用高斯随机投影与软局部敏感哈希近似注意力输出,避免构造完整注意力矩阵。在语言建模、掩码语言建模及文本/图像分类任务中,RACE在长达64K序列下表现优于或匹配强基线,同时显著降低运行时间和内存消耗。此外,我们在单层注意力上进行了受控缩放实验,证明在单次前向反向传播中可处理高达1200万词元(NVIDIA GH200 GPU)和7500万词元(Intel Xeon Gold 5220R CPU),远超当前最先进实现能力。RACE为现有硬件上的长上下文训练提供了实用且理论严谨的解决方案。代码已开源:https://github.com/sahiljoshi515/RACE_Attention。

原文摘要 · Abstract (English)

Softmax Attention has a quadratic time complexity in sequence length, which becomes prohibitive to run at long contexts, even with highly optimized GPU kernels. For example, FlashAttention-2/3 (exact, GPU-optimized implementations of Softmax Attention) cannot complete a single forward-backward pass of a single attention layer once the context exceeds ~4 million tokens on an NVIDIA GH200 (96 GB). We introduce Repeated Arrays-of-Count Estimators (RACE) Attention, a kernel-inspired alternative to Softmax Attention that is strictly linear in sequence length and embedding size. RACE Attention replaces the exponential kernel with a sharpened angular similarity, and approximates attention outputs via Gaussian random projections and soft Locality-Sensitive Hashing (LSH), avoiding construction of the full attention matrix. Across language modeling, masked language modeling, and text/image classification, RACE Attention matches or outperforms strong baselines up to 64K seqeuence length while reducing wall-clock time and memory usage. In addition, we conduct a controlled scaling study on a single attention layer and demonstrate processing of up to 12 million tokens on an NVIDIA GH200 GPU and 75 million tokens on an Intel Xeon Gold 5220R CPU in a single forward-backward pass, which is well beyond the capabilities of current state-of-the-art attention implementations. RACE Attention thus offers a practical and theoretically grounded mechanism for long-context training on today's hardware. We release our code at https://github.com/sahiljoshi515/RACE_Attention.

注意力机制长上下文线性复杂度大模型训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。