提出新型因果注意力算法，运算量减少10%。
Exact Causal Attention with 10% Fewer Operations
- 基于代数恒等式与搜索优化，改进三角矩阵乘法。
- 在因果注意力中实现10%的运算量降低。
- 适合计算密集型场景，不依赖GPU融合内核。
我们提出精确因果注意力（ECA），一种类斯特拉森算法，可将精确因果注意力的运算量减少10%。该方法针对一类特殊矩阵乘法进行优化，其中任一操作数或输出矩阵为上三角或下三角矩阵，涵盖因果注意力前向和反向传播中的所有矩阵乘法操作，如掩码乘积 $ ext{Mask}(QK^T)$。ECA 的设计基于机器学习与组合搜索发现的代数恒等式。值得注意的是，ECA 无法加速 GPU 上的融合内核（如 FlashAttention），因其需在内存中存储大型中间表达式，而 FlashAttention 无需此步骤。但 ECA 为计算密集型应用提供了替代方案，尤其适用于对浮点运算次数有严格要求的场景。
原文摘要 · Abstract (English)
We present Exact Causal Attention (ECA), a Strassen-style algorithm that computes exact Causal Attention using 10\% fewer operations. ECA improves a special class of matrix multiplications where either one operand or the output matrix is upper- or lower-triangular. This includes all matrix multiplication operations in the forward and backward pass of Causal Attention, such as masked product $\mathrm{Mask}(QK^{T})$. ECA is built upon algebraic identities discovered via machine learning and combinatorial search. We note that ECA cannot accelerate fused kernels such as FlashAttention on GPU. This is because ECA requires materialization of large intermediate expressions in the memory, while FlashAttention does not. However, it provides an alternative approach for compute-bound applications and can potentially be useful in scenarios with FLOPs considerations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。