提出新型幂注意力机制，实现长序列建模的线性开销与高效训练。
Scaling Context Requires Rethinking Attention
- 设计可调状态大小的幂注意力层，独立于参数控制内存开销。
- 在长序列上下文学习任务中，性能超越指数与线性注意力模型。
- 开源高效GPU核函数，通过操作融合解决内存瓶颈问题。
我们指出，Transformer和非二次方架构均不适用于长序列训练：前者处理上下文成本过高，后者则过低。滑动窗口注意力虽降低每标记成本，但损害了上下文学习能力，也不适用。为此，我们提出幂注意力（power attention），一种线性开销序列建模的架构层，其状态大小可独立于参数调整，从而在实际场景中释放线性注意力优势。我们开发并开源了一套高效幂注意力的GPU内核，发现一种新的操作融合模式，有效避免内存与带宽瓶颈。在长上下文学习实验中，该模型在长序列训练下显著优于指数注意力和线性注意力。
原文摘要 · Abstract (English)
We argue that neither transformers nor sub-quadratic architectures are well suited to training at long sequence lengths: the cost of processing the context is too expensive in the former, too inexpensive in the latter. Approaches such as sliding window attention which reduce the cost-per-token of a transformer impair in-context learning, and so are also unsuitable. To address these limitations, we introduce power attention, an architectural layer for linear-cost sequence modeling whose state size can be adjusted independently of parameters, unlocking the advantages of linear attention on practical domains. We develop and open-source a set of GPU kernels for efficient power attention, identifying a novel pattern of operation fusion to avoid memory and bandwidth bottlenecks. Our experiments on the in-context learning of power attention shows that these models dominate both exponential attention and linear attention at long-context training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。