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

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(⟨⟨\exp ({\mathsf{A}}_{j_0}x),1_n{⟩}^{-1}\exp ({\mathsf{A}}_{j_0}x),A_3{Y}_{\ast ,i_0}⟩-b_{j_0,i_0}{)}^2$. Here $\mathsf{A}\in {\mathbb{R}}^{n^2\times d^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}}_{j_0}\in {\mathbb{R}}^{n\times d^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_{j_0,i_0}\in \mathbb{R}$ is one entry at $j_0$-th row and $i_0$-th column of $B$, $Y_{\ast ,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\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 $Q{K}^{\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 $\tilde{O}(({\mathcal{T}}_{\mathrm{mat}}(n,n,d)+{\mathcal{T}}_{\mathrm{mat}}(n,d,d)+d^{2\omega })\log (1/\epsilon ))$ time. Here ${\mathcal{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.