Solving Attention Kernel Regression Problem via Pre-conditioner

TL;DR

The paper tackles the computational bottleneck of attention in large models by introducing proxy problems that approximate the attention matrix: a matrix-exponential proxy using powers of $A^\top A$ and an attention-kernel proxy via $\exp(AA^\top)$. It develops fast, randomized algorithms based on sketching and preconditioning to solve regression problems against these proxies, achieving near-linear time in $n$ and $d$ under practical regimes, with provable error bounds and high-probability guarantees. The core contributions include (i) efficient algorithms for regressing against $(A^\top A)^j$ and $A(A^\top A)^j$ with explicit time bounds, (ii) a method to approximate $\exp(AA^\top)$ via spectral-approximate sketches and a preconditioner to enable fast gradient-based solves, and (iii) a thorough analysis relating these subproblems to matrix-exponential expansions, offering an alternative perspective on efficient attention-approximation techniques. These results have potential practical impact for speeding up attention-related computations in transformers and other attention-based architectures, especially in large-scale or structured-data settings.

Abstract

The attention mechanism is the key to large language models, and the attention matrix serves as an algorithmic and computational bottleneck for such a scheme. In this paper, we define two problems, motivated by designing fast algorithms for proxy of attention matrix and solving regressions against them. Given an input matrix $A\in \mathbb{R}^{n\times d}$ with $n\gg d$ and a response vector $b$, we first consider the matrix exponential of the matrix $A^\top A$ as a proxy, and we in turn design algorithms for two types of regression problems: $\min_{x\in \mathbb{R}^d}\|(A^\top A)^jx-b\|_2$ and $\min_{x\in \mathbb{R}^d}\|A(A^\top A)^jx-b\|_2$ for any positive integer $j$. Studying algorithms for these regressions is essential, as matrix exponential can be approximated term-by-term via these smaller problems. The second proxy is applying exponential entrywise to the Gram matrix, denoted by $\exp(AA^\top)$ and solving the regression $\min_{x\in \mathbb{R}^n}\|\exp(AA^\top)x-b \|_2$. We call this problem the attention kernel regression problem, as the matrix $\exp(AA^\top)$ could be viewed as a kernel function with respect to $A$. We design fast algorithms for these regression problems, based on sketching and preconditioning. We hope these efforts will provide an alternative perspective of studying efficient approximation of attention matrices.

Figures (7)

• Figure 1: The visualization of the matrix $D(X) \in \mathbb{R}^{n \times n}$. Given $Q, K , V \in \mathbb{R}^{d \times d}$ and $X \in \mathbb{R}^{n \times d}$, we first compute $XQK^\top X^\top \in \mathbb{R}^{n \times n}$. Then, we find $\exp(XQK^\top X^\top) \in \mathbb{R}^{n \times n}$. After that, we multiply $\exp(XQK^\top X^\top) \in \mathbb{R}^{n \times n}$ with the vector ${\bf 1}_n \in \mathbb{R}^{n}$. Finally, we use $\mathop{\mathrm{diag}}\nolimits(\cdot)$ to transform $\exp(XQK^\top X^\top) {\bf 1}_n \in \mathbb{R}^n$ into a diagonal matrix, which is $D(X) \in \mathbb{R}^{n \times n}$. In this figure, green matrices/vectors represent the terms that are given; the purple matrix represents the term after one operation; the red vector represents the term after two operations; the blue matrix represents the term after three operations.
• Figure 2: The visualization of the attention computation (see Eq. \ref{['eq:attention']}). Since we present the visualization of how we get $D(X) \in \mathbb{R}^{n \times n}$ and $\exp(XQK^\top X^\top) \in \mathbb{R}^{n \times n}$ in Figure \ref{['fig:DX']}, we regard them as given. Moreover, we are also given $V \in \mathbb{R}^{d \times d}$ and $X \in \mathbb{R}^{n \times d}$. We compute their product, namely $D(X)^{-1} \exp(X Q K^\top X^\top) X V$. In this figure, green matrices represent the terms that are given, and the purple matrix represents the term after one operation.
• Figure 3: The visualization of the matrix $D \in \mathbb{R}^{n \times n}$. Given $Q, K, V \in \mathbb{R}^{n \times d}$, we first compute $QK^\top \in \mathbb{R}^{n \times n}$. Then, we find $\exp(QK^\top) \in \mathbb{R}^{n \times n}$. After that, we multiply $\exp(QK^\top) \in \mathbb{R}^{n \times n}$ with the vector ${\bf 1}_n \in \mathbb{R}^n$. Finally, we use $\mathop{\mathrm{diag}}\nolimits(\cdot)$ to transform $\exp(QK^\top) {\bf 1}_n \in \mathbb{R}^n$ into a diagonal matrix, which is $D \in \mathbb{R}^{n \times n}$. In this figure, green matrices/vectors represent the terms that are given; the purple matrix represents the term after one operation; the red vector represents the term after two operations; the blue matrix represents the term after three operations.
• Figure 4: The visualization of the simplified version of attention computation in as23bsz23 (see Eq. \ref{['eq:attention_in_AS23_BSZ23']}). Since we present the visualization of how we get $D \in \mathbb{R}^{n \times n}$ and $\exp(QK^\top) \in \mathbb{R}^{n \times n}$ in Figure \ref{['fig:D']}, we regard them as given. Moreover, we are also given $V \in \mathbb{R}^{n \times d}$. We compute their product, namely $D^{-1} \exp(Q K^\top) V \in \mathbb{R}^{n \times d}$. In this figure, green matrices represent the terms that are given, and the purple matrix represents the term after one operation.
• Figure 5: The visualization of the simplified version of attention computation in gsyz23_quantum (see Eq. \ref{['eq:attention_in_gsyz23a']}). Since we present the visualization of how we get $D \in \mathbb{R}^{n \times n}$ and $\exp(QK^\top) \in \mathbb{R}^{n \times n}$ in Figure \ref{['fig:D']}, we regard them as given. We compute their product, namely $D^{-1} \exp(Q K^\top) \in \mathbb{R}^{n \times n}$. In this figure, green matrices represent the terms that are given, and the purple matrix represents the term after one operation.

Theorems & Definitions (55)

• Definition 1.1
• Definition 1.2
• Definition 1.3: Attention Kernel Regression (or Exponential Regression)
• Theorem 1.4: Informal Version of Theorem \ref{['thm:even']}
• Theorem 1.5: Informal Version of Theorem \ref{['thm:odd']}
• Theorem 1.6: Informal Version of Theorem \ref{['thm:formal_exp']}
• Theorem 4.1: Main Result for Matrix Exponential Proxy and Even/Odd Power Regression
• Theorem 4.2: Main Result for Attention Kernel Regression
• Definition A.1
• Definition A.2: Hadamard matrix