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∈ ℝn× d with ngg 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: minx∈ ℝd|(A^top A)jx-b|2 and minx∈ ℝ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 minx∈ ℝ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.