Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time

We thank the constructive comments from the reviewer. We provide some clarification as follows.

W1 & Q1: Lack of empirical results and comparison with state-of-art methods in terms of speed and accuracy.

Thanks for your valuable suggestion. Firstly, we would like to point out that the main focus of our work is theoretical analysis. Nevertheless, to demonstrate the effectiveness of our proposed method, we have also conducted experimental evaluations. See Global Response Part 2: Empirical evaluation.

W2 & Q2: How does the approximation error affect the performance of the downstream tasks?

Thanks for your insightful comments. For downstream tasks, particularly those that demand high precision, we recommend

• Replacing some part of the approximated attention module to standard attention computation and fine-tune on specific datasets to bypass the slight approximation error.
• Also, the approximation error diminishes quickly with the input length $n$. Furthermore, existing empirical work shows [1, 2, 3] that the approximation error of the attention matrix doesn't affect the downstream performance that much, e.g., approximating the softmax attention doesn't change the output probabilities that much.

W3 & Q3: What are the detailed challenges for implementing this method on GPUs? Can the author provide more concrete guidance on how to implement this method?

Thanks for pointing this out. In Part 2: Empirical evaluation. in the Global Response, we present preliminary results demonstrating the efficiency of our method. However, as outlined in Part 3: Potential implementation challenges on GPUs., implementing the full gradient computation involves several challenges.

[1] Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., ... & Chen, B. H2o: Heavy-hitter oracle for efficient generative inference of large language models. NeurIPS’23.

[2] Jiang, H., Li, Y., Zhang, C., Wu, Q., Luo, X., Ahn, S., ... & Qiu, L. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. NeurIPS’24.

[3] Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., ... & Chen, D. Snapkv: Llm knows what you are looking for before generation. NeurIPS’24.

We thank the intuitive comments from the reviewer. Here are some clarifications from us.

W1 & Q1.1: What is the difference between Theorem 4.1 and Theorem 1.6 of [1]?

Theorem 4.1 presented in our paper offers two distinct advantages over Theorem 1.6 from [1]:

• Support gradient computation for three variables: Theorem 4.1 demonstrates our ability to compute the gradient with respect to three variables—intermediate variable $T_i$, weight matrix $W_i$, and value matrix $W_{V_i}$—in almost linear time, $n^{1+o(1)}$. This contrasts with Theorem 1.6 in [1], which only facilitates the almost linear time calculation of the gradient for the weight matrix $W_i$. Note that analyzing the gradient for $T_i$ is crucial and non-trivial since it is the core component that supports gradient backpropagation between multiple layers of the transformer.
• Compatibility with any differentiable loss function: Our theorem (Theorem 4.1) extends support to any differentiable loss function, thereby accommodating the commonly employed cross-entropy loss function. In contrast, Theorem 1.6 in [1] is limited to the $\ell_2$ loss for the attention matrix, precluding its direct application to the training of contemporary large language models (LLMs) that necessitate the use of cross-entropy loss.

Q1.2: How is Theorem 1.6 of [1] used in your proof?

We leverage Theorem 1.6 of [1] in our proof for the gradient computation of the weight matrix $W_i$, with a detailed explanation provided in Section F of the Appendix.

W2 & Q2: How did $n/poly(n)$ end up becoming $1/poly(n)$? Can the author write down the degree of $poly(n)$? Where is the $m$ used in the induction?

Rigorously speaking, we acknowledge that the polynomial term $poly(n)$ in the expressions $n/poly(n)$ and $1/poly(n)$ does not correspond to the polynomial with the same degree. Given that we treat the number of layers $m$ as a constant in our study, the $poly(n)$ term will indeed be asymptotically larger than $n^m$. Consequently, in our proof, we treat $n/poly(n)$ as effectively equivalent to $1/poly(n)$.

The number of transformer layers, denoted as $m$, is indeed a factor in the inductive process. When accounting for $m$, the computational time required to calculate the gradient across an $m$-layer transformer is $m n^{1+o(1)}$. However, in practical scenarios, the sequence length $n$ is exceedingly large, whereas the number of layers $m$ remains relatively small and can be treated as a constant. This is the rationale behind the absence of $m$ in our final runtime analysis, as its impact is negligible compared to the sequence length.

W3 & Q3: What are the key challenges and theoretical contributions in proving the results of the Causal attention mask?

Causal masking poses a challenge for low-rank approximations because it results in a full-rank attention matrix. We used the algorithm from [2] that separately computes entries of the low-rank representations of the attention matrix. However, we face a new challenge of a new structure of gradient compared to [2]. Due to the complicated form of the gradient, we first categorize gradient components into dot product and Hardamad product while [2] did not have this complexity. In detail, in Lemma I.6 Line 3758, [2] only addresses the dot product form of Part 1 while in our case we have to deal with the Hadamard form Part 2. Moreover, we utilize row-wise Kronecker product $\oslash$ (Line 3784) to solve Part 2 and [2] did not analyze this.

Minor suggestions

Thanks for pointing out the typo in Theorem 4.1. We have removed $m =n^{o(1)}$ in Theorem 4.1. And thanks for pointing out the duplication of line 401-410 and 409-416. We have fixed this duplication, and have uploaded the revised paper.

[1] Alman, J., & Song, Z. The fine-grained complexity of gradient computation for training large language models. NeurIPS’24.

[2] Liang, Y., Liu, H., Shi, Z., Song, Z., Xu, Z., & Yin, J. Conv-basis: A new paradigm for efficient attention inference and gradient computation in transformers. arXiv preprint arXiv:2405.05219.

New Q1.1: Comment on Remark 1.3 of [1].

Thanks for your insightful observations. Firstly, we would like to point out that the notation system used by [1] is different from our notation system. As discussed in the sentence above Definition 1.2 in [1], their $X \in \mathbb{R}^{d \times d}$ is defined by $X = Q K^\top$ and their $Y \in \mathbb{R}^{d \times d}$ is defined by $Y = V$.

We provide a table for the notations used in our paper and [1] as follows:

Then, we will discuss the Remark 1.3 of [1]. For simplicity, the following discussion is based on our notation system. Firstly, the contribution of [1] (Theorem 1.6 of [1]) is that they accelerate the gradient computation for $W_i$, whereas in our paper (Theorem 4.1 in our paper), we proposed the acceleration for the gradient on $X$, $W_i$ and $W_{V_i}$, which means the result of [1] is only a subset of our Theorem 4.1.

New Q1.2 Does this remark (Remark 1.3 of [1]) simply apply per-layer with $T_i = Y$?

We would like to point out that our $T_i$ and their $Y$ are not comparable. They do not have the same dimension. Please refer to the new Q1.1, where we have pointed out that our $T_i \in \mathbb{R}^{n \times d}$ is the intermediate variable/hidden state for each layer, and their $Y \in \mathbb{R}^{d \times d}$ in [1] only represents value weight matrix $V \in \mathbb{R}^{d \times d}$ in [1].

New Q2: The relationship between any differentiable loss function $F$ and the $\ell_2$ loss.

Thank you for pointing this out! We acknowledge your derivation. We agree that any differentiable loss function is inherently equivalent to $\ell_2$ loss. Thus, we edited the draft based on your valuable comments, in the Line 130 & 355-356 & 361-362 & 400-404.

New Q3: The $poly(n)$ term.

Thank you for your suggestion. We added the dependency of $m$ is $n^m$ is our revised paper. Please refer to Line 344-347.

[1] Alman, J., & Song, Z. The fine-grained complexity of gradient computation for training large language models. NeurIPS’24.

We extend our gratitude to the reviewer for their meticulous feedback. We offer the following elucidations:

W1: Lack of intuition regarding low-rank approximation techniques in the main text

Thanks for your suggestions. We have revised our paper; for further details, please see the paragraph from Line 161-168. The revised paragraph articulates the intuition behind the polynomial approximation technique, which is used for the low-rank approximation of the attention matrix. Also, we provide a short summary below.

As outlined in Section 3 of [1], the polynomial approximation for low-rank approximation can be succinctly summarized as follows:

• Firstly, the authors of [1] demonstrate that for any real number $x \in \mathbb{R}$, if $x$ is constrained within the interval $[-B, B]$ for some constant $B$, then $e^x$ can be effectively approximated by a polynomial of degree $g$.
• Additionally, the author establishes that a matrix with bounded entries, following the application of an exponential function, can be low-rank approximated by matrices $U_1$ and $U_2$, both belonging to $\mathbb{R}^{n \times r}$, where $r$ is on the order of $n^{o(1)}$.

W2: Lack of explanation for the steps of Algorithm 1 in Section 5.

Thank you for pointing out this point. We have added an explanation for Algorithm 1 in Line 279-285. We believe this newly added explanation can help readers have a better understanding of our algorithm.

W3: Lack of proof sketch for Lemma 5.1.

Thanks for your valuable suggestion. We revised our paper and have incorporated a proof sketch for Lemma 5.1 (Line 437-445). This sketch is intended to provide a concise overview of the proof’s intuition. This addition will enhance the reader’s understanding by offering a clearer outline of our proof strategy.

W4: Paragraph 401-407 and 409-413 are duplicated

Thanks for pointing this out. We have fixed this duplication, and have uploaded the revised paper.

Q1: What assumptions does Theorem 4.2 require?

There are three assumptions required by Theorem 4.2. We list them as follows:

• The number of layers of the multi-layer transformer $m$ is constant.
• The embedding dimension $d$ satisfies $d = O(\log (n))$.
• Each entry in the matrix is represented by at most log(n) bits ( $\log(n)$ bits model). This assumption is well-accepted and widely used in the computational complexity community, e.g., [2,3,4]. It aligns well with practical scenarios, where machine precision (e.g., 16 or 32 bits) is typically far smaller than the input context length (e.g., 128k).

The above three assumptions are also listed as the conditions of Theorem 4.2.

Q2: Practical regime of dimension $d$ and number of layers $m$.

Indeed, the embedding dimensions employed in contemporary LLMs typically range from $1024$ to $4096$. However, we contend that the multi-head attention mechanism effectively reduces the effective embedding dimension during computation. Consider the Llama3 model as a case in point. Although Llama3 has an embedding dimension of $4096$, it is divided among $32$ attention heads. Consequently, the computational dimension per head is only $4096 / 32 = 128$, which is relatively modest in comparison to the length of the input sequence. Moreover, for LLama 3.2 1B, they have $d=64$. Regarding the number of layers, $m$, which is $16$ for LLama 3.2 1B, we treat it as a constant in our analysis.

Q3: I feel that the benefits will only be apparent when n is very large (e.g., at least 100k), which requires significant computational resources and may be difficult to run. Have the authors had a chance to experiment with this?

We appreciate the reviewer's suggestion to conduct experiments with a long context length. Regrettably, our current computational resources do not permit us to carry out extensive experiments. However, we perform preliminary experiments to show some promising results, suggesting a potential fast almost linear time training method bypassing the quadratic complexity bottleneck of LLMs. See Global Response Part 2: Empirical evaluation for more details.

[1] Alman, J., & Song, Z. Fast attention requires bounded entries. NeurIPS’23.

[2] Feng, G., Zhang, B., Gu, Y., Ye, H., He, D., & Wang, L. Towards revealing the mystery behind chain of thought: a theoretical perspective. NeurIPS’23.

[3] Liu, B., Ash, J. T., Goel, S., Krishnamurthy, A., & Zhang, C. Transformers Learn Shortcuts to Automata. ICLR’23.

[4] Merrill, W., & Sabharwal, A. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics (ACL’23).

We thank the reviewer for the helpful comments. We provide some additional information below.

W1 & Q1: Lack of practical scenarios where our method outperforms other quadratic complexity methods?

Thanks for your valuable suggestions. We would like to point out that the main focus of our work is theoretical analysis. We provide rigorous theoretical analysis about the reduction of computation complexity of multi-layer transformers. However, we perform preliminary experiments to show some promising results. See Global Response Part 2: Empirical evaluation.

W2 & Q2: How does the slight approximation error affect the performance in tasks which require high precision?

Thank you for your insightful comments. For downstream tasks, particularly those that demand high precision, we recommend replacing some parts of the approximated attention module to standard attention computation and fine-tune on specific datasets to bypass the slight approximation error.

W3 & Q4: What will happen when the attention matrix is full rank?

Thank you for your suggestions. We would like to highlight that our method uses low-rank computation techniques for approximating the attention, but is NOT assuming the attention is close to low rank. Indeed, our method works for full-rank attention.

For an in-depth examination of how to deal with full rank attention matrix, please refer to Line 161-168 in our revised paper and also Part 1: Can our low-rank approximation deal with the attention matrix with full-rank? in Global Response.

W4: Potential implementation challenges on GPUs

Many thanks to your suggestions. As our study centers on the backpropagation process of multi-layer transformers, applying our method to GPUs would necessitate modifying the gradient computation functions within PyTorch and might even require writing custom CUDA code. For more details, please refer to Part 3 of the global rebuttal. And we provide some empirical results in Global Response Part 2: Empirical evaluation.

Q3: How does this method perform when in tandem with pre-training setups?

Thank you for pointing this out. Our method indeed is compatible for tandeming with pre-training. This can be achieved by “hacking” the backpropagation function of pytorch, and replacing the backpropagation of transformer layers with the algorithm provided in our work.

Q5: Quantitative comparison with recent works, FlashAttention[2], Hyperattention[3].

Thanks for your suggestion. We offer a qualitative comparison with FlashAttention and Hyperattention.

• FlashAttention is an empirical study that leverages GPU caching to enhance the speed of attention computations. In comparison, our work centers on the acceleration of attention computations from a theoretical perspective.
• In relation to Hyperattention, our research diverges in two key respects: (1) we concentrate on expediting gradient computations across multi-layer transformers, whereas Hyperattention is concerned solely with the forward pass; (2) Hyperattention employees Locality Sensitive Hashing (LSH) to approximate the attention matrix, which only ensures a small error with high probability. In contrast, our method does not employ hashing techniques, allowing us to guarantee a minimal approximation error with certainty. We are willing to discuss more about LSH per the reviewer’s request.

[1] Alman, J., & Song, Z. Fast attention requires bounded entries. NeurIPS’23.

[2] Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness. NeurIPS’22.

[3] Han, I., Jayaram, R., Karbasi, A., Mirrokni, V., Woodruff, D. P., & Zandieh, A. Hyperattention: Long-context attention in near-linear time. ICLR’24.