快速优化视角：基于张量与 SVM 技巧重构 LLM 中的单层注意力并以其矩阵乘法时间求解

大语言模型（LLMs）在变革我们日常生活的诸多方面发挥了关键作用。求解注意力回归是优化 LLM 的一项基本任务。在本工作中，我们致力于为单层注意力网络目标函数 $L(X,Y) = \sum_{j_0 = 1}^n \sum_{i_0 = 1}^d ( \langle \langle \exp( \mathsf{A}_{j_0} x ) , {\bf 1}_n \rangle^{-1} \exp( \mathsf{A}_{j_0} x ), A_{3} Y_{*,i_0} \rangle - b_{j_0,i_0} )^2$ 提供可证明的保证。此处 $\mathsf{A} \in \mathbb{R}^{n^2 \times d^2}$ 是 $A_1 \in \mathbb{R}^{n \times d}$ 与 $A_2 \in \mathbb{R}^{n \times d}$ 间的克罗内克积。$A_3$ 是 $\mathbb{R}^{n \times d}$ 中的矩阵，$\mathsf{A}_{j_0} \in \mathbb{R}^{n \times d^2}$ 是 $\mathsf{A}$ 的第 $j_0$ 块。$X, Y \in \mathbb{R}^{d \times d}$ 为我们欲学习的变量。$B \in \mathbb{R}^{n \times d}$ 且 $b_{j_0,i_0} \in \mathbb{R}$ 是 $B$ 的第 $j_0$ 行第 $i_0$ 列元素，$Y_{*,i_0} \in \mathbb{R}^d$ 是 $Y$ 的第 $i_0$ 列向量，$x \in \mathbb{R}^{d^2}$ 是 $X$ 的向量化。在多层的 LLM 网络中，矩阵 $B \in \mathbb{R}^{n \times d}$ 可视为某层的输出，$A_1= A_2 = A_3 \in \mathbb{R}^{n \times d}$ 可视为某层的输入。$x$ 的矩阵形式可视为 $QK^\top$，$Y$ 可视为 $V$。我们提供一种迭代贪婪算法，以 $\widetilde{O}( ({\cal T}_{\mathrm{mat}}(n,n,d) + {\cal T}_{\mathrm{mat}}(n,d,d) + d^{2\omega}) \log(1/\epsilon) )$ 时间将损失函数 $L(X,Y)$ 训练至 $\epsilon$ 误差内。此处 ${\cal T}_{\mathrm{mat}}(a,b,c)$ 表示将一个 $a \times b$ 矩阵与另一个 $b \times c$ 矩阵相乘的时间，$\omega\approx 2.37$ 表示矩阵乘法的指数。