通过预条件子求解注意力核回归问题

注意力机制是大语言模型的关键，而注意力矩阵是此类方案的算法与计算瓶颈。在本文中，我们定义了两个问题，其动机在于为注意力矩阵的代理设计快速算法并求解相应的回归。给定输入矩阵 $A\in \mathbb{R}^{n\times d}$（其中 $n\gg d$）和响应向量 $b$，我们首先考虑矩阵 $A^\top A$ 的矩阵指数作为代理，进而针对两类回归问题设计算法：对任意正整数 $j$，求解 $\min_{x\in \mathbb{R}^d}\|(A^\top A)^jx-b\|_2$ 与 $\min_{x\in \mathbb{R}^d}\|A(A^\top A)^jx-b\|_2$。研究这些回归的算法至关重要，因为矩阵指数可通过这些较小问题逐项逼近。第二类代理是对 Gram 矩阵逐元素取指数，记为 $\exp(AA^\top)$，并求解回归 $\min_{x\in \mathbb{R}^n}\|\exp(AA^\top)x-b \|_2$。我们称此问题为注意力核回归问题，因为矩阵 $\exp(AA^\top)$ 可视为关于 $A$ 的核函数。我们基于草图法与预条件化，为这些回归问题设计了快速算法。我们希望这些努力能为研究注意力矩阵的高效逼近提供另一种视角。