Towards Making Linear Attention Usable
Keywords: Linear Attention, Kernel Separation, Transformers, Memory Reduction
TL;DR: We introduce a method to reduce the memory complexity of Kernel Separation in Transformer attention from $O(ND^2)$ to $O(ND)$, where $N$ is the number of tokens and $D$ dimension per attention head. We also introduce an alternative dropout mechanism.
Abstract: The original Transformer attention mechanism, based on Softmax, has time and memory complexities of $O(N^2D)$ and $O(N^2)$, where $N$ is the number of tokens and $D$ the dimension per attention head. As current LLM applications trend towards processing larger token sequences, and Transformers gain popularity in image, video, and audio processing, addressing this quadratic cost becomes imperative. Since the introduction of Transformers, numerous approaches have been proposed to linearize this scaling. One such method is Linear Attention, which captures all-to-all token pair attention in $O(ND^2)$ time. However, its drawback lies in its high memory footprint of $O(ND^2)$. While Linear Attention has shown promise in small-scale benchmarks, the high memory demand has prevented Linear Attention to be studied in context of large benchmarks and practical use cases. In this work, we demonstrate how to reduce the memory complexity to $O(ND)$ by approaching calculations from a novel perspective. Additionally, since Linear Attention does not compute the attention matrix directly, it precludes the use of traditional dropout. To address this, we introduce an alternative dropout mechanism. Our study confirms linear scaling in both wall-clock time and memory usage. We also compare our method with Flash Attention and conduct an ablation study on our proposed dropout alternative.
Primary Area: infrastructure, software libraries, hardware, systems, etc.
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 8345
Loading