A Fast Optimization View: Reformulating Single Layer Attention in LLM Based on Tensor and SVM Trick, and Solving It in Matrix Multiplication Time

Abstract: 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}ⁿ \sum_{i_0 = 1}^d ( \langle \langle \exp( \mathsf{A}<em>{j_0} x ) , {\bf 1}_n \rangle^{-1} \exp( \mathsf{A}</em>{j_0} x ), A_{3} Y_{<em>,i_0} \rangle - b_{j_0,i_0} )<sup>2$$. Here $$\mathsf{A} \in \mathbb{R}^{n² \times d<sup>2}$$ is Kronecker product between $A_1 \in \mathbb{R}^{n \times d}$ and $A_2 \in \mathbb{R}^{n \times d}$. $A_3$ is a matrix in $\mathbb{R}^{n \times d}$, $$\mathsf{A}<em>{j_0} \in \mathbb{R}^{n \times d<sup>2}$$ is the $j_0$-th block of $\mathsf{A}$. The $X, Y \in \mathbb{R}^{d \times d}$ are variables we want to learn. $B \in \mathbb{R}^{n \times d}$ and $b</em>{j_0,i_0} \in \mathbb{R}$ is one entry at $j_0$-th row and $i_0$-th column of $B$, $Y_{</em>,i_0} \in \mathbb{R}<sup>d$ is the $i_0$-column vector of $Y$, and $x \in \mathbb{R}<sup>{d<sup>2}$ is the vectorization of $X$. In a multi-layer LLM network, the matrix $B \in \mathbb{R}^{n \times d}$ can be viewed as the output of a layer, and $$A_1= A_2 = A_3 \in \mathbb{R}^{n \times d}$$ can be viewed as the input of a layer. The matrix version of $x$ can be viewed as $QK<sup>\top$ and $Y$ can be viewed as $V$. We provide an iterative greedy algorithm to train loss function $L(X,Y)$ up $\epsilon$ that runs in $$\widetilde{O}( ({\cal T}<em>{\mathrm{mat}}(n,n,d) + {\cal T}</em>{\mathrm{mat}}(n,d,d) + d^{2\omega}) \log(1/\epsilon) )$$ time. Here ${\cal T}_{\mathrm{mat}}(a,b,c)$ denotes the time of multiplying $a \times b$ matrix another $b \times c$ matrix, and $\omega\approx 2.37$ denotes the exponent of matrix multiplication.