Exact Causal Attention with 10% Fewer Operations
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866914087350304768 |
|---|---|
| author | Rybin, Dmitry Zhang, Yushun Tian, Ding Lin, Zhihang Luo, Zhi-Quan |
| author_facet | Rybin, Dmitry Zhang, Yushun Tian, Ding Lin, Zhihang Luo, Zhi-Quan |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2510_05175 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Exact Causal Attention with 10% Fewer Operations Rybin, Dmitry Zhang, Yushun Tian, Ding Lin, Zhihang Luo, Zhi-Quan Machine Learning Discrete Mathematics Data Structures and Algorithms 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. |
| title | Exact Causal Attention with 10% Fewer Operations |
| topic | Machine Learning Discrete Mathematics Data Structures and Algorithms |
| url | https://arxiv.org/abs/2510.05175 |