Exact Causal Attention with 10% Fewer Operations
This provides an incremental improvement for compute-bound applications in machine learning where FLOPs reduction is critical, though it does not accelerate fused kernels like FlashAttention on GPUs.
The paper tackles the computational efficiency of Causal Attention by presenting Exact Causal Attention (ECA), a Strassen-style algorithm that reduces operations by 10% for exact computation, applicable to matrix multiplications involving triangular matrices in both forward and backward passes.
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.