Solving Attention Kernel Regression Problem via Pre-conditioner

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^{\mathrm{\top }}A$ as a proxy, and we in turn design algorithms for two types of regression problems: ${\min }_{x\in {\mathbb{R}}^d}\mathrm{\parallel }(A^{\mathrm{\top }}A{)}^{j}x-b{\mathrm{\parallel }}_2$ and ${\min }_{x\in {\mathbb{R}}^d}\mathrm{\parallel }A(A^{\mathrm{\top }}A{)}^{j}x-b{\mathrm{\parallel }}_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 (A{A}^{\mathrm{\top }})$ and solving the regression ${\min }_{x\in {\mathbb{R}}^n}\mathrm{\parallel }\exp (A{A}^{\mathrm{\top }})x-b{\mathrm{\parallel }}_2$. We call this problem the attention kernel regression problem, as the matrix $\exp (A{A}^{\mathrm{\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.