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

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) = sumj0 = 1n sumi0 = 1d ( langle langle exp( mathsfAj0 x ) , bf 1n rangle-1 exp( mathsfAj0 x ), A3 Y*,i0 rangle - bj0,i0 )2. Here mathsfA ∈ ℝn2 × d2 is Kronecker product between A1 ∈ ℝn × d and A2 ∈ ℝn × d. A3 is a matrix in ℝn × d, mathsfAj0 ∈ ℝn × d2 is the j0-th block of mathsfA. The X, Y ∈ ℝd × d are variables we want to learn. B ∈ ℝn × d and bj0,i0 ∈ ℝ is one entry at j0-th row and i0-th column of B, Y*,i0 ∈ ℝd is the i0-column vector of Y, and x ∈ ℝd2 is the vectorization of X. In a multi-layer LLM network, the matrix B ∈ ℝn × d can be viewed as the output of a layer, and A1= A2 = A3 ∈ ℝn × d can be viewed as the input of a layer. The matrix version of x can be viewed as QK^top and Y can be viewed as V. We provide an iterative greedy algorithm to train loss function L(X,Y) up ε that runs in widetildeO( (cal Tmat(n,n,d) + cal Tmat(n,d,d) + d2ω) log(1/ε) ) time. Here cal Tmat(a,b,c) denotes the time of multiplying a × b matrix another b × c matrix, and ω≈ 2.37 denotes the exponent of matrix multiplication.