优化线性注意力的计算与内存，速度提升3.3倍，适合大模型高效训练。
Transformer Based Linear Attention with Optimized GPU Kernel Implementation
- 设计新型线性注意力前向与反向传播方法，适配GPU加速。
- 在14亿参数模型上实现3.3倍加速、3.6倍内存降低。
- 适用于需要高效率的长序列建模场景，如大语言模型训练。
Transformer架构中的原始软最大注意力机制在处理N个标记、每个标记维度为D的注意力时，时间复杂度为$O(N^2D)$。尽管线性注意力（LA）理论上可将复杂度降至$O(ND^2)$且保持相近精度，但实际性能仍落后于理论预期。本文提出一种新的线性注意力前向与反向传播方法，并实现高度优化的CUDA内核。实验表明，该方法在单层和端到端设置下分别实现3.3倍的速度提升，内存消耗降低3.6倍。通过训练一个14亿参数的语言模型验证，其在主流推理基准上表现与常规注意力相当。
原文摘要 · Abstract (English)
The original softmax-based attention mechanism (regular attention) in the extremely successful Transformer architecture computes attention between $N$ tokens, each embedded in a $D$-dimensional head, with a time complexity of $O(N^2D)$. Given the success of Transformers, improving their runtime during both training and inference is a popular research area. One such approach is the introduction of the linear attention (LA) mechanisms, which offers a linear time complexity of $O(ND^2)$ and have demonstrated comparable accuracy to regular attention. However, LA in practice lags behind its theoretical efficiency. We propose a novel method for LA's forward and backward passes, along with a highly-optimized CUDA implementation. Our approach outperforms the state-of-the-art by 3.3 times in speed and reduces memory consumption by 3.6 times. We validate these improvements in both single-layer and end-to-end settings by training a 1.4 billion parameter language model, which demonstrates similar expressivity to regular attention on major reasoning benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。