[論文レビュー] 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

ひとこと要約

この論文は、テンソル代数とSVMにインspiredされたスパース化技術を用いて注意計算を再定式化することで、大規模言語モデル（LLMs）における1層の自己注意機構の最適化フレームワークを提案する。行列乗算時間内に収束し、$;{O}((Τ_{ρατ}(n,n,d) + Τ_{ρατ}(n,d,d) + d^{2\omega})\log(1/\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.

研究の動機と目的

• LLMsにおける1層注意機構の証明可能な効率的最適化手法を提供すること。トレーニングおよび推論における計算ボトル neck を解決する。
• テンソル代数とクロネッカー積を用いて注意の目的関数を再定式化し、高速な計算を可能にする。
• 証明可能な収束保証を有するグリーディな反復アルゴリズムを活用し、行列乗算時間内に収束を達成する。
• 高度なスパース化技術を用いて、注意回帰の計算複雑度を二次関数的から近似的最適な行列乗算時間に低減する。
• $(l,M)$-良い損失関数と適切に初期化された点を用いた、注意機構最適化の理論的枠組みを確立する。

提案手法

• クロネッカー積を用いて1層注意の目的関数を行列最適化問題に再定式化する：$\mathsf{A} = A_1 \otimes A_2 \in \mathbb{R}^{n^2 \times d^2}$ により、テンソルベースの計算を可能にする。
• TensorSketchによるSVMのテクニックのテンソル化を適用し、注意計算をスパース化することで、次元の有効縮小を実現しつつ精度を保持する。
• 近似ヘッセ行列の逆行列を用いたグリーディな反復アルゴリズムにより、パラメータ $X$ と $Y$ を更新し、$(l,M)$-良い損失条件の下で収束を保証する。
• 行列乗算に基づくフレームワークを採用し、時間計算量が $\mathcal{T}_{\mathrm{mat}}(a,b,c)$ で抑えられることを保証する。これは $a\times b$ と $b\times c$ の行列の乗算に要する時間である。
• 初期点の条件として $r_0 M \leq 0.1l$ を導入し、ヘッセ行列が正定値である最小値への局所収束を保証する。
• 補題14.3による反復的縮小を用い、最適解からの距離が各ステップで幾何的に減少することを保証する。$r_{t+1} \leq 0.4 r_t$。

実験結果

リサーチクエスチョン

• RQ1 LLMsにおける1層注意最適化問題は、テンソル代数を用いて再定式化可能であり、これにより高速計算が可能になるか？
• RQ2 SVMにインスパイアされたスパース化技術は、注意計算の複雑度を低減しつつ精度を保持するように適応可能か？
• RQ3 注意回帰問題に対して、証明可能な保証を伴う行列乗算時間内での収束は達成可能か？
• RQ4 損失関数にどのような条件が課されると、注意機構の最適化において安定的かつ高速な収束が保証されるか？
• RQ5 アルゴリズムの再定式化とスパース化により、注意層の計算ボトル neck をどのように軽減できるか？

主な発見

• 提案されたアルゴリズムは、$\widetilde{O}((\mathcal{T}_{\mathrm{mat}}(n,n,d) + \mathcal{T}_{\mathrm{mat}}(n,d,d) + d^{2\omega})\log(1/\epsilon))$ 時間で1層注意最適化問題を解き、$\epsilon$-精度を達成する。
• 標準的な仮定の下で損失関数が $(l,M)$-良いことが証明され、最小値におけるヘッセ行列が正定値であることを保証する。
• アルゴリズムは幾何的収束を保証する：$r_{t+1} \leq 0.4 r_t$ であり、反復中常に $M \cdot r_t \leq 0.1l$ が維持される。
• TensorSketchに基づくスパース化の使用により、$n^2 \times n^2$ の全行列の明示的計算を回避でき、メモリと時間コストを削減する。
• 標準的な注意計算の依存を減らすことで、従来の注意計算より証明可能な高速化を実現する。
• この手法は標準的なLLMのコンponentsに適用可能である：$X$ は $QK^T$ に対応し、$Y$ は $V$ に対応し、$B$ はターゲット出力に対応する。したがって、Transformerのトレーニングに直接関連する。