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

Summary

This paper studies the fine-grained complexity of gradient computation for training attention networks. In particular, the paper shows that under SETH assumption, when the norm of certain matrices are bounded by $B$, the gradient computation lower bound is almost linear. an almost optimal algorithm is used in the theorem to provide the upper bound which almost matches the lower bound.

Strengths

The paper is theoretically well-rounded. The work is among the first few to study the computational complexity lower bound for gradient of attention networks.

Weaknesses

(Pleae reply to the Questions section directly) The work is pure theoretical and lack of experiment verifications; The work is dedicated for one layer attention networks and may arguably not for "large" language models

Questions

First of all, I am not very familiar with the strong exponential time hypothesis (SETH) or the related literature. I evaluate the work on LLM and numerical analysis perspectives, which might not be adequate. In general I think this is an interesting work with potentials of guiding the practices, yet it seems not fully explored. I have the following questions for the authors 1. I think the paper is somehow over-claimed for "training LLMs" since the authors only consider one layer of attention network. How is the gradient computation error behave if we use backpropogation with multi-layers? There should be a remark or corollary for it. 2. Can the authors discuss the implication of the algorithm used for construction the upepr bound in Theorem 1.6 for real applications? For example, since when $B$ is small we have a theoretical efficient algorithm, shall we adapt certain strategy to keep the norm smaller than certain constant? Also shall we follow the theoretically-efficient algorithm when computing the gradient in practice? 3. Continue the above point, is there any numerical evidence to show that having a small norm or the algorithm used in constructing the upper bound have advantage in practice? This may greatly strengthen the point made in this work.

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors are pretty clear about their limitations.

Summary

This paper studies the complexity of gradient computation in training models with attention layers. The paper claims that their results support the same complexity as the forward computation for the backward/gradient computation. The results show the impact of the size of the entries of the attention metrics on LLM training.

Strengths

* The complexity of gradient computation for attention layers is an interesting and important problem for LLM training. The approach to complexity analysis can provide useful insights to the community. * The complexity results are theoretically sound, supported by proof sketches.

Weaknesses

* The clarity of the paper's writing could be improved. The paper could define the problem more clearly in the introduction and present the contributions clearly. It would be better to follow a flow that first introduces the computation problem, briefly describes the techniques and assumptions used to derive the complexity results, and then presents the complexity results. Currently, these aspects are intermixed. * The significance of the paper is unclear. The complexity of gradient computation aims to solve Approximate Attention Loss Gradient Computation, but the accuracy of this approximation compared to practical gradient computation is not studied. Additionally, the analysis requires the SETH assumption, the validity of which is not provided in the paper. Thus, it is unclear if the complexity presented in this paper accurately reflects practical cases.

Questions

* Could the authors explain the key differences in technique between this paper and existing work that also explores the complexity of forward and backward/gradient computations? * Could we verify the SETH assumption for attention models?

Rating

4

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

NA

Summary

This paper presents a comprehensive analysis of the fine-grained complexity involved in gradient computation for training large language models (LLMs). Building on previous work that characterized the complexity of forward computations in LLMs, this study extends the analysis to backward computations. The authors establish a threshold for gradient computation analogous to those found in forward computation. They develop a near-linear algorithm for scenarios where the bound is within the threshold and prove the impossibility of a subquadratic algorithm under the Strong Exponential Time Hypothesis (SETH) when the bound exceeds the threshold.

Strengths

1. The analysis provides theoretical evidence supporting recent successes in the efficient computation of transformers. 2. The conclusions offer valuable insights for future algorithm design. 3. The paper is well-organized and clearly presented.

Weaknesses

While the theoretical contributions are strong, the paper lacks empirical validation of the conclusions and the proposed algorithm.

Questions

1. How does the proposed algorithm perform when the bound exceeds the threshold? 2. Based on your conclusions, is there a principle in practice for adjusting the strategy for forward and backward computations? 3. What is the memory cost associated with the proposed algorithm?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

As the authors discussed the limitations in Section 6.

Summary

The authors study Approximate Attention Loss Gradient Computation (AALGC) (e.g. complexity of training LLMs with backprop): - they show that the magnitude of entries of certain matrices (layer inputs times QKV matrices) affect whether LLM backprop is "almost-linear" or not. This connects techniques like quantization to complexity. - if the max entry of the products is bounded by B and B is o(root(logn)) it's almost-linear time, and when it's more than that, it's impossible under some assumptions (SETH) - for the almost-linear case, they construct a novel algorithm for AALGC - SETH: the assumptions are the Strong Exponential Time Hypothesis: for any eps>0 exists k>=3 for which solving k-sat with n variables in O(2^{(1-eps)n}) is impossible (even with randomized algorithms) - part of the construction involves a polynomial approximation to softmax that works well for non-large entries. - the authors use some nice techniques (like the tensor trick) that not everyone between theory + ML is familiar with, so the techniques in the paper are useful in their own right - more generally, the style of the paper is nice in that the results are clear + it's very pedagogical

Strengths

- the results, as mentioned in the introduction: extending results in complexity of forward computation to complexity of backprop - discussion of how complexity changes when certain assumptions are made (magnitude of entries) or approximations are made (polynomial softmax) - connection to practice: discussion of e.g. quantization practices in the context of theory

Weaknesses

- the softmax approximation passes by a little too quickly in the paper, I had to scroll up and down throughout the paper several times to piece the definitions together. If you control+F (search) for "softmax" or "polynomial" in the paper, only a few results come up and it's mostly references. It would be really great to add a committed subsection that says "This is the precise softmax approximation we use" and later in the algorithm section point out where exactly it's used in a clear way. - could use a little more discussion on whether such bounds on B hold in practice (e.g. in some open source models)

Questions

- I can't tell if it's just notation or not, but at some points like Definition 1.4 it seems like "n" refers to (1) the sequence length (2) the number of bits in the B bound B = o(root(log(n)) , and also (3) d = O(log n) where d is the embedding dimension ? Could you elaborate on which of these are tied together? - You allude to the near-linear time algorithms possibly being tough and needing more algorithms/hardware innovations to make practical. Could you elaborate on whether there are any particular things that come to mind? - AALGC seems to switch between an epsilon parameter and B parameter around Definition 1.4

Rating

8

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

Yes

Thank you for your response

Thanks to the author for answering my questions. I currently have no more questions. I will maintain my score.