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

• URL: http://arxiv.org/abs/2405.05219v1
• Date: Wed, 8 May 2024 17:11:38 GMT
• Title: Conv-Basis: A New Paradigm for Efficient Attention Inference and Gradient Computation in Transformers
• Authors: Jiuxiang Gu, Yingyu Liang, Heshan Liu, Zhenmei Shi, Zhao Song, Junze Yin,
• Abstract summary: We leverage the convolution-like structure of attention matrices to develop an efficient approximation method for attention using convolution matrices. Our algorithm achieves almost linear time, i.e., $n1+o(1)$. This work provides a new paradigm for accelerating attention computation in transformers to enable their application to longer contexts.
• Score: 27.54512534985192
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• Abstract: Large Language Models (LLMs) have profoundly changed the world. Their self-attention mechanism is the key to the success of transformers in LLMs. However, the quadratic computational cost $O(n^2)$ to the length $n$ input sequence is the notorious obstacle for further improvement and scalability in the longer context. 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, "similar" to the rank basis, and show that any lower triangular (attention) matrix can always be decomposed as a sum of $k$ structured convolution matrices in this basis system. We then design an algorithm to quickly decompose the attention matrix into $k$ convolution matrices. Thanks to Fast Fourier Transforms (FFT), the attention {\it inference} can be computed in $O(knd \log n)$ time, where $d$ is the hidden dimension. In practice, we have $ d \ll n$, i.e., $d=3,072$ and $n=1,000,000$ for Gemma. Thus, when $kd = n^{o(1)}$, our algorithm achieve almost linear time, i.e., $n^{1+o(1)}$. Furthermore, the attention {\it training forward} and {\it backward gradient} can be computed in $n^{1+o(1)}$ as well. Our approach can avoid explicitly computing the $n \times n$ attention matrix, which may largely alleviate the quadratic computational complexity. Furthermore, our algorithm works on any input matrices. This work provides a new paradigm for accelerating attention computation in transformers to enable their application to longer contexts.

Related papers

• Provably learning a multi-head attention layer [55.2904547651831]
  Multi-head attention layer is one of the key components of the transformer architecture that sets it apart from traditional feed-forward models. In this work, we initiate the study of provably learning a multi-head attention layer from random examples. We prove computational lower bounds showing that in the worst case, exponential dependence on $m$ is unavoidable.
  arXiv  Detail & Related papers  (2024-02-06T15:39:09Z)
• Structured Semidefinite Programming for Recovering Structured Preconditioners [41.28701750733703]
  We give an algorithm which, given positive definite $mathbfK in mathbbRd times d$ with $mathrmnnz(mathbfK)$ nonzero entries, computes an $epsilon$-optimal diagonal preconditioner in time. We attain our results via new algorithms for a class of semidefinite programs we call matrix-dictionary approximation SDPs.
  arXiv  Detail & Related papers  (2023-10-27T16:54:29Z)
• How to Capture Higher-order Correlations? Generalizing Matrix Softmax Attention to Kronecker Computation [12.853829771559916]
  We study a generalization of attention which captures triple-wise correlations. This generalization is able to solve problems about detecting triple-wise connections that were shown to be impossible for transformers. We show that our construction, algorithms, and lower bounds naturally generalize to higher-order tensors and correlations.
  arXiv  Detail & Related papers  (2023-10-06T07:42:39Z)
• Solving Attention Kernel Regression Problem via Pre-conditioner [9.131385887605935]
  We design algorithms for two types of regression problems: $min_xin mathbbRd|(Atop A)jx-b|$ for any positive integer $j$. The second proxy is applying exponential entrywise to the Gram matrix, denoted by $exp(AAtop)$ and solving the regression $min_xin mathbbRn|exp(AAtop)xb |$.
  arXiv  Detail & Related papers  (2023-08-28T04:37:38Z)
• Fast $(1+\varepsilon)$-Approximation Algorithms for Binary Matrix Factorization [54.29685789885059]
  We introduce efficient $(1+varepsilon)$-approximation algorithms for the binary matrix factorization (BMF) problem. The goal is to approximate $mathbfA$ as a product of low-rank factors. Our techniques generalize to other common variants of the BMF problem.
  arXiv  Detail & Related papers  (2023-06-02T18:55:27Z)
• Fast Attention Requires Bounded Entries [19.17278873525312]
  inner product attention computation is a fundamental task for training large language models such as Transformer, GPT-1, BERT, GPT-2, GPT-3 and ChatGPT. We investigate whether faster algorithms are possible by implicitly making use of the matrix $A$. This gives a theoretical explanation for the phenomenon observed in practice that attention computation is much more efficient when the input matrices have smaller entries.
  arXiv  Detail & Related papers  (2023-02-26T02:42:39Z)
• Sketching Algorithms and Lower Bounds for Ridge Regression [65.0720777731368]
  We give a sketching-based iterative algorithm that computes $1+varepsilon$ approximate solutions for the ridge regression problem. We also show that this algorithm can be used to give faster algorithms for kernel ridge regression.
  arXiv  Detail & Related papers  (2022-04-13T22:18:47Z)
• Learning a Latent Simplex in Input-Sparsity Time [58.30321592603066]
  We consider the problem of learning a latent $k$-vertex simplex $KsubsetmathbbRdtimes n$, given access to $AinmathbbRdtimes n$. We show that the dependence on $k$ in the running time is unnecessary given a natural assumption about the mass of the top $k$ singular values of $A$.
  arXiv  Detail & Related papers  (2021-05-17T16:40:48Z)
• Quantum algorithms for spectral sums [50.045011844765185]
  We propose new quantum algorithms for estimating spectral sums of positive semi-definite (PSD) matrices. We show how the algorithms and techniques used in this work can be applied to three problems in spectral graph theory.
  arXiv  Detail & Related papers  (2020-11-12T16:29:45Z)
• Constant-Depth and Subcubic-Size Threshold Circuits for Matrix Multiplication [1.9518237361775532]
  Recent advances in large-scale neural computing hardware has made their practical implementation a near-term possibility. We describe a theoretical approach for multiplying two $N$ by $N$ matrices that integrates threshold gate logic. Dense matrix multiplication is a core operation in convolutional neural network training.
  arXiv  Detail & Related papers  (2020-06-25T18:28:10Z)
• Sketching Transformed Matrices with Applications to Natural Language Processing [76.6222695417524]
  We propose a space-efficient sketching algorithm for computing the product of a given small matrix with the transformed matrix. We show that our approach obtains small error and is efficient in both space and time.
  arXiv  Detail & Related papers  (2020-02-23T03:07:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.