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.

大型语言模型对世界产生了深远影响，其中自注意机制是 transformer 在大型语言模型中取得成功的关键。然而，长度为 n 的输入序列带来的二次计算复杂度 O (n^2) 一直是对提高和扩展更长上下文的一个难以克服的障碍。本文利用注意力矩阵的卷积结构开发了一种有效的注意力计算近似方法，提出了一个类似秩基系统的卷积基系统，表明任意下三角（注意力）矩阵总是可以在该基系统中分解为 k 个结构化卷积矩阵的和。我们然后设计了一个算法，快速将注意力矩阵分解为 k 个卷积矩阵。借助快速傅里叶变换（FFT），可以在 O (knd log n) 的时间内计算注意力推理，其中 d 是隐藏维度。实际上，我们有 d ≪ n，即 d = 3,072 和 n = 1,000,000 用于 Gemma。因此，当 kd = n^{o (1)} 时，我们的算法实现了几乎线性时间，即 n^{1+o (1)}。此外，注意力的训练前向和反向梯度也可以在 n^{1+o (1)} 内计算。我们的方法可以避免明确计算 n×n 的注意力矩阵，这可能会大大减轻二次计算复杂度。此外，我们的算法适用于任何输入矩阵。本研究提供了一种加速 transformer 注意力计算的新范式，以实现其对更长上下文的应用。