The Fine-Grained Complexity of Gradient Computation for Training Large Language Models

Thank you for the thoughtful review. Indeed our focus is on the theoretical foundations of these LLM computational tasks, although we are optimistic that these results will help future empirical work as well. For example, the prior work on the theory of these problems which we extend here [AS23] has motivated a number of new fast algorithms for attention ([KMZ23,HJK+23,ZBKR24,YAH+24]).

To answer your questions:

$\bullet$ Our algorithms and lower bounds can both apply separately to each attention head in a larger LLM. In other words, in training, you can separately apply the algorithm for each attention head and attention layer, still giving a fast algorithm for any size model. We will expand on this in the final version.

$\bullet$ Indeed, a main takeaway of our results is that one should focus on maintaining parameters with smaller norms throughout the training process in order to achieve faster algorithms. This mirrors the phenomenon observed in practice that quantization or other parameter approximations often lead to faster algorithms [AS23]. And indeed, our approach could be used for gradient computations in practice, although as we discuss in section 6, a practical implementation would likely need more algorithms engineering work.

$\bullet$ Indeed, prior work on LLM implementations has observed a similar phenomenon, that algorithmic techniques like quantization [ZBIW19] and low-degree polynomial approximation [KVPF20], which require bounded or low-precision entries, can substantially speed up LLM operations. We discuss this in section 1 (lines 66-72) of our paper, and the introductions of those papers discuss the phenomenon in more detail.

[AS23] Josh Alman, and Zhao Song. Fast attention requires bounded entries. In NeurIPS 2023.

[KMZ23] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast trans- formers via sketches for polynomial kernels. In ICML 2024

[HJK+23] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In ICLR 2024.

[KVPF20] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Trans- formers are rnns: Fast autoregressive transformers with linear attention. In ICML 2020.

[ZBIW19] Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8bert: Quantized 8bit bert. In 2019 Fifth Workshop on Energy Efficient Machine Learning and Cognitive Computing-NeurIPS Edition (EMC2-NIPS), pages 36–39. IEEE, 2019.

[ZBKR24] Michael Zhang, Kush Bhatia, Hermann Kumbong and Christopher Re. The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry. In ICLR 2024.

[YAH+24] Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, Liwei Wang. Do Efficient Transformers Really Save Computation? In ICML 2024.

Thank you for the thoughtful review. Indeed our focus is on the theoretical foundations of these LLM computational tasks, although we are optimistic that these results will help future empirical work as well. For example, the prior work on the theory of these problems which we extend here [AS23] has motivated a number of new fast algorithms for attention ([KMZ23,HJK+23,ZBKR24,YAH+24]).

To answer your questions:

1. Indeed, when the bound exceeds the threshold, our algorithm may give incorrect answers, as it is not designed to handle larger inputs. However, our lower bound result shows that this phenomenon is necessary: it is not possible to design a fast algorithm beyond the threshold using any algorithmic technique.

2. One of the major insights behind our algorithm is that: we prove that the matrices in training/inference of the attention network can be expressed as the multiplication of two low-rank matrices. Such an idea has been widely used in fine-tuning throughout the study of LLMs (such as LoRA [HSW+21]). In a sense, our method gives a theoretical foundation for such a practical method. LoRA is mainly used in fine-tuning, but our method is for backward computation, indicating that this low-rank idea may be applied successfully to the training stage as well.

3. The memory cost is also almost linear: the algorithm simply performs simple matrix-vector multiplications where the memory and time usage are both almost linear.

[AS23] Josh Alman, and Zhao Song. Fast attention requires bounded entries. In NeurIPS 2023.

[KMZ23] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. In ICML 2024.

[HJK+23] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In ICLR 2024.

[ZBKR24] Michael Zhang, Kush Bhatia, Hermann Kumbong and Christopher Re. The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry, In ICLR 2024.

[YAH+24] Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, Liwei Wang. Do Efficient Transformers Really Save Computation? In ICML 2024.

[HSW+21] Edward J Hu and Yelong Shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Lu Wang and Weizhu Chen. Lora: Low-rank adaptation of large language models. In ICLR 2022.

Thank you for the thoughtful review. We aimed to organise the paper by first introducing the problem we study (Definition 1.4 in section 1.1) then explaining our two main results and how they relate to each other (section 1.2), explaining the context of other work in this area (section 2), and then describing how we achieved these results (section 3 for standard tools we use, and then section 4 introduces our new idea). We chose this flow, of stating and explaining both the algorithmic and complexity results before getting into technical details, because we think that the contrast between the two results is one of the most important aspects of the paper that we want to get across: that having a bound on the entries is both necessary (from the complexity result) and sufficient (from the algorithm) for fast LLM training.

To answer your other questions:

$\bullet$ The existing work requires $n^2$ time to compute the gradient. Our new technique, of applying the polynomial/low-rank factorization idea to speedup the gradient computation, is novel, and has not been used before in gradient computation to our knowledge. Moreover, to the best of our knowledge, we’re the first to give an almost linear time ( $n^{1+o(1)}$ ) algorithm to compute the gradient with provable guarantees.

$\bullet$ SETH is the most popular conjecture in the area of fine-grained complexity, and forms the backbone of many results in that area, showing hardness of problems from ML, graph algorithms, string algorithms, similarity search, and more. The fact that it gives tight lower bounds for a wide variety of different algorithms problems, and that we don’t know how to beat SETH lower bounds for any of these problems, gives strong evidence for SETH. As we discuss in section 2 and section 3.2, SETH is actually a strengthening of the P != NP conjecture, and so verifying it would require, at the very least, first solving the P vs NP problem (which is far beyond the scope of this paper).

Thank you for the thoughtful review. We especially appreciate the suggestions in the ``weaknesses'' section, and we will make these suggested changes in the final version.

To answer your questions:

$\bullet$ You’re right, our main problem definition (Definition 1.4) defines these as separate parameters, and $n$ is always the sequence length. Our results discuss when fast algorithms are possible, and the answer depends on how these different parameters relate to each other, which is why we set some of the other parameters in terms of $n$. For instance, we show that if $B$ is bigger than $\sqrt{\log n}$ then no faster algorithm is possible, but if $B$ is smaller then we achieve an almost linear time algorithm.

$\bullet$ Actually, some research groups are currently working to implement these polynomial approximation approaches in practice. See, for example, the papers [MKZ23,HJK+23] we cite which discusses this in a lot more detail.

$\bullet$ In our paper, the parameter B gives a maximum norm for the input matrices, whereas epsilon is the allowed output error. E.g., in Theorem 1.6, “1/poly($n$) accuracy” means that the epsilon is any 1/poly($n$). We’ll try to clarify this more in the writing.

[KMZ23] Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. In ICML 2024.

[HJK+23] Insu Han, Rajesh Jarayam, Amin Karbasi, Vahab Mirrokni, David P. Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. In ICLR 2024.