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 providing a provable guarantee for the one-layer attention network objective function: given input matrices of a layer, A_1, A_2, A_3 ∈ℝ^n × d, our goal is to minimize the loss function: L(X,Y) = ∑_j_0=1^n ∑_i_0=1^d ( ⟨⟨exp( 𝖠_j_0 x ) , 1_n ⟩^-1·exp( 𝖠_j_0 x ), A_3 Y_,i_0 ⟩- b_j_0,i_0 )^2, where 𝖠_j_0∈ℝ^n × d^2 is the j_0-th block of the Kronecker product of A_1 and A_2. The matrix B ∈ℝ^n × d represents the output of a layer, and b_j_0,i_0∈ℝ is the (j_0, i_0)-th entry of B. Y_*,i_0∈ℝ^d is the i_0-th column vector of Y, and x ∈ℝ^d^2 is the vectorization of X. In self-attention, Q, K, V ∈ℝ^d × d represent the weight matrices for the query, key, and value, respectively. Unlike prior works that relied on simplified and single-variable versions of the attention computation problem, our multivariate loss function analyzes a complete and unsimplified attention layer, treating all these weights as variables, where X = QK^⊤∈ℝ^d × d and Y = V ∈ℝ^d × d. We propose an iterative greedy algorithm to train a neural network using the loss function L(X,Y), achieving an error bound of ϵ∈ (0, 0.1). The algorithm runs in O( ( T_mat(n,n,d) + T_mat(n,d,d) + d^2ω) log(1/ϵ) ) time, where T_mat(a,b,c) denotes the time complexity of multiplying an a × b matrix with a b × c matrix, and ω≈ 2.37 is the exponent of matrix multiplication.