[论文解读] A Fast Optimization View: Reformulating Single Layer Attention in LLM Based on Tensor and SVM Trick, and Solving It in Matrix Multiplication Time

Yeqi Gao, Zhao Song+2 more|arXiv (Cornell University)|Sep 14, 2023

Topic ModelingComputer Science被引用 3

一句话总结

本文通过使用张量代数和受支持向量机启发的稀疏化技术，重新表述了大规模语言模型（LLMs）中单层自注意力机制的注意力计算，提出了一种新颖的优化框架。该方法在矩阵乘法时间内实现收敛，以可证明的保证在 $\widetilde{O}((\mathcal{T}_{\mathrm{mat}}(n,n,d) + \mathcal{T}_{\mathrm{mat}}(n,d,d) + d^{2\omega})\log(1/\epsilon))$ 时间内求解注意力回归问题，达到 $\epsilon$-精度。

ABSTRACT

Large language models (LLMs) have played a pivotal role in revolutionizing various facets of our daily existence. Solving attention regression is a fundamental task in optimizing LLMs. In this work, we focus on giving a provable guarantee for the one-layer attention network objective function $L(X,Y) = \sum_{j_0 = 1}^n \sum_{i_0 = 1}^d ( \langle \langle \exp( \mathsf{A}_{j_0} x ) , {\bf 1}_n angle^{-1} \exp( \mathsf{A}_{j_0} x ), A_{3} Y_{*,i_0} angle - b_{j_0,i_0} )^2$. Here $\mathsf{A} \in \mathbb{R}^{n^2 imes d^2}$ is Kronecker product between $A_1 \in \mathbb{R}^{n imes d}$ and $A_2 \in \mathbb{R}^{n imes d}$. $A_3$ is a matrix in $\mathbb{R}^{n imes d}$, $\mathsf{A}_{j_0} \in \mathbb{R}^{n imes d^2}$ is the $j_0$-th block of $\mathsf{A}$. The $X, Y \in \mathbb{R}^{d imes d}$ are variables we want to learn. $B \in \mathbb{R}^{n imes d}$ and $b_{j_0,i_0} \in \mathbb{R}$ is one entry at $j_0$-th row and $i_0$-th column of $B$, $Y_{*,i_0} \in \mathbb{R}^d$ is the $i_0$-column vector of $Y$, and $x \in \mathbb{R}^{d^2}$ is the vectorization of $X$. In a multi-layer LLM network, the matrix $B \in \mathbb{R}^{n imes d}$ can be viewed as the output of a layer, and $A_1= A_2 = A_3 \in \mathbb{R}^{n imes d}$ can be viewed as the input of a layer. The matrix version of $x$ can be viewed as $QK^ op$ and $Y$ can be viewed as $V$. We provide an iterative greedy algorithm to train loss function $L(X,Y)$ up $ε$ that runs in $\widetilde{O}( ({\cal T}_{\mathrm{mat}}(n,n,d) + {\cal T}_{\mathrm{mat}}(n,d,d) + d^{2ω}) \log(1/ε) )$ time. Here ${\cal T}_{\mathrm{mat}}(a,b,c)$ denotes the time of multiplying $a imes b$ matrix another $b imes c$ matrix, and $ω\approx 2.37$ denotes the exponent of matrix multiplication.

研究动机与目标

• 为大规模语言模型中的单层注意力提供一种可证明高效的优化方法，解决训练和推理中的计算瓶颈。
• 通过张量代数和克罗内克积重新表述注意力目标，以实现更快的计算。
• 通过利用具有可证明收敛保证的贪心迭代算法，实现在矩阵乘法时间内的收敛。
• 通过使用先进的稀疏化技术，将注意力回归的计算复杂度从二次方降低到接近最优的矩阵乘法时间。
• 为具有 $(l,M)$-良好损失函数和良好初始化点的注意力机制优化建立理论框架，以确保稳定收敛。

提出的方法

• 使用克罗内克积将单层注意力目标重新表述为矩阵优化问题：$\mathsf{A} = A_1 \otimes A_2 \in \mathbb{R}^{n^2 \times d^2}$，从而支持基于张量的计算。
• 通过 TensorSketch 的张量化版本应用受支持向量机启发的稀疏化技术，对注意力计算进行稀疏化处理，降低有效维度的同时保持精度。
• 使用一种基于近似 Hessian 逆的贪心迭代算法来更新参数 $X$ 和 $Y$，在 $(l,M)$-良好损失条件下确保收敛。
• 采用基于矩阵乘法的框架，其中时间复杂度受 $\mathcal{T}_{\mathrm{mat}}(a,b,c)$ 限制，即相乘 $a\times b$ 和 $b\times c$ 矩阵所需的时间。
• 引入良好初始化点条件，其中 $r_0 M \leq 0.1l$，以确保在正定 Hessian 矩阵下实现局部收敛。
• 通过引理 14.3 的迭代收缩，保证每一步到最优解的距离几何级递减，满足 $r_{t+1} \leq 0.4 r_t$。

实验结果

研究问题

• RQ1 是否可以通过张量代数重新表述大规模语言模型中单层注意力优化问题，以实现更快的计算？
• RQ2 受支持向量机启发的稀疏化技术是否可以被适配以降低注意力计算的复杂度，同时保持精度？
• RQ3 是否可能在具有可证明保证的前提下，实现在矩阵乘法时间内求解注意力回归问题？
• RQ4 损失函数需要满足何种条件，才能确保注意力机制优化过程的稳定且快速收敛？
• RQ5 如何通过算法重构和稀疏化缓解注意力层的计算瓶颈？

主要发现

• 所提出的算法在 $\widetilde{O}((\mathcal{T}_{\mathrm{mat}}(n,n,d) + \mathcal{T}_{\mathrm{mat}}(n,d,d) + d^{2\omega})\log(1/\epsilon))$ 时间内求解单层注意力优化问题，达到 $\epsilon$-精度。
• 在标准假设下，损失函数被证明为 $(l,M)$-良好，确保在最小值处具有正定 Hessian 矩阵，从而实现局部收敛。
• 该算法保证几何收敛：$r_{t+1} \leq 0.4 r_t$，且在整个迭代过程中保持 $M \cdot r_t \leq 0.1l$。
• 采用基于 TensorSketch 的稀疏化技术，使方法能够避免显式计算完整的 $n^2 \times n^2$ 矩阵，从而降低内存和时间开销。
• 该框架通过减少对昂贵的 softmax 和矩阵求逆操作的依赖，实现了相对于标准注意力计算的可证明加速。
• 该方法可直接应用于标准 LLM 组件：$X$ 对应 $QK^T$，$Y$ 对应 $V$，$B$ 对应目标输出，因此与 Transformer 训练具有直接相关性。