Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers (2025)

Abstract: The self-attention mechanism is key to the success of transformers in recent large language models (LLMs). However, the quadratic computational cost, O(n^2) , with respect to the input sequence length n poses a significant obstacle to further improvement and scalability in longer contexts.In this work, we leverage the convolution-like structure of attention matrices to develop an efficient approximation method for attention computation using convolution matrices. We propose a \mathsf{conv} basis system, analogous to the rank basis, and show that any lower triangular matrix can be decomposed as a sum of structured convolution matrices in this basis. We then design a fast algorithm to approximate the attention matrix using a sum of k convolution matrices. This enables us to compute attention during inference via Fast Fourier Transforms (FFT) in O(knd \log n) time, where d is the hidden dimension, achieving nearly linear time complexity, n^{1+o(1)} , in practical scenarios where kd = n^{o(1)} . Furthermore, both training forward and backward gradient computations can be performed in n^{1+o(1)} time as well.We provide theoretical guarantees on runtime and approximation error and conduct preliminary experiments to evaluate the effectiveness of our approach. We hope this new paradigm for accelerating attention computation in transformer models facilitates their application to longer contexts.

## introduction numerous notable large language models (llms) in natural language processing (nlp) have emerged over the past two years, including mistral @xcite , gemini @xcite , claude3 @xcite , gpt-4 @xcite , and llama3 (ai, 2024), among others. these models have profoundly impacted the world and are widely used in various domains, such as education @xcite , law @xcite , finance @xcite , bioinformatics @xcite , coding @xcite , and even creative writing @xcite , including top ai conference reviews @xcite . the key to the success of generative llms is the decoder-only transformer architecture introduced by @xcite . the transformer employs a self-attention mechanism, enabling the model to capture long-range dependencies in sequences. self-attention computes a weighted sum of input tokens, where the weights are determined by the similarity between each token pair. this allows the model to focus on relevant information across different parts of the sequence when generating output. however, the computational complexity of selfattention grows quadratically, o(n 2 ), with the input length n, limiting their scalability for long contexts. for example, gpt-4 @xcite , claude3 @xcite , and gemma @xcite support input lengths of 128k, 200k, and 1000k tokens, respectively. the complexity o(n 2 ) arises from computing the similarity between each pair of tokens, resulting in an n × n matrix. more specifically, let d be the hidden dimension, and let q, @xmath0. attention then requires computing the softmax function on qk ⊤ ∈ r n×n . although qk ⊤ has at most rank d, the softmax operation in attention can produce softmax(qk ⊤ ) ∈ r n×n with full rank. to overcome the computational obstacle of softmax(qk ⊤ ), many studies have proposed more efficient attention computation methods that scale gracefully with sequence length while maintaining the model's performance. @xcite shows that if all entries of qk ⊤ are bounded and @xmath1, softmax(qk ⊤ ) will be "close" to a low-rank matrix. they then present an algorithm that approximates attention computation in almost linear time. under the assumptions of uniform softmax column norms and sparsity, @xcite solves the attention computation in almost linear time by identifying large entries in the attention matrix and focusing only on them. another line of work @xcite @xcite finds that the attention pattern exhibits a convolutional-like (or "diagonalized") structure (see figure 1 (b) ). mathematically, a i,j ≈ a i ′ ,j ′ when @xmath2, where ij represents the positional distance between two tokens. this is related to the bag-of-words or n-gram concept, i.e., sequences of n adjacent symbols or words in nlp. furthermore, the convolutional-like structure connects to convolutional recurrent models @xcite , hyena hierarchy models @xcite , and structured state space models (ssms) such as mamba @xcite . more specifically, multiple convolution matrices can be used to approximate an attention matrix, following an intuition similar to low-rank approximation in terms of computational acceleration. notably, the product of a convolution matrix and a vector can be computed using the fast fourier transform (fft) in o(n log n) time, whereas the naive approach requires o(n 2 ) time (see details in figure 1 (a) ). therefore, it is natural to ask: can we exploit the convolutional structure to accelerate the attention computation? in this paper, we use multiple convolution matrices to efficiently approximate attention computation. informally speaking, we present the following results, which apply to any q, @xmath3. theorem 1.1 (main result, informal version of theorem 3.4). let ϵ > 0, @xmath4, and q, @xmath5. if qk ⊤ is ϵ-close in the ℓ ∞ norm to a matrix with a k-conv basis (definition 3.1), then the exact attention computation (definition 2.3) can be solved in o(knd log n) time via fft, with an error of at most o(ϵ). when @xmath6, our method achieves an almost linear runtime of n 1+o(1) . similar to lowrank approximation, we construct a conv basis system, analogous to the rank basis, and show that any lower triangular matrix @xmath7, where [n] = 1, 2, . . . , n (lemma 2.12 and theorem 3.3). then, algorithm 2 efficiently decomposes qk ⊤ into k convolution matrices when qk ⊤ satisfies certain non-degeneracy properties (see definition 3.1). finally, using fft, we achieve a time complexity of o(knd log n) to solve the task (algorithm 1 and theorem 3.4), whereas naive methods require o(n 2 d). thus, our algorithm enables attention inference in o(knd log n) without any parameter updates, such as retraining or fine-tuning. additionally, our theorems apply to accelerating attention training, requiring o(knd log n + nd 2 ) time for forward computation and o(knd 2 log n) time for backward gradient computation (theorem 4.4). we also conduct preliminary experiments to evaluate its effectiveness (section 5). furthermore, our technique extends the low-rank approximation of attention matrices @xcite to more general settings (theorem d.5). specifically, @xcite focuses only on attention approximation without an attention mask, whereas our approach applies to various types of attention masks, including the widely used causal attention mask (definition 2.2). this demonstrates the broad applicability of our analysis. our contributions. our contributions are summarized as follows: our motivation lies in addressing the limitations of prior methods that rely on restrictive assumptions. in contrast, we develop a theoretically grounded, assumption-light approach with provable guarantees. while our main focus is theoretical, we include experiments to demonstrate the method's practicality and ease of integration into real-world models like llama 3. detailed comparison with previous works. our results are beyond or comparable to the two brilliant previous works. (1) to ensure a small approximation error for the attention matrix, @xcite requires bounded entries and assumes @xmath8, while @xcite assumes uniform softmax column norms and sparsity. however, our algorithm guarantees a small approximation error without any of these assumptions (corollary 3.5). that is, our approach applies to any q, k, including unbounded matrices, dense matrices, and any hidden dimension d. (2) alman and song (2023) assumes @xmath9. in contrast, our algorithm achieves n 1+o(1) runtime as long as @xmath10, imposing far fewer restrictions on d. moreover, our time complexity ranges from n 1+o(1) to n 2-ω(1) depending on d, whereas alman and @xcite is limited to @xmath11. ## preliminaries and frobenius norm is ∥a∥ f := i,j a 2 i,j , where a ij is an entry at the i-th row and j-th column. t mat (n, d, k) is the time of an n × d matrix times a d × k matrix. ## conv approximation during inference in section 3.1, we introduce the basic definitions to support our algorithmic analysis in this section. in section 3.2, we present the binary search and recover k-conv algorithms and present their theoretical guarantees. in section 3.3, we provide the formal version of our main result. ## conv approximation for training we accelerate both forward and back propagation of attention. we first define attention optimization, which is also used in @xcite . ## experiments now, we show our experimental results for convolution attention computing in language models, offering empirical backing to our theoretical claims. setup. we use the latest llama3 8b instruct model 2 (ai, 2024) as our foundation, modifying its attention mechanism with our convolution-based approach using varying numbers of convolution bases k. we evaluate our method on the imdb dataset @xcite , which consists of labeled movie reviews. our assessments rely on two key metrics: (1) the relative difference between our final layer output y and the original model output y , i.e., ∥y -y ∥ 2 f /∥y ∥ 2 f ; and (2) classification accuracy. this dual approach enables us to assess both internal representations and the overall 2 https://huggingface.co/meta-llama/ meta-llama-3-8b-instruct predictive performance of our convolution-based attention compared to the standard mechanism. implementation details. to ensure a fair comparison and prevent memory issues, we set the model's context length to 2048 tokens and incrementally increased the number of conv bases k. when @xmath12, our convolution attention produces an identical output to the original attention mechanism. we use an instruction-based approach to evaluate generation accuracy, formatting our input as review: <@xmath13: is this review positive or negative? answer:. this technique allows us to systematically assess the performance of our convolution-based attention across various complexity levels while maintaining comparability with the original model. we randomly sample 5 sample groups, with 200 samples per group, and report the results averaged across each group. all evaluations are conducted by directly applying our approximation to the pretrained llama3 8b instruct model without any retraining or fine-tuning, isolating the effect of the attention modification. results. figure 3 (left) shows that the relative mse drops rapidly as k increases, even with small values like @xmath14, indicating convergence toward the original attention. the right plot shows accuracy also improves with k and matches the original at @xmath15, demonstrating that our method retains strong performance with lower computational costs. these results imply that our approach effectively approximates the original attention mechanism, offering a promising trade-off between accuracy and efficiency, particularly in resource-constrained scenarios. ## further discussion comparison with efficient attention methods. recent works such as linformer @xcite and longformer @xcite reduce attention complexity by imposing structural assumptions like low-rank projections or sparsity patterns on the attention matrix. in contrast, our approach leverages a convolutional basis decomposition, which does not rely on any low-rank hypothesis. the conv-basis framework is provably applicable to any lower triangular matrix, including masked attention scenarios. importantly, our theoretical results (theorem 3.4 and corollary 3.5) guarantee that the approximation error is tightly controlled and can be made arbitrarily small by increasing k, with provable bounds on ∥y -y ∥ ∞ . longlora. our conv and low-rank approximation can be applied to longlora @xcite , whose mask is shown in the left of figure 4. 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 8× a100 machine. as the "diagonalized" mask structure, we can directly apply our algorithm 1 by replacing the causal attention mask (definition 2.2) with their sparse mask for the conv approximation with time complexity o(knd log(n)). similarly, for the low-rank approximation, we directly use the second statement in theorem d.5 by considering row change by amortized constant mask defined in definition d.1 with time complexity o(knd), where b @xmath16. rope. the rotary position embedding (rope) @xcite designs a rotation matrix r (m) ∈ r d×d , for all @xmath17, which can effectively encode the positional information into embedding q, @xmath18. in detail, let q i , k @xmath19, where q ⊤ i and k ⊤ j be the i-th and j-th row of q, k respectively, for any i, @xmath20. by the property of rotation matrix, we have we define q ′ , k ′ ∈ r n×d , and let q ′ i , k ′ @xmath21, where q ′ ⊤ i and k ′ ⊤ j be the i-th and j-th row of q ′ , k ′ respectively, for any i, @xmath22, we know that we can get q ′ , k ′ in o(nd) time. thus, we can apply q ′ , k ′ in our theorem 3.4 and theorem d.5 to get the same approximation error guarantee and the same time complexity. extend to full self-attention. we can easily extend our method to full self-attention. our proposed approach can be extended to accelerate full self-attention as well, not just the causal attention mechanism. note that the full self-attention matrix can be split into a lower triangular matrix l and an upper triangular matrix u . then, our conv-basis approximation method can be applied separately to l and the transpose of u . this allows the algorithm to handle both the lower and upper triangular components of the full attention matrix. the diagonal normalization step d -1 would need to be adjusted to account for the full matrix rather than just the lower triangular portion. finally, we combine the approximations of l and u ⊤ to reconstruct the full self-attention output. ## conclusion we introduce a novel method for efficient attention computation in transformers by leveraging a structured convolution matrix decomposition. unlike previous approaches that often rely on restrictive assumptions or focus exclusively on approximating specific types of attention, our algorithm provides a general, assumption-light framework with provable guarantees. it achieves nearly linear time complexity for both attention inference and gradient computation, offering stronger theoretical assurances than existing methods. this work establishes a new paradigm for accelerating attention operations, enabling transformers to handle longer contexts more efficiently and paving the way for further advancements in large language models.