Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers

We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. You have precisely identified the strengths of our work, namely, that our main contribution is reducing the time complexity of attention approximation without relying on restrictive assumptions, enabling successful implementation in real-world applications. Below, we address your concerns regarding our work.

W1 & Q3 & W5: The experiments section is rather light, and contains only one experiment.

Thank you for your question. We would like to emphasize that our work is theoretical in nature. We have provided a detailed comparison of our work with recent papers, such as [1] and [2], in Line 129–145. Our motivation stems from the fact that their assumptions are too restrictive, making their algorithms challenging to implement successfully. In contrast, our work develops an implementable, assumption-light algorithm that can be applied in real-world scenarios while maintaining a comparable running time to [1] and [2]. The experiment in our paper serves as an example to demonstrate the implementation of our theoretical contribution and to highlight its independence from the restrictive assumptions of [1]. However, it is not the primary focus of our work. Our experiment illustrates that our theoretical results can be directly implemented in a popular transformer model, Llama 3, whereas [1] lacks any experimental validation. Additionally, our work is theoretically comparable to—and, in many aspects, surpasses—the prior theoretical contributions of [1] and [2] (see Line 129–145 and our response to Reviewer vPbo for further details). We also refer the reviewer to the Global response Empirical evaluation section for more details.

W2 & Q2: Training or finetuning

Thank you for your question. For the Llama3 experiment, we do not have the training, and we directly do approximation, i.e., our methods are training-free. We don't require retraining/fine-tuning, which is one of the advantages of our method.

We provide training analysis in our main analysis. However, practical training or finetuning requires resources and technical hardness. We leave them as future work.

W3 & Q1: What relationship does the conv-basis rank (k) have to the context length?

The parameter $k$ is not sensitive to the context length. Theoretically, the error bound remains the same when $k$ remains fixed while the context length increases (see Theorem 4.4). The algorithm runtime remains almost linear time as long as $k$ is sublinearly with the length of the input token $n$. This highlights another advantage of our algorithm: it is suitable for different scenarios and offers flexibility in parameter selection. Empirically, our experiment shows a very good performance when $k = 512$.

W4: Confusion terms in Algorithm 2

• Both $u$ and $v$ are local variables. They are introduced in Line 2 of Algorithm 2. They are initialized as 0 vectors. $u$ tracks the cumulative sum of conv bases, and $v$ tracks the first $T$ entries of $u$ (used for binary search, Algorithm 3).
• $s$ in Algorithm 2 denotes the starting index for binary search, i.e., the binary search of Algorithm 3 happens between index $s$ and $t$.
• After achieving $\tilde{H}_s$, we can get conv-basis for this index. See Lemma B.19 for more details.

We thank you again for your insightful comments.

We express our deepest gratitude to the reviewer for the time and effort in reviewing our work. It is very encouraging to see that our theoretical results are regarded as comprehensive, interesting, and well-presented. Additionally, the problem we study is recognized as important. Below, we address your concerns regarding our work.

W1 & Q2: Empirical results

Thank you for your question. We would like to emphasize that our work is theoretical in nature. We have provided a detailed comparison of our work with recent papers, such as [1] and [2], in Line 129–145. Our motivation stems from the fact that their assumptions are too restrictive, making their algorithms challenging to implement successfully. In contrast, our work develops an implementable, assumption-light algorithm that can be applied in real-world scenarios while maintaining a comparable running time to [1] and [2]. The experiment in our paper serves as an example to demonstrate the implementation of our theoretical contribution and to highlight its independence from the restrictive assumptions of [1]. However, it is not the primary focus of our work. Our experiment illustrates that our theoretical results can be directly implemented in a popular transformer model, Llama 3, whereas [1] lacks any experimental validation. Additionally, our work is theoretically comparable to—and, in many aspects, surpasses—the prior theoretical contributions of [1] and [2] (see Line 129–145 and our response to Reviewer vPbo for further details). We also refer the reviewer to the Global response Empirical evaluation section for more details.

Q1: Is there some way a conv-approximation of an existing attention matrix can be interpreted as a set of specific kernels for a subquadratic model? In particular, I wonder if the conv-basis concept relates to the semiseparable matrices of Mamba2 [3]?

Thank you for your brilliant question. The conv basis $H = \sum_{i} \mathrm{conv}(b_i, m_i)$ can be viewed as learning $m_i$ length-selective filters $b_i$ that capture different ranges of dependencies in the sequence. This connects to how SSMs use structured state transitions to model dependencies efficiently. While related, our approach differs from Mamba's semiseparable matrices.

• When defining semiseparable matrices, they rely on the matrix rank. However, our convolution matrix definition is completely unrelated to rank. We are more related to the Toeplitz matrix, in Figure 3 of Mamba2 [3].
• Mamba2 can be viewed as linear attention in some sense, while our work is approximating standard attention.
• The theoretical guarantees on approximation error are different.

We thank you again for your insightful comments.

[1] Josh Alman, and Zhao Song. "Fast attention requires bounded entries." NeurIPS’23.

[2] Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. "Hyperattention: Long-context attention in near-linear time." ICLR’24.

[3] Dao, Tri, and Albert Gu. "Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality." ICML’24.

We express our deepest gratitude to the reviewer for taking the time and effort to review our work. It is very encouraging to see that our study is recognized for addressing a timely question, providing empirical results, and focusing on the theoretical aspects. Below, we address your concerns regarding our work.

W1: The algorithm 3 is deferred to the appendix while being a key building block of the algorithm. In general, a more detailed explanation of how the binary search is carried out will improve the clarity of the paper.

Thank you for pointing out this! We completely agree that Algorithm 3, as a key building block of our technique, is crucial. We have updated the corresponding part in the revision Line 318-323. We also made a summary below.

In Algorithm 3, we use binary search to efficiently locate the convolution basis position by leveraging the non-degenerate property of the attention matrix. This allows us to find $k$-$\mathsf{conv}$-basis in our main Algorithm 1, enabling better control over the running time while bounding the error. The choice of $k$ thus balances the trade-off between accuracy and efficiency. Technically, Algorithm 3 identifies positions in the attention matrix where the $\ell_1$ norm of remaining attention values exceeds the threshold $\delta - 2 T \epsilon$. The non-degenerate property enables the binary search algorithm to find the next convolution basis position in $O(\log n)$ steps.

W2: Because the current work is built upon the previous work with low rank assumption, a more direct comparison (either empirically or theoretically) between the current method and previous one should be presented.

Thanks for pointing out this question. We would like to point out that our work does need a low-rank assumption. Note that due to Casual Mask, the attention matrix is full-rank indeed. In our paper, informally, our assumption is that the attention matrix can be decomposed into $k$-$\mathsf{conv}$-basis for some small $k$.

For the theoretical comparison, we have provided in Line 129-145. For the empirical comparison, we leave it as our future work.

We thank you again for your valuable suggestions.