LevAttention: Time, Space and Streaming Efficient Algorithm for Heavy Attentions

We thank the reviewer for their comments. We address the main ones below:

-This paper does not include an overview of the recent works, which might make it hard for readers outside this field to understand the importance of the contributions. Works like [ZGD+20] develop the sparse attention, reducing the computation complexity to $O(n)$, and other works like [SYZ23, GSWY23, ZHDK23] are more related to this work: applying the numerical linear algebra techniques, including tensor trick, preconditioner, and sketching to reduce the computational complexity.

We have now cited all of these works and gave a more comprehensive discussion in our updated draft - please see the second Appendix, Section 7, with a comparison of our work to each of these works. We would like to mention that due to the hardness of Alman and Song, for softmax with the exponential function and without any assumptions, it is impossible to obtain a high accuracy approximation to the attention matrix in subquadratic time. Thus, the works above come with different assumptions.

-The experiment only focuses on the pre-trained ViT model. This paper does not study any other language models. Therefore, it is hard to see whether or not the framework developed in this paper can be generalized to more language models.

In this paper, we focus on non-causal attention and define the universal set with respect to this. Many large language models can significantly benefit from this approach, e.g., "Lookahead When It Matters: Adaptive Non-causal Transformers for Streaming Neural Transducers" in ICML, 2023. We leave the application of these ideas to language models with causal masking as future work. There are potential challenges here in order to enforce causality. One possibiity is to use online leverage scores rather than leverage sores to ensure past tokens do not depend on future tokens through the leverage score computation.

-$\epsilon$ is used to determine the large attention scores. How can we decide the value of $\epsilon$?

One thing to note is the size of our universal set is $O(d/\epsilon)$ for $p = 2$ (and $O(d^{p/2}/\epsilon)$ in general), and thus for any $\epsilon \gg d/n$, we obtain a reduction the total number of columns of the attention. As the context length $n$ is typically much larger than the dimension $d$, this results in a significant column reduction while still finding even only moderately heavy entries. A common practice is to start with a small value for $\epsilon$ and gradually increase it while monitoring the impact on accuracy.

We thank the reviewer for their comments. We address the main ones below:

• Evaluation on Large Language Models

In this paper, we focus on non-causal attention and define the universal set with respect to this. Many large language models can significantly benefit from this approach, e.g., "Lookahead When It Matters: Adaptive Non-causal Transformers for Streaming Neural Transducers" in ICML, 2023. We leave the application of our ideas to improve the efficiency of language models with causal masking as future work. One possibility is to use online leverage scores rather than leverage sores to enforce causality.

-Limited ablation studies on prediction quality

We have performed experiments with local tokens added along with the universal set of tokens but have observed that the performance doesn't improve much over the end-to-end accuracies that we report in Table 1. This shows that while adding local tokens improves the amount of attention weight captured, it does not seem to improve the accuracy in the image classification task. We are happy to add more ablation study details, such as varying the number of local tokens or different local token selection strategies.

-What are the trade-offs between exact and approximate computation of attention scores?

The short answer, ignoring polynomial factors in epsilon and logarithmic factors in $n$, the time for exact computation is $nd^2$ while for approximate computation it is $nd$.

In more detail, given a key matrix $K$, one can compute its SVD $K = U \Sigma V^T$, where $U$ is $n \times d$, $\Sigma$ is $d \times d$, and $V^T$ is $d \times d$. This takes $O(nd^2)$ time and is done once for a given key matrix $K$. Given a query $q$, we can compute the normalization factor for the $q$-th row of the attention matrix as $\|Kq\|_2^2$. Importantly, this equals $\|\Sigma V^T q\|_2^2$ since $U$ has orthonormal columns. Consequently, for each query $q$ (row of attention matrix), we can compute its normalization factor exactly in only $O(d^2)$ time. As there are only $O(d/\epsilon)$ columns of the attention matrix that could contain an epsilon-heavy entry, we just need to evaluate $\langle q, k \rangle^2$ for each key $k$ corresponding to one of these $O(d/\epsilon)$ columns.  This takes $O(d^2/\epsilon)$ time. Thus, in $O(d^2/\epsilon)$ time we can compute all heavy entries in a single row of the attention matrix exactly, and in $O(nd^2/\epsilon)$ time all heavy attention scores in the entire attention matrix exactly.

This can be sped up using sketching. By using the Johnson-Lindenstrauss transform $R$ with $O(\log n /\epsilon^2)$ rows, we have $\|RKq\|_2^2 = (1+-\epsilon) \|Kq\|_2^2$ simultaneously for all $n$ queries. Further, $R \cdot K$ can be computed in $\tilde{O}(nd)$ time for any $\epsilon > 1/\sqrt{d}$ using fast Johnson Lindenstrauss transforms. As $R \cdot K$ is a small matrix, for each of $n$ new queries $q$ we compute $\|RKq\|_2^2$ in $O((\log n) d/\epsilon^2)$ time, so $O(nd (\log n)/\epsilon^2)$ time in total to compute the normalization factor for all queries up to a multiplicative $1+\epsilon$ factor.

For the entries of $\exp(QK^T)$ before normalization, we can find a superset $S$ of $O(d/\epsilon)$ columns containing the large leverage scores in $\tilde{O}(nd + d^{\omega})$ time using sketching, where $\omega  < 2.37$ is the exponent of fast matrix multiplication (one could set $\omega = 3$ since the context length $n$ may be much larger than $d$). See "Fast Algorithm for Finding $U$" in Section 1.1 for references. Now we only need to compute the large entries of $Q \cdot S$, where $Q$ is the $n \times d$ query matrix, and $S$ is a $d \times O(d/\epsilon)$ matrix with the keys corresponding to the universal set. We can again use sketching to instead compute $Q \cdot S \cdot T$, where $T$ is an $O(d/\epsilon) \times O((\log n)/\epsilon)$ CountSketch matrix which can be used to find the heavy entries in each row of $Q \cdot S$. Importantly we compute $S \cdot T$ first in $d^2 \textrm{poly}((\log n)/\epsilon)$ time, at which point $S \cdot T$ is a $d \times O((\log n)/\epsilon)$ sized matrix, and we can compute $Q \cdot (ST)$ in only $O(nd (\log n)/\epsilon)$ time. We can then divide each row by the normalizations found in the previous paragraph.

To summarize, while approximate computation reduces time complexity from $nd^2$ to $nd \cdot \textrm{poly}((\log n)/\epsilon)$, it introduces a small approximation error controlled by epsilon.

We thank the reviewer for their comments. We address the main ones below:

• The main question I have is around the weakness that I brought up regarding the experimental section - that is, how would the case be handled when $K$ has high rank and every key has high sensitivity ?

Our theoretical results give a small universal set size (of size roughly $d^{p/2}/\epsilon$) when applied to activation functions that have at most polynomial growth, and hold for any possibly worst case attention matrix.

Unfortunately with softmax, the unbounded degree of the exponential function means that in certain cases, there provably cannot exist a universal set of size less than $n$, where $n$ is the total number of columns of the attention matrix. Here is a simple example which illustrates that the universal set can have cardinality $n$: suppose $Q = K$ is a random $n \times d$ sign matrix, where $d = C \ln n$, for a large constant $C > 0$. Then the diagonal entries of $QK^T$ equal $C \ln n$, whereas by standard concentration bounds, with high probability the off-diagonal entries are simultaneously all at most $\sqrt{C \ln n} \sqrt{\log n}$ in absolute value. Thus, softmax$(\exp(QK^T))$ will have diagonal entries that are close to $1$, while all off-diagonal entries will be $O(1/n)$. Consequently, the only universal set of softmax$(\exp(QK^T))$ is the entire set of n columns.

The example in the previous paragraph is worst-case, and the goal of our experiments was to investigate whether the insights gained from polynomial attention, where we have stronger theoretical guarantees, carry over to the more widely used softmax attention for real-world, practical use cases. Despite the worst case theoretical example for softmax given above, we observed a promising trend in the experiments: even with a relatively small universal set comprised of keys with high leverage scores (corresponding to polynomial activation with $p = 2$ in our theoretical results), we were able to capture a significant portion of the {\bf softmax attention} mass for each token by taking the universal set that we found for $p = 2$.

We also performed experiments with local tokens added along with the universal set of tokens but have observed that the performance does not improve much over the end-to-end accuracies that we report in Table 1. This shows that while adding local tokens improves the amount of attention weight captured, it does not seem to improve the accuracy in the image classification task. 

Ultimately, even if the attention matrix has high rank due to these local dependencies, the ability to approximate it with a sparse representation using both universal and local tokens offers significant computational advantages compared to standard softmax attention.

• Moreover in the experimental section have the authors considered evaluation with other methods which are popularly used in attention approximation such as flash attention ?

Our work is somewhat orthogonal to FlashAttention, which is an efficient hardware-based implementation of the quadratic-time softmax attention. Our goal is to understand structurally how much of the heavy attention mass can be captured using a small universal set of columns, together with a few local tokens. The universal set that we find, together with local tokens, could in principle be used as a preprocessing step for other attention mechanisms.

We do not claim that our implementations of leverage score based attention are highly optimized, and therefore do not expect to be faster than FlashAttention at the context lengths we consider. One should note though that as context lengths become larger, even a highly optimized quadratic-time implementation such as FlashAttention will be prohibitive, which motivates the design of subquadratic time algorithms.

We thank the reviewer for their comments. We address the main ones below:

• The result in this paper, as the authors stated, applies to any $Q$. It seems like this is because of the way $f$-sensitivity is defined (by taking the sup over all $y$). In practice $Q$ might be structured (I think there is a shared-$QK$ transformer where $Q$ and $K$ are identical), and I wonder if we can say anything about it.

Our structural results apply regardless of whether $Q = K$ or not, but we think the reviewer is asking if one can perhaps obtain a small universal set for softmax itself when $Q = K$. Unfortunately with softmax, the unbounded nature of the exponential function means that in certain cases, there cannot exist a universal set of size less than $n$, where $n$ is the total number of columns in the matrix. To see this, suppose $K = Q$ is a random $n \times d$ sign matrix, where $d = C \ln n$, for a large constant $C > 0$. Then the diagonal entries of $KQ^T$ equal $C \ln n$, whereas by standard concentration bounds, with high probability the off-diagonal entries are simultaneously all at most $\sqrt{C \ln n} \sqrt{\log n}$ in absolute value. Thus, softmax$(\exp(KQ^T))$ will have diagonal entries that are close to $1$, and off-diagonal entries that are at most $1/n$. Consequently, the only universal set of softmax$(\exp(KQ^T))$ is the entire set of n columns.

• The paper mentions that the standard attention mechanism cannot be well-approximated unless SETH is false (Alman & Song 2023), so I guess the point is that we are hoping for a good/efficient approximation for other attention variants. Although $f(x) = |x|^p$ has been proposed and used in previous work (PolySketchFormer, TensorSketch), I don’t think they are the “gold standard” for attention, and therefore the motivation of studying this polynomial attention is not so clear to me.

As you note, Softmax with the exponential function unfortunately requires quadratic time. Even now, for very long contexts, it is too expensive in terms of both running time and cost to spend $n^2$ time on state of the art models in every head and in every layer, and this will only become more problematic as $n$ grows larger. Our work thus fits into a growing body of methods for bypassing the quadratic time barrier.

In terms of specific motivation for looking at $f(x) = |x|^p$, we first give both empirical and theoretical motivation. Empirically, our experiments show that leverage-score based pruning ($f(x) = x^2$ based) is effective. From the theory side, in addition to our main results, we formulate a stochastic model of $Q$ and $K$ and show that under this model, $f(x) = |x|^p$ yields strong results. We make a beginning with our planted model which we think will inspire the study of more detailed models.

• Even if we can find all the columns containing all the large entries, does that imply a reasonable approximation after we multiply the attention matrix with the value matrix $V$? In addition, if we consider multiple self-attention layers, even though we can find the important columns of all attention matrices, it is not clear to me whether we can still have any kind of guarantee on the final output.

We are currently not aware of any work that guarantees any end-to-end guarantees for a full transformer, i.e., with multiple self-attention layers. Recent work, such as HyperAttention, provides end-to-end guarantees when multiplying by $V$ in a single layer, but it requires multiple assumptions and it is not clear if these assumptions always hold in practice. That said, HyperAttention works by separately finding the heavy attention scores and estimating the contribution from the light attention scores, and so we can use our universal set as a preprocessing step to HyperAttention, since any heavy attention must be inside our universal set. We can also adjust eps in practice, i.e., even $\epsilon = 1/\sqrt{n}$ gives a universal set of size $n^{1/2} d$, which can result in a significant savings in runtime for HyperAttention.

• Overall I think the problem that this paper proposes is an interesting problem in numerical linear algebra, but I am not really convinced about its impact on transformer theory.

We believe the concepts of universal keys and the planted model could lead to new theoretical bounds on the complexity of attention mechanisms or inspire new approximation algorithms with provable guarantees.