Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers
Yingyu Liang, Heshan Liu, Zhenmei Shi, Zhao Song, Zhuoyan Xu, Junze Yin
Introduction
Numerous notable large language models (LLMs) in natural language processing (NLP) have emerged in these two years, such as BERT , OPT , PaLM , Mistral , Gemini , Gemma , Claude3 , GPT-3 , GPT-4 , Llama , Llama2 , Llama3 and so on. These models have profoundly changed the world and have been widely used in human activities, such as education , law , finance , bio-informatics , coding , and even creative writing such as top AI conference reviews . The key component of the generative LLMs success is the decoder-only transformer architecture introduced by . The transformer uses the self-attention mechanism, allowing the model to capture long-range dependencies in the input sequence. Self-attention computes a weighted sum of the input tokens, where the weights are determined by the similarity between each pair of tokens. This enables the model to attend to relevant information from different parts of the sequence when generating the output. However, the computational complexity of the self-attention in transformers grows quadratically with the input length , limiting their applicability to long context, e.g., 128k, 200k, 1000k input tokens for GPT4 , Claude3 , Gemma respectively.
To overcome the complexity obstacle of , many works study more efficient attention computation methods that can scale almost linearly with the sequence length while maintaining the model’s performance. show that if all entry of is bounded and , will be “close” to a low-rank matrix. Then, they present an algorithm that can approximate attention computation in almost linear time. Similarly, by uniform softmax column norms assumption and sparse assumption, solve attention computation in almost linear time, where they identify large entries in the attention matrix and only focus on them.
On the other hand, many works find that the attention pattern has -like (or “diagonalized”) structure (see Figure 1 (b)), mathematically, when , where we can see as the position distance between two tokens. It is relevant to the bag-of-words or -gram concept, i.e., adjacent symbols or words in NLP. Furthermore, the -like structure can be connected to convolution recurrent models and structured state space models (SSMs) as Mamba . More specifically, we can use multiple convolution matrices to approximate an attention matrix, whose intuition is similar to the low-rank approximation in the sense of computation acceleration. Note that the matrix product of a convolution matrix and a vector can be computed by Fast Fourier Transform (FFT) with time complexity , while the naive way takes time (see details in Figure 1 (a)). Therefore, it is natural to ask:
Can we use the convolution matrix to accelerate the attention computation?
Thus, our algorithm can achieve attention inference in , without needing any parameter updates e.g., re-train or finetune. Our theorems can also applied to accelerate attention training . In detail, our methods take time for forward computation and time for backward gradient computation (Theorem 5.6). Furthermore, applying our technique to low-rank approximation of attention matrices , we can extend their results to more general settings (Theorem 6.5). In detail, only works on attention approximation without an attention mask, while ours can be applied to different kinds of attention masks including the most popular causal attention mask (Definition 3.2). It shows the broad applicability of our analysis. We summarize our contributions below.
Our results are beyond or comparable to the two brilliant previous works in the following sense. (1) To guarantee a small approximation error, for the attention matrix, needs bounded entries assumption and assumption, while needs uniform softmax column norms assumption and sparse assumption. However, without all these assumptions, our algorithm can still guarantee a small approximation error (Corollary 4.5), i.e., our algorithm can apply to any including unbounded matrices, dense matrices, and any hidden dimension . (2) To guarantee a truly subquadratic running time, needs to assume to get time complexity. However, for our algorithm, as long as and , we achieve running time . This has much less restriction on . Moreover, our time complexity covers from to with different , while can only handle . (3) To guarantee a truly subquadratic running time, needs to assume , as they get time complexity where is the number of large entries in attention matrices. Our work gets time complexity and we need to get truly subquadratic running time. For the situation and , both our algorithm and run in time. For the situation , and , running time in will be truly super-linear while our algorithm remains almost linear timeConsidering the case where attention matrix is all lower triangular matrix, we have and ..
Our contributions.
Our algorithm can quickly decompose any lower triangular matrix into its convolution basis (Algorithm 2), so, via FFT, we can solve Exact Attention Computation task in (Algorithm 1 and Corollary 4.5). When , our method takes almost linear time . Our results are beyond or comparable to previous works (see comparison above).
During attention inference, our algorithm takes , without needing any parameter updates e.g., re-train or fine-tune (Theorem 4.4). It may enable further improvement and scalability of LLMs in the longer context.
During attention training, our methods take time for forward computation and time for backward gradient computation (Theorem 5.6). It may save time, resources, and energy for nowadays LLMs training.
Our broad applicable technique can be applied to the low-rank approximation of attention matrices and extend existing results to more general settings (Theorem 6.5).
Roadmap.
In Section 2, we introduce the related work. In Section 3, we present the background for this paper. In Section 4, we present our algorithms and mathematical analysis for the convolution approximation for attention computation. In Section 5, we present the approximation during training forward and backward gradient. In Section 6, we explain how to apply our technique to the low-rank approximation. In Section 7, we discuss two case studies: Longlora and RoPE.
Related Work
Very recent works study the -like attention matrix. find that in-context learning is driven by the formation of “induction heads”–attention heads that copy patterns from earlier in the input sequence. This is reflected in the attention matrix becoming more diagonal, with tokens attending primarily to preceding tokens that match the current token. In Figure 6, they show a similar -like attention pattern for other important attention circuits. Figure 3 of shows that in a minimal classification task, the abrupt emergence of in-context learning coincides with the formation of an induction head, characterized by a diagonal attention pattern. proves that for a simplified task, gradient descent causes a transformer to encode the causal graph structure of the task in the attention matrix. This results in tokens attending primarily to their causal parents reflected in a sparse diagonal structure (Figure 2).
Fast attention computation and long context LLM.
The development of efficient attention computation has been an active area of research in recent years. The standard self-attention mechanism, introduced in the transformer architecture , has a quadratic complexity with respect to the sequence length, which limits its applicability to long sequences. To address this limitation, various approaches have been proposed to improve the efficiency of attention computation. One line of research focuses on patterns of sparse attention that reduce the number of computations . Another approach is to use low-rank approximations or random features for the attention matrix , which reduces the computational complexity to linear in the sequence length. In addition, using linear attention as a proxy of softmax attention is a rich line of work . These developments in efficient attention computation have enabled transformer-based models to process longer sequences and have opened up new possibilities for their application in various domains .
Convolution in language model and FFT.
There are many subquadratic-time architectures are proposed to address Transformers’ computational inefficiency on long sequences, gated convolution recurrent models , and structured state space models (SSMs) . They can use global or local convolution operations to replace attention while keeping a comparable performance. The convolution operation can be computed by fast Fourier transform (FFT) efficiently . Moreover, the development of efficient convolution algorithms like Winograd and FFT-based convolutions has further optimized the computation, reducing the memory footprint and improving the overall speed. There are many other works studying Fourier transform .
Preliminary
In Section 3.1, we introduce the basic definitions and mathematical properties. In Section 3.2, we give the formal definition of the sub-convolution matrix and present it basic properties.
1 Basic Definitions and Facts about Attention and 𝖼𝗈𝗇𝗏𝖼𝗈𝗇𝗏\mathsf{conv}
Now, we present basic definitions. We start by introducing the input and weight matrix.
It is straightforward to see . In generative LLMs, there is a causal attention mask to guarantee the later tokens cannot see the previous tokens during generation. It is defined as follows:
We define the causal attention mask as , where if and otherwise. We define be the -th column of .
Now, we are ready to introduce the mathematical definition of the exact attention computation with a mask.
In Definition 3.3, we divide the operation into an element-wise operation and a diagonal normalization matrix to obtain a clear formulation.
Efficiently computing the attention needs to exploit structured matrices that enable fast multiplication algorithms. Here, we define the convolution matrix, which is a structured matrix where each row vector is rotated one element to the right relative to the preceding row vector, which is defined as follows:
By the following fact, we know that the rank of a convolution matrix can be an arbitrary number. Thus, our -basis is totally different from the rank basis.
The proof is trivial by Definition 3.5. ∎
Efficient computation of the convolution operation is crucial for many applications. The convolution theorem states that the circular convolution of two vectors can be computed efficiently using the Fast Fourier Transform (FFT). This leads to the following claim (see proof in Appendix A.1):
One property of convolution matrices is that they are additive with respect to the input vectors. In other words, the convolution of the sum of two vectors is equal to the sum of the convolutions of the individual vectors. This is stated formally in the following claim:
This is trivial by Definition 3.5 and matrix product operation is additive. ∎
2 Sub-convolution Matrix: Definitions and Properties
If we would like to use as a basis system, we need to introduce some new concepts. Recall that, in general, the sum of two rank-1 matrices is a rank-2 two matrix. Due to being additive, the sum of two convolution matrices is another convolution matrix, which does not hold the above property. Thus, we need to introduce sub-convolution matrices to be the basis.
Similarly, sub-convolution can be computed in time via FFT.
This is trivial by considering the calculation between the truncated matrix of and the truncated vector of with Claim 3.7. ∎
Here, we present the definition of the matrix with - basis which is non-reducible.
The following lemma establishes that any non-zero lower triangular matrix can be represented as a matrix with a - basis for some unique between and . The proof is Appendix D.1.
𝖼𝗈𝗇𝗏𝖼𝗈𝗇𝗏\mathsf{conv} Approximation during Inference
In Section 4.1, we introduce the basic definitions to support our algorithmic analysis in this section. In Section 4.2, we present the binary search and recover - algorithms and present their theoretical guarantees. In Section 4.3, we provide the formal version of our main result.
Recall that any non-zero lower triangular matrix can be represented as a matrix with a - basis for some unique between and (Lemma 3.12). However, exactly getting is hard and the definition is too strict for the algorithm design. Thus, for more flexibility, we introduce a more general definition of non-degenerate - basis as below, which is a proxy notion to relax the conditions required in algorithm design.
Here -non-degenerate - basis means that each basis cannot be “covered” by the other basis easily.
The following theorem establishes that any non-zero lower triangular matrix can be represented as an -close -non-degenerate - basis matrix. There may be many different choices of , which provide flexibility for our Algorithm 1.
By Lemma 3.12, we have is a matrix with - basis for some . We finish the proof by setting and . ∎
2 Algorithms and Their Properties
Now, we present our main Algorithm 1. We also present Algorithm 2 and Algorithm 3 being used.
3 Main Theoretical Result
In this section, we present our main result.
whose time complexity is given .
The proof idea is that using binary search to recover all non-degenerate basis (Lemma A.15), which takes time and has upto error (Lemma A.16). Then, via FFT (Claim 3.10), we finish the proof. ∎
Furthermore, we can exactly get , i.e., , through Algorithm 1 with time complexity in the worst case.
We set , , and as the input of Algorithm 1. Then, the proof follows Theorem 4.3 and Theorem 4.4 . ∎
By Theorem 4.4, when , we directly get the attention inference time complexity is with error upto as claimed in Section 1. It may enable further improvement and scalability of LLMs in the longer context.
𝖼𝗈𝗇𝗏𝖼𝗈𝗇𝗏\mathsf{conv} Approximation during Training Forward and Backward Gradient
We can apply our algorithm to accelerate attention training including forward and back propagation. We first define the attention training task, which is also used in .
Recall that during inference, we have the size matrix . Similarly, in gradient calculation, we have an size matrix, and we denote it as .
Then, we are ready to present our main results for attention training.
During backward computation, we can convey the properties of low-rank and convolution at the same time (Lemma B.13 and Lemma B.15). Then, by tensor trick, we can compute the attention gradient based on attention inference (Lemma B.9). We finish the proof by Theorem 4.4. ∎
Note that only needs to convey the low-rank property, while our proof needs to convey the properties of low-rank and convolution simultaneously, which is a more general analysis.
Our Theorem 5.6 shows that our algorithm can accelerate Transformer training as well. It may save time, resources, and energy for nowadays LLMs training.
Low Rank Approximation
We can apply our analysis technique to low-rank approximation setting in . In detail, only works on attention approximation without an attention mask. Equipped with our mask analysis trick, we can generalize their results with different kinds of attention masks including the most popular causal attention mask. We first introduce some practical attention masks.
We define the continuous row mask as , where for each , we are given such that if and otherwise.
Then, we have the following main results for the low-rank setting. The proof is in Appendix C.2.
The time complexity to get is
when is a causal mask defined in Definition 3.2.
when is a row change mask defined in Definition 6.1.
when is a continuous row mask defined in Definition 6.2.
when is a distinct columns / rows mask defined in Definition 6.3 / Definition 6.4.
Our Theorem 6.5 has the same error guarantee as . For the normal mask, e.g., casual attention mask (Definition 3.2), Theorem 6.5 shares the same time complexity as theirs.
Discussion
In this section, we discuss how to implement our methods in some popular long-context LLMs, e.g., LongLora and the RoPE model family.
Our and low-rank approximation can be applied to LongLora , whose mask is shown in the left of Figure 3. They use this kind of sparse mask to extend the context sizes of pre-trained large language models, with limited computation cost, e.g., extending Llama2 70B from 4k context to 32k on a single A100 machine. As the “diagonalized” mask structure, we can directly apply our Algorithm 1 by replacing the causal attention mask (Definition 3.2) with their sparse mask for the approximation with time complexity . Similarly, for the low-rank approximation, we directly use the second statement in Theorem 6.5 by considering row change by amortized constant mask defined in Definition 6.1 with time complexity , where for any .
RoPE.
Acknowledgement
Research is partially supported by the National Science Foundation (NSF) Grants 2023239-DMS, CCF-2046710, and Air Force Grant FA9550-18-1-0166.
In Section A, we present additional details and proofs related to the convolution approximation approach. In Section B, we introduce the approximation in gradient. In Section C, we include supplementary material for the low-rank approximation. In Section D, we present a collection of useful tools and lemmas that are referenced throughout the main text and the appendix. In Section E, we provide more related work.
Appendix A Technical Details About 𝖼𝗈𝗇𝗏𝖼𝗈𝗇𝗏\mathsf{conv} Approximation
In Section A.1, we present the background of Toeplitz, circulant, and convolution matrices. In Section A.2, we develop more mathematical tools for studying the approximation. In Section A.3, we give the key lemmas we used. In Section A.4, we use these tools and lemmas to prove our main theorem for the approximation. In Section A.5, we analyze our case study.
The integer may have different ranges. We will specify these ranges in later text, corresponding to different contexts.
The Toeplitz matrix is one such structured matrix that has constant values along its diagonals. We define it as follows:
Furthermore, we define the circulant matrix, which is a structured matrix where each row vector is rotated one element to the right relative to the preceding row vector, which is defined as follows:
Finally, we present a basic fact about the Hadamard product.
Below, we explore the properties of , , , and .
The proof directly follows the Definition 3.2, Definition A.2, and Definition 3.5. ∎
where the first step follows Claim A.6, i.e., , the second step follows Fact A.7 and the last step follows Fact A.8. We finish the proof by for FFT. ∎
A.2 Mathematical Tools Development for k𝑘k-𝖼𝗈𝗇𝗏𝖼𝗈𝗇𝗏\mathsf{conv} Basis
When a lower triangular matrix is expressed as the sum of convolution matrices, it is useful to understand the structure of the entries in . The following claim provides an explicit formula for the entries of in terms of the basis vectors of the convolution matrices.
For any , we have .
This is trivial by following , the Definition 3.5 and Definition 3.9. ∎
We present the property of as follows:
with time complexity , where denotes the -th row of .
where the first step follows from the definition of (see Definition A.10), the second step follows from simple algebra, the third step follows from the fact that the -th column of is equal to the -th row of .
For any vector , we need time to get .
Thus, in total, the time complexity is . ∎
The key idea behind our approach is to express the matrix exponential of a matrix with - basis as the sum of sub-convolution matrices involving the basis vectors. This allows us to efficiently approximate the exponential of the attention matrix. We show how to compute the new basis vectors of the convolution matrices from the original basis vectors below.
and with time complexity .
As is an element-wise function, when we have and
When we have .
Thus, we have .
We need time to get for any . Then, we need time for element-wise and minus operation for terms. Thus, in total, we need time complexity. ∎
where the first step follows the lemma statement, the second step follows the property of Hadamard product and the last step follows the lemma statement. ∎
A.3 Lemma Used in Main Theorem Proof
Part 1: and
Part 3:
Part 4: for any .
We use the math induction to prove the correctness.
Part 1: and
Part 2: (Denote , after the -th loop.)
Part 3:
Part 4: for any
We have and by the line 9 and line 10 in Algorithm 2.
We denote the output of Search() as . Now, we prove .
It is clear that . For any , we have line 7 in Algorithm 3 as
where the first step follows from Definition A.10 (), the second step follows from Part 1, and the last step follows from Definition 4.2 ().
where the first step follows from the triangle inequality, the second step follows from Definition 4.2 (), the third step follows from , the fourth step follows from Definition 3.9, the fifth step follows from Part 3, and the last step follows from Definition 4.2 ().
Similarly, when , we have Eq. (2) as
where the first step follows from the triangle inequality, the second step follows from Definition 4.2 (), the third step follows from simple algebra, the fourth step follows from the triangle inequality, the fifth step follows from Part 3, and the last step follows from Definition 4.1.
Thus, we can claim, when , we have , and we have otherwise. Therefore, by binary search, we can get .
We have and at line 8 in Algorithm 2. Thus, we have
where the first step follows from simple algebra, the second step follows from Algorithm 2 (line 8), the third step follows from , the fourth step follows from Definition A.10 (), the fifth step follows from Definition 4.2 (), the sixth step follows from , the seventh step follows from Definition 3.9, the eighth step follows from simple algebra, and the last step follows from Definition 4.2 ().
We can get for any similarly as Proof of Part 3.
We can check the initial conditions hold. Thus, we finish the whole proof by math induction. ∎
Building upon Lemma A.15, we now analyze the overall error of our approach for approximating the attention computation. Recall that our goal is to efficiently approximate the matrix , where and . We will show that by using the approximate basis vectors recovered by Algorithm 2, we can construct matrices and such that the approximation error is bounded. The following lemma provides this error analysis:
where the first step follows from triangle inequality, the second step follows from and the last step follows from and Eq. (3).
by Lemma A.13 and line 12 in Algorithm 2.
In each loop, we call times of binary search function. In each binary search function, we take time for line 6 in Algorithm 3 by Lemma A.12. Thus, we take in total for the search (Algorithm 3) in each loop.
In each loop, we take time for line 7 in Algorithm 2 by Lemma A.12.
Thus, we take total for the whole loop.
We take time for the line 12 in Algorithm 2 by Lemma A.13.
In total, we take time. ∎
We are now ready to prove our main result for the approximation approach. Theorem A.17 brings together the key components we have developed: the existence of a - basis for the attention matrix (Definition 4.2), the ability to efficiently recover an approximate - basis (Algorithm 2 and Lemma A.15), and the bounded approximation error when using this approximate basis (Lemma A.16). The theorem statement is a formal version of our main result, Theorem 4.4 and Algorithm 1, which was presented in the main text. It specifies the input properties, the approximation guarantees, and the time complexity of our approach.
A.4 Proof of Main Theorem
whose time complexity is given .
Denote . By Claim 3.10, we take time to get via FFT as - basis and columns in . Similarly, by Claim 3.10, we take time for via FFT as - basis. Finally, we take time to get as is a diagonal matrix.
Thus, in total, we take time complexity. ∎
A.5 Construction for Case Study
For each , let and for all
Then we have for all , for all , for some function .
where the first step follows from the assumption that for each and , and , the second step follows from simple algebra, the third step follows from the , and the last step follows from the definition of the function .
and . For all , we have .
Then we have for all , for all , for some function .
where the first step follows from construction condition, the second step follows from simple algebra, and the last step follows from the trigonometric properties.
Let be a permutation of .
Then we have for all , for all , for some function .
When is odd, we can show similar results by the same way. Thus, we complete the proof. ∎
if
Then, there is a vector such that
where the second step follows from the fact that is applied entry-wisely to a vector.
By the assumption from the Lemma statement that if and if , we get
which is exactly equal to (see Definition A.3).
Therefore, combining with Eq. (A.5), we have
For each , .
Then, there is a vector such that
Let defined in Definition A.24 satisfy the properties in Lemma A.20.
Then, there is a vector such that
By Lemma A.20, we have for all , for all ,
where the first two steps from Definition A.24, and the last step from Lemma A.20. We finish the proof by denote as in Lemma A.22. ∎
Appendix B 𝖼𝗈𝗇𝗏𝖼𝗈𝗇𝗏\mathsf{conv} Approximation in Gradient
In Section B.1, we present the basic definitions. In Section B.2, we combine all these definitions to form the loss function. In Section B.3, we analyze the running time. In Section B.4, we present the proof of the main theorem of approximation in gradient.
B.2 Loss Functions
Now, we start the construction of the loss function.
For every , for every , we define to be .
The proof is trivial by element-wise multiplication. ∎
From the Lemma statement, by Lemma B.8, we have
where the first step is from the chain rule and the second step follows from .
where the last step is by simple algebra.
Let be defined as in Definition B.6:
Let be define as in Definition B.7:
where the 1st step is because of Definition 5.1, the second step follows from Eq. (B.2), the third step follows from Eq. (8), the fourth step follows from Eq. (9), and the fifth step follows from Fact D.9. ∎
B.3 Running Time
In this section, we analyze the running time of the approximation approach for computing the training forward pass and backward gradient. We build upon the key definitions and loss functions introduced in the previous sections to derive the running time of the algorithm.
Suppose is a - matrix defined in Definition 3.11 with known basis.
which can be done in time by Fact A.5.
The second part is trivial by Definition B.3. ∎
Suppose takes time.
can be expiciltiy computed in time.
Firstly we can compute , this can be done in , since we run times a vector oracle (Lemma B.10) for times.
’s rank- factorization can be explilcitly computed in time.
Note that . Since both and are known. Thus, the result is trivial. ∎
Let denote a rank- matrix with known low-rank factorizations.
Using a standard linear algebra trick, we can show that
Let .
Let denote a rank- matrix with known low-rank factorizations.
Let . It is easy to see that can be written as .
We firstly compute , since has columns, each column will take time, so in total it takes time.
Then, we know that which takes time per . There are different , so it takes time.
Overall it takes time.
Let (This is obvious from definition of )
For any vector , we firstly compute , then we compute .
Firstly, we can compute , this takes time.
Second, we can compute , this takes time.
Putting it all together we complete the proof. ∎
B.4 Proof of Main Theorem
In this section, we present the formal proof of our main theorem regarding the approximation approach for efficiently computing the training forward pass and backward gradient of the attention mechanism.
Suppose is a - matrix defined in Definition 3.11 with known basis. Then there is an algorithm that runs in time time to compute the gradient of attention loss defined in Definition 5.1.
We need to choose , thus total running time is
by putting everything together from Lemma B.9, Lemma B.10, Lemma B.11, Lemma B.12, Lemma B.13, Lemma B.14, Lemma B.15, Lemma B.16. ∎
For the forward, we directly get the correctness by Theorem 4.4. For the backward, we directly run error propagation analysis which is similar to and proof of Lemma A.16.
Appendix C Incorporating Weighted Low Rank Approximation
In Section C.1, we introduce the preliminary for this section. In Section C.2, we present the proof of our main result for the low-rank approximation. In Section C.3, we present the algorithm and its mathematical properties for causal attention mask. In Section C.4, we analyze the algorithm and its mathematical properties for row change by amortized constant mask. In Section C.5, we study the algorithm and its mathematical properties for continuous row mask. In Section C.6, we analyze the property of the mask matrix with distinct columns or distinct rows.
with any arbitrary .
In the following lemma, we prove the validity of the statement that if there exists an algorithm whose output is in time, then there exists an algorithm outputs in time. We will combine everything together and show the soundness of this statement later in the proof of Theorem C.4.
which takes time, then, there exists an algorithm promise that
Suppose there exists an algorithm whose output is satisfying and takes time. We denote this algorithm as Alg.
Let . Let . Then, .
Computing and takes time. Computing takes time. Therefore, it takes time in total. ∎
C.2 Proof of Main Results
The time complexity to get is
when is a causal mask defined in Definition 3.2.
when is a row change mask defined in Definition 6.1.
when is a continuous row mask defined in Definition 6.2.
when is a distinct columns / rows mask defined in Definition 6.3 / Definition 6.4.
where the first step follows and , the second step follows mask is element-wise operation, the third step follows Definition C.1, and the last step follows .
By Lemma C.2, the matrices and defining can be computed in time.
By Lemma C.3, if we can compute in time, we can compute in time.
Finally, we finish the proof by following Lemma C.6 for the causal mask, Lemma C.8 for row change by amortized constant mask, Lemma C.9 for continuous row mask, and Lemma C.12 for distinct columns mask or distinct rows mask. ∎
C.3 Causal Attention Mask
In this section, we present the causal attention mask.
Let denote the -th row of .
Let be the support set defined in Lemma C.5. Note that for the causal attention mask, we have for any . Thus, by Lemma C.5, we have
Computing , for all takes time.
Note that by the definition of inner product
Therefore, it also takes to compute for all .
Therefore, it takes times in total. ∎
C.4 Row Change by Amortized Constant Mask
In this section, we analyze the row change by amortized constant mask.
Let be the causal attention mask defined in Definition 3.2. Then we have is a row change by amortized constant mask defined in Definition 6.1, where , .
The proof directly follows the two Definitions. ∎
which takes time.
We will prove it by induction. It is obvious that base case is correct, because .
For a fixed , we suppose has the correct answer. This means is correct for that , i.e., .
Now we use and to generate by adding terms in and deleting terms in ,
where the first step follows Algorithm 5 line 10 and line 12, the second step follows , and are disjoint, the third step follows simple algebra, and the last step follows the as the second step.
Therefore, we have is correct, i.e., . Thus, is also correct by Lemma C.5. Finally, we finish proving the correctness by math induction.
Note that there are two for-loops in this algorithm. Inside the inner for-loops, it takes time to compute
The inner for-loop has iterations, and the outer for-loop has iterations.
Therefore, it takes time in total. ∎
C.5 Continuous Row Mask
In this section, we study the continuous row mask.
The correctness is trivially from the construction of the segment tree.
The running time is dominated by . This time comes from two parts, where the first is from building the segment tree by , and the second part is from for-loop by . ∎
C.6 Distinct r𝑟r Columns or Rows
Now, we analyze the mask matrix with distinct columns.
Let be the distinct columns mask defined in Definition 6.3. Let denote disjoint subsets and be defined in Definition 6.3. Let denote that and is the smallest index in .
Now, we analyze the mask matrix with distinct rows.
Let be the distinct rows mask defined in Definition 6.4. Let denote disjoint subsets and be defined in Definition 6.4. Let denote that and is the smallest index in .
Therefore, we have shown Eq. (10), which completes the proof. ∎
The correctness and running time is directly follows Lemma C.10 for the column case and Lemma C.11 for the row case. ∎
Appendix D Supporting Lemmas and Technical Results
In Section D.1, we present the matrix and vector properties. In Section D.2, we analyze and develop the tools for error analysis.
As , it must have at least basis, and we proved the first part.
Now, we prove the second part by math induction.
where the first step follows from the fact that is a lower triangular matrix and Definition 3.9, the second step follows from simple algebra, and the last step follows from simple algebra.
As and are lower triangular matrices, we have that is a lower triangular matrix. Thus, we proved the following statement.
D.2 Tools for Error Analysis
It is trivial by . ∎
where the first step follows simple algebra, and the last step follows triangle inequality.
For the first part, for any , we have
where the first step follows simple algebra, the second step follows triangle inequality, the third step follows simple algebra, the fourth step follows , , , , the fifth steps follows triangle inequality, the sixth step follows Lemma D.3 and the last step follows and .
For the second part, for any , we have
where the first step follows simple algebra, the second step follows triangle inequality, the third step follows , , the fourth step follows Lemma D.3, and the last step follows .
Let and . If , then
It is trivial by considering two cases when and . ∎
where the first step follows simple algebra, and the last step follows triangle inequality.
For the first part, for any , we have
where the first step follows simple algebra, the second step follows triangle inequality, the third step follows simple algebra, the fourth step follows , , the fifth step follows triangle inequality, the sixth step follows Lemma D.5 and the last step follows and .
For the second part, for any , we have
where the first step follows simple algebra, the second step follows triangle inequality, the third step follows Lemma D.5, and the last step follows .
D.3 Tensor Tools for Gradient Computation
where the first step follows from the definition of outer product, the second step follows from the definition of vectorization operator which stacks rows of a matrix into a column vector, and the last step follows from Definition 5.4. ∎
where the first step follows from that matrix can be written as a summation of vectors, the second step follows from Fact D.8, the third step follows from that matrix can be written as a summation of vectors, and the last step follows from the definition of vectorization operator . ∎
Appendix E More Related Work
introduced sparse factorizations of the attention matrix, which scale linearly with the sequence length. proposed a combination of local and global attention, where local attention captures short-range dependencies, and global attention captures long-range dependencies. introduced a random attention pattern that scales linearly with the sequence length while maintaining the ability to capture global dependencies. introduced a method called Performer, which uses random feature maps to approximate the attention matrix, resulting in linear complexity.
Other works have focused on studying the attention regression problem. The work by demonstrates that the forward computation of attention can be performed in sub-quadratic time complexity without explicitly constructing the full attention matrix. answer the question of how efficiently the gradient can be computed in the field of fast attention acceleration. There are many works simplifying the attention regression problem. For example, studies the softmax regression problem, proposes and analyzes the attention kernel regression problem, studies the rescaled hyperbolic functions regression, all of which fill the need of using iterative methods analyzing the attention regression related problems. After that, provides an un-simplified version of the attention regression problem, and construct and study two-layer regression problems related to attention.
Moreover, applies the attention computation to design a decentralized LLM. shows that attention layers in transformers learn 2-dimensional cosine functions, similar to how single-hidden layer neural networks learn Fourier-based circuits when solving modular arithmetic tasks. analyzes the data recovery by using attention weight. designs the quantum algorithm for solving the attention computation problem. analyzes the sparsification of the attention problem. applies the Zero-th Order method to approximate the gradient of the attention regression problem. analyzes the techniques for differentially private approximating the attention matrix. extends the work of to study the multiple softmax regression via the tensor trick, which is the same as our Fact D.9. analyzes the hardness for dynamic attention maintenance in large language models. studies how to balance the trade-off between creativity and reality using attention computation. formulates the , and regression problems from the original attention computation. replace the softmax unit by the polynomial unit. introduces an outlier-efficient modern Hopfield model and corresponding Hopfield layers that provide an alternative to the standard attention mechanism in transformer-based models.
Convolution computation.
Convolution computation has been a focal point of research in the field of deep learning and computer vision. Convolutional Neural Networks (CNNs) have demonstrated remarkable performance in various tasks such as image classification , object detection , and semantic segmentation . The efficiency and effectiveness of convolution computation play a crucial role in the success of CNNs.
One notable development in convolution computation is the introduction of depthwise separable convolutions, proposed by . This technique separates the convolution operation into two steps: depthwise convolution and pointwise convolution, which significantly reduces the computational cost and model size while maintaining comparable accuracy. This advancement has enabled the deployment of CNNs on resource-constrained devices such as mobile phones and embedded systems. Another important contribution is the development of strided convolutions, which allow for downsampling the spatial dimensions of feature maps without the need for additional pooling layers. This technique has been widely adopted in modern CNN architectures like ResNet and Inception , leading to more efficient and compact models.
To further improve the computational efficiency of convolutions, several works have explored the use of group convolutions. ResNeXt and ShuffleNet have demonstrated that by dividing the input channels into groups and performing convolutions within each group, the computational cost can be significantly reduced while maintaining or even improving the model’s performance. In addition to architectural innovations, there have been advancements in the implementation of convolution computation on hardware. The introduction of specialized hardware such as GPUs and TPUs has greatly accelerated the training and inference of CNNs . Recent research has also focused on reducing the computational cost of convolutions through network pruning and quantization techniques. proposed a pruning method that removes unimportant connections in the network, resulting in sparse convolutions that require less computation. Quantization techniques, such as those introduced by , reduce the precision of weights and activations, allowing for faster convolutions with minimal loss in accuracy.