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

Authors: Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, Yufa Zhou
Primary area: learning theory
Keywords: Multi-Layer Transformers, Attention Acceleration, Gradient Computation

Abstract:
The computational complexity of the self-attention mechanism in popular transformer architectures poses significant challenges for training and inference, and becomes the bottleneck for long inputs. Is it possible to significantly reduce the quadratic time complexity of computing the gradients in multi-layer transformer models? This paper proves that a novel fast approximation method can calculate the gradients in almost linear time $n^{1+o(1)}$ where $n$ is the input sequence length, while it maintains a polynomially small approximation error $1 / \mathrm{poly}(n)$ across the entire model. 
Our theory holds for general loss functions and when the multi-layer transformer model contains many practical sub-modules, such as residual connection, casual mask, and multi-head attention. 
By improving the efficiency of gradient computation, we hope that this work will facilitate more effective training and deployment of long-context language models based on our theoretical results.

Decision: reject
Reviewer scores: 6, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730054928334

thread_type:
official_review

<thread_object>
creator: reviewer_zjzg

<content>
summary:
This paper presents a novel method to approximate the gradient computation in multi-layer transformer models with almost linear time complexity. Traditionally, the self-attention mechanism in transformers has a quadratic time complexity with respect to the input sequence length, which becomes a bottleneck during training, especially for models with long input contexts. The authors propose a fast approximation algorithm that calculates gradients in time proportional to $𝑛^{1+o(1)}$ , where n is the input length, and achieves a polynomially small approximation error. This method is applicable to general loss functions and transformers with practical components such as residual connections, multi-head attention, and causal masks. The results hold significant potential for improving the training efficiency of large language models without sacrificing much accuracy.

strengths:
*  The paper introduces a method to approximate gradient computation in multi-layer transformers in almost linear time, breaking the quadratic time complexity barrier. This efficiency boost is crucial for training large language models, especially for long-context tasks, and could reduce both computational and energy costs.

* The proposed algorithm is versatile, supporting general loss functions and working with commonly used transformer components such as residual connections, multi-head attention, and causal masks. This ensures that the method can be applied to a wide range of practical transformer architectures and tasks without major adjustments.

weaknesses:
* While the paper provides a strong theoretical foundation for the proposed method, there is no empirical evidence or experimental results presented to validate its performance in real-world transformer models. This leaves a gap in understanding how well the method performs in practice compared to theoretical expectations.

* Although the approximation error is stated to be polynomially small, the paper does not deeply explore how this approximation might affect model performance in different tasks. There could be edge cases where the accuracy trade-offs are more significant, especially in critical applications requiring precise outcomes.

* The paper acknowledges that practical implementation, especially on GPUs, may face coding and optimization challenges, such as needing to redefine tensor operations and re-implement certain backpropagation functions. This could limit the immediate adoption of the method by practitioners looking for a ready-to-use solution.

questions:
* Could the authors provide empirical results or benchmarks to demonstrate the practical efficiency and effectiveness of the proposed gradient approximation method on real-world large language models? How does it compare to state-of-the-art methods in terms of both speed and accuracy?

* The paper claims that the approximation error is polynomially small, but how does this error affect downstream tasks such as text generation or classification? Are there scenarios where this approximation might lead to a noticeable degradation in performance, especially in sensitive applications?

* The authors mention potential coding challenges for implementing this method, particularly on GPUs. Have the authors considered these challenges in detail, and can they provide more concrete guidance or solutions for implementing this algorithm efficiently in real-world systems, especially for large-scale models like GPT-4 or LLaMA?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730665768728

thread_type:
official_review

<thread_object>
creator: reviewer_o3cw

<content>
summary:
In this paper, the authors provide a number of algorithms that enable one to approximately compute the gradient (and forward pass) of a multi-layer Transformer in time that is nearly-linear in $n$, the input sequence length. The key challenge is in dealing with the self-similarity matrix, which has $O(n^2)$ entries. The dimension $d$ is assumed to be $O(log(n))$.

In Theorem 4.1, authors establish nearly-linear time guarantees for computing the gradient of a single-layer Transformer.

In Theorem 4.2, the authors generalize this result to a Transformer with up to $m=n^{o(1)}$ layers.

In both Theorem 4.1 and 4.2, the guarantee is roughly $1/poly(n)$ error in $n^{1+o(1)}$ time. 

In section 6, the authors discuss further extensions of their results to various settings, including multi-headed attention, causal mask, and prompt tuning.

strengths:
Overall, I think being able to compute the attention (both forward and backward) is a question of great theoretical and practical interest.

I think the extensions to a $m$-layer Transformer (Theorem 4.2) is interesting, well-motivated, and proof is non-trivial. (However, I do have trouble verifying some parts of the proof, see below).

weaknesses:
My main criticism of the paper is that, while the results are interesting and important, I am uncertain about the contributions of the present paper, given the existing papers on this topic. 

Theorem 4.1 shows fast gradient computation for a single layer. This result seems identical (at least in big-O notation) to the result in Theorem 1.6 of [1]. On line 446, authors claim that they generalize the results in [1] in two ways: (a) accomodate general loss function L(X), and (b) generalize their result to both $W_i$ and $W_{V_i}$. However, I do not see why (a) is meaningful -- for the purpose of computing gradients, one can wlog consider a linear $L(X)$, which is exactly handled by [1]. I also do not see why (b) is meaningful -- [1] already handles derivative wrt $W_{i}$, whereas the derivative of $W_{V_i}$ is just a minor extension of existing results (see Remark 1.3 of [1]).

I feel that the generalization of 4.1 to the multi-layer setting is potentially quite nontrivial -- the error from a single-layer's grradient computation can blow up over layers, and thus handling this is non-trivial. However, I had a bit of trouble verifying the part of the proof that bounds this error propagation. The main result seems to be Lemma 5.5 (Theorem H.4). In this proof, the key seems to be bounding error accumulation in the gradient computation over $m$ layers, whose proof I have some trouble verifying. (see item 2 under questions below)

Finally, the authors discuss a number of extensions of their work to handling things like multi-headed attention, and causal masks. I think the handling of causal mask is also interesting, but the result seem to be almost immediate following [2].

questions:
1. In Theorem 1.6 of [1], authors show that there is a $n^{1+o(1)}$ time algorithm that computes the gradient up to 1/poly(n) accuracy. Can the authors please comment on the difference between this result and Theorem 4.1? Both in terms of guarantees, as well as in terms of proof technique. Does the present paper's proof use Theorem 1.6 in any way?
2. Can the authors elaborate on the proof of Theorem H.4 for the multi-layer analysis? Specifically, I refer to line 3646-3650: how did $n/poly(n)$ end up becoming $1/poly(n)$? Since the authors are using an inductive proof, I don't think one can rigorously prove the inductive step with explicitly writing down the degree in $poly(n)$. Related to this problem, I don't see where the authors use the assumed bound on $m$, surely, the inductive proof has to use this fact?
3. Can the authors please comment on what are the key challenges and theoretical contributions in proving their result on causal masks? In particular, please highlight what additional work was done on top of the results by [2].

**minor suggestions that did not affect score**
1. I suggest to remove mention of $m =n^{o(1)}$ in Theorem 4.1 if you are going to assume $m=1$ in the same theorem.
2. line 401-410 seems to be a duplication of 409-416. 

[1] [The Fine-Grained Complexity of Gradient Computation for Training Large Language Models](https://arxiv.org/pdf/2402.04497), Alman and Song 2024

[2] [Conv-basis: A new paradigm for efficient attention inference and gradient computation in transformers.](https://arxiv.org/pdf/2405.05219) Liang et al 2024a
</content>
</thread_object>

</thread>
<thread>
creation date:
1731255314535

thread_type:
official_review

<thread_object>
creator: reviewer_hh9h

<content>
summary:
This paper proposes a method to compute gradients in multi-layer transformers in almost linear time, addressing the quadratic complexity bottleneck of traditional self-attention mechanisms. By leveraging a low-rank approximation technique, the approach aims to make large language models more efficient, particularly for long-context tasks, while maintaining a small, polynomial approximation error. The algorithm is versatile, supporting general loss functions and widely used transformer components such as residual connections, multi-head attention, and causal masks, making it theoretically applicable to a broad range of transformer-based architectures.

strengths:
Strengths:

1.Innovation in Complexity Reduction: The primary contribution of an almost linear time approximation method for gradient computation addresses a significant bottleneck in transformer models, advancing efficiency in training large language models.
2.Comprehensive Theoretical Support: The paper offers a well-developed theoretical foundation, demonstrating that the proposed method maintains accuracy within a polynomially small approximation error.
3.General Applicability: The approach is compatible with various sub-modules (e.g., residual connections, multi-head attention) and loss functions, enhancing its utility across diverse transformer-based applications.
4.Potential Impact on Long-Context Models: Given the paper’s focus on scaling efficiency for long input sequences, the findings are likely to benefit future developments in large language models that need to handle extensive contexts effectively.

weaknesses:
1.Lack of Empirical Validation: Although the paper establishes a strong theoretical foundation, it lacks empirical evaluation to validate the proposed method's performance in real-world applications. While there is discussion of related methods, a deeper quantitative or empirical comparison (e.g., speed, memory usage) with current acceleration techniques could further support the claims of the paper. Without experiments or benchmarks, it is difficult to assess the practical effectiveness of the method, especially in comparison with other acceleration techniques.
2.Unclear Approximation Error Impact: While the approximation error is stated to be polynomially small, the paper does not explore how this error might affect model performance in downstream tasks. In applications requiring high precision, even small errors could be significant; this is particularly relevant in domains where precision is critical, such as healthcare or finance.
3.Limitations of Low-Rank Approximation: The proposed approach relies on low-rank approximation of the attention matrix. However, if the attention matrix is effectively full-rank, this assumption may not hold, potentially negating the efficiency gains of the method. Further clarification on the handling of full-rank attention maps would improve the paper.
4.Potential Implementation Challenges: The paper notes that implementing the approach, particularly on GPUs, may involve re-defining tensor operations and re-implementing specific backpropagation steps. This could pose a barrier to immediate adoption by practitioners, as it requires significant engineering effort. This is a good paper and I would like to raise my score if the author can provide more empirical supports.

questions:
1. Could the authors provide examples or scenarios where this method outperforms existing quadratic complexity methods in practical applications?
2. For cases where the multi-layer approximation results in slight errors, how do these impact tasks with high sensitivity to precision (e.g., in medical or financial text processing)?
3. Can the authors clarify the compatibility of this method with common pre-training setups in large language models? Specifically, how would this method perform in tandem with models like LLaMA or GPT-4?
4. What happens if the attention matrix is full-rank? Does the proposed method lose its advantage in such cases?
5. A more detailed, quantitative comparison with recent works on transformer efficiency, such as FlashAttention[1], hyperattention[2], or tensor-based approximation methods, would enhance the positioning of the proposed method.

[1] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
[2] HyperAttention: Long-context Attention in Near-Linear Time
</content>
</thread_object>

</thread>
<thread>
creation date:
1731262426343

thread_type:
official_review

<thread_object>
creator: reviewer_fvhv

<content>
summary:
The paper addresses the important problem of reducing the quadratic computational complexity of gradient calculation in transformers, proposing an approximation algorithm that computes gradients in almost linear time with a bounded polynomial error. The work builds on recent work in reducing the quadratic complexity of the forward pass and gradient computation for single-layer transformers to linear complexity.

strengths:
The work focuses on a highly relevant topic, given the growing context length in the training of LLMs, and highlights the need for algorithms with solid theoretical backing to improve the quadratic complexity of gradient computations in the attention mechanism.

The results seem broadly applicable, as they hold across different loss functions (unlike previous studies) and can incorporate important elements of multi-layer transformers, such as MLPs, residual connections, and causal masking.

weaknesses:
The work is very interesting, but I have a few questions and points to raise.

In general, I feel that there are missing intuitions about the approach in the main body. For example, the technical overview section mentions that the attention matrix can be approximated by low-rank matrices. The related work section (line 157) states that the work uses polynomial approximation techniques (Aggarwal and Alman, 2022), but there is no further discussion about these techniques in the main body. I strongly believe that a small discussion section should be included, as this addresses the heart of the challenge.

Similarly, for Algorithm 1, all the steps reference lemmas in the appendix, but there is no discussion of these steps in Section 5. While Section 5 discusses accelerating the gradient for individual variables, there should be a connection and reference to Algorithm 1 in this discussion. Furthermore, I strongly suggest including a small proof sketch for the key lemma, Lemma 5.1, which shows the linear time complexity of the gradient computation with respect to $T_i$.

Lastly, paragraphs 401-407 and 409-413 essentially say the same thing with minor rewording—please address this repetition.

Some questions:

—What assumptions do the results of Theorem 4.2 (linear time complexity and $1/poly(n)$ error) require to hold?

—The results operate in a very practical regime of $d≪n$, but $d=O(log⁡ (n))$ seems too small, considering practical values like 1024–4096. I understand that such assumptions are often necessary for theoretical results, but how sensitive do you think the results are to this choice? The same question applies to the sub-polynomial choice for the number of layers $m$ and the logarithmic bit precision. Does this imply a practical limit on the number of layers for the algorithm to remain efficient?

—It would be useful to see the algorithm in practice and compare it to quadratic attention. However, 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?

questions:
See Weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731820125581

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Rebuttal - 2

comment:
On the other hand, we kindly wish to emphasize that the main focus of our work is theoretical analysis. The importance of theoretical contributions is widely acknowledged in esteemed conferences such as NeurIPS, ICLR, and ICML. For example, papers [1, 2, 3, 4] accepted at ICLR are solely based on theoretical analysis without any experimental data. Likewise, studies [4, 5, 6, 7, 8] that focus on designing efficient algorithms for low-rank and attention approximations also bypass empirical results.

Moreover, our work tackles the quadratic time complexity associated with the backpropagation in multi-layer transformers. Our analytical framework accommodates various elements of attention mechanisms, including multi-head attention, residual connections, and causal masking. We are confident that our theoretical findings will contribute to the advancement of practical algorithm design in the future.

### Part 3: Potential implementation challenges on GPUs.
As discussed in **Part 2**, we give a simplified demo of our method implemented on attention computation. We have proved the effectiveness of our method through preliminary experiments. There are several potential challenges:
- 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. 
- As shown in our mathematical proof, the backpropagation gradient is very complicated. Therefore, we can only implement part of the gradient computation and cannot fully implement it during the rebuttal phase.
- The low rank approximation algorithm we use involves many for loop operations. Parallelizing these operations on the GPU, including SM scheduling and shared memory allocation, will take lots of time and effort, and this part may be not realistic to complete during the rebuttal. 

Some possible solutions in our mind:
By harnessing the extensibility of PyTorch, we would customize the backpropagation algorithm and devise custom CUDA kernels. Subsequently, we would enhance the performance of these kernels by fine-tuning the Streaming Multiprocessor (SM) scheduling and adeptly managing shared memory, thereby maximizing hardware utilization efficiency. Therefore, our method can achieve optimal acceleration in performance.

Due to limited time and resources, we are not able to do the full implementation. However, we have provided the preliminary experimental as above to show the potential of practical implementation

[1] Zhan, W., Uehara, M., Kallus, N., Lee, J. D., & Sun, W. Provable Offline Preference-Based Reinforcement Learning. ICLR’24.

[2] Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., & Zhang, A. R. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. ICLR’23.

[3] Wen, K., Ma, T., & Li, Z. How Sharpness-Aware Minimization Minimizes Sharpness?. ICLR’23.

[4]  Alman, J., & Song, Z. How to capture higher-order correlations? generalizing matrix softmax attention to kronecker computation. ICLR’24.

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

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

[7] Sarlos, T., Song, X., Woodruff, D., & Zhang, R. Hardness of low rank approximation of entrywise transformed matrix products. NeurIPS’24.

[8] Dexter, G., Drineas, P., Woodruff, D., & Yasuda, T. Sketching algorithms for sparse dictionary learning: PTAS and turnstile streaming. NeurIPS’24.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731820198528

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Rebuttal - 1

comment:
We gratefully thank all reviewers for their valuable and constructive feedback. 

We appreciate the reviewers FvhV, hH9h, o3CW, and zJzg for recognizing the significance of our work in addressing the quadratic complexity bottleneck of gradient computations in attention mechanisms, particularly for long-context tasks in LLMs. They commend our almost linear time approximation method for its innovation, versatility across loss functions and transformer components (e.g., residual connections, multi-head attention, causal masking), and potential to reduce computational and energy costs. The reviewers also value the comprehensive theoretical foundation, with hH9h and o3CW noting the accuracy guarantees and the non-trivial extension to multi-layer transformers (Theorem 4.2). Overall, they highlight the practical and theoretical impact of our contributions to scaling efficiency in LLM training.

We have updated a **revision** for our draft. We also update the code about experiments in the supplemental material. We summarize all the updates (in brown color) we made in the revision. All line numbers in the rebuttal correspond to the revised version.
- Line 161-168: Add a discussion for the intuition behind the low-rank approximation method used in our paper. 
- Line 279-285: Add an explanation for our main algorithm (Algorithm 1). We believe this explanation can help readers to understand our intuition more easily. 
- Line 437-445: Add a proof sketch for Lemma 5.1, which displays the high-level idea of our proof. 
- Line 130 & 355-356 & 361-362 & 400-404: Modify statements of general loss function. 

Then, we will cover some questions that reviewers commonly ask. 

### Part 1: Can our low-rank approximation deal with the attention matrix with full-rank?

The full rank nature of the attention matrix poses no detriment to our approach. 

The fundamental of our low-rank approximation approach for the attention matrix is the polynomial approximation technique introduced in Section 3 of [1]. Put another way, rather than directly decomposing the attention matrix in its full rank nature, our method involves constructing low-rank matrices to serve as the “approximation’’ for the original matrix $\exp(QK^\top) \in \mathbb{R}^{n \times n}$. For an in-depth examination of this process, please refer to Line 161-168 in our revised paper. 

### Part 2: Empirical evaluation.
Thank you for the reviewers’ valuable suggestions to provide empirical evidence for our LLM training speedup method. However, we would like to emphasize that our work is intended to be positioned as a theoretical study, focusing primarily on theoretical analysis. Nevertheless, to demonstrate the effectiveness of our method, we conducted preliminary experiments within the constraints of time and resources. These experiments were performed across various embedding dimensions, denoted as $d$, and input sequence lengths, denoted as $n$. 

We focus on the key step of approximating the attention output, $\mathsf{Softmax}(Q K^\top / d) V$, where $Q, K, V \in \mathbb{R}^{n \times d}$ are random matrices, using the conventional attention computation and our approximation technique. As indicated in the subsequent table, we have maintained a low relative error, which intriguingly diminishes as the input sequence length increases. Notably, our method exhibits a significant acceleration in computation time when contrasted with the standard attention mechanism. 

|Feature dim d  | Input token length n | Relative Error (%) | Original Running Time (ms) | Our Running Time (ms) | Speedup |
|:--: |:---------------:|:---------------:|:--: |:--: |:--: |
|4|10*2**d=160|4.673|1.010|0.710|1.423|
|5|10*2**d=320|2.505|1.090|0.630|1.730|
|6|10*2**d=640|2.846|2.190|0.850|2.576|
|7|10*2**d=1280|1.838|3.780|0.970|3.897|
|8|10*2**d=2560|2.005|12.160|1.290|9.426|
|9|10*2**d=5120|1.620|47.300|1.790|26.425|
|10|10*2**d=10240|1.286|192.330|7.200|26.713|
|11|10*2**d=20480|1.558|807.740|20.030|**40.327**|
|12|10*2**d=40960|0.773|3297.410|41.750|**78.980**|

In the above Table, we use $n=10 \times 2^d$ as $d = O(\log(n))$ in our assumptions. And the approximation error diminishes rapidly as the input sequence length $n$ increases. This aligns with our theoretical analysis that the approximation error of the entire model can be bounded by $1 / poly(n)$.

 The code is submitted as supplementary material in the revised version of this paper. We hope this additional experiment can address reviewers’ concerns about the practical efficacy of our method.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The reviewers and AC found this work to be an interesting contribution. In terms of rating, this paper was borderline, slightly below the acceptance threshold. However, the final decision is rejection for the following reasons:

- As initially raised by Reviewer o3CW, this work appears somewhat incremental over [Alman & Song, 2024a]. The authors state: 
"Fast gradient computation. The prior study in Alman & Song (2024a) demonstrated that the gradient of $W_i$ can be computed in almost linear time. We extend their findings by adapting their approach to accommodate general loss function $L(X)$ (as defined in Definition 3.1) and further generalize their results to include the gradient computation for both $W_i$ and $W_{V_i}$ in each transformer layer (Lemma 5.2 and 5.3)". First, I am not sure how much complexity $W_{V_i}$ brings. Second, I don't see either how loss function itself make a major difference in gradient compute complexity. At the minimum, the authors should have done a much better job distinguishing novelty and proof techniques over [Alman & Song, 2024a].

- As pointed out by Reviewer FvhV, the constraint on $d=log(n)$ is not very realistic and not fully in line with practical LLMs. There should have been comprehensive discussions of practical considerations: What the typical scaling of $d$ and the number of heads are, and how the results change if $d$ grows polynomially rather than logarithmic.

- Additional drawbacks of this work include the lack of thorough empirical validation, which is partly addressed during the author response, and connections to practical LLM training.

Overall, this submission could be strengthened by addressing the above points and will represent a strong contribution to an upcoming venue. However, current weaknesses prevent its acceptance as is.