In-Context Learning for Attention Scheme: from Single Softmax Regression to Multiple Softmax Regression via a Tensor Trick
share
Large language models (LLMs) have brought significant and transformative changes in human society. These models have demonstrated remarkable capabilities in natural language understanding and generation, leading to various advancements and impacts across several domains. We consider the in-context learning under two formulation for attention related regression in this work. Given matrices A_1 ∈ℝ^n × d, and A_2 ∈ℝ^n × d and B ∈ℝ^n × n, the purpose is to solve some certain optimization problems: Normalized version min_X D(X)^-1exp(A_1 X A_2^⊤) - B _F^2 and Rescaled version exp(A_1 X A_2^⊤) - D(X) · B _F^2. Here D(X) := diag( exp(A_1 X A_2^⊤) 1_n ). Our regression problem shares similarities with previous studies on softmax-related regression. Prior research has extensively investigated regression techniques related to softmax regression: Normalized version ⟨exp(Ax) , 1_n ⟩^-1exp(Ax) - b _2^2 and Resscaled version exp(Ax) - ⟨exp(Ax), 1_n ⟩ b _2^2 In contrast to previous approaches, we adopt a vectorization technique to address the regression problem in matrix formulation. This approach expands the dimension from d to d^2, resembling the formulation of the regression problem mentioned earlier. Upon completing the lipschitz analysis of our regression function, we have derived our main result concerning in-context learning.
READ FULL TEXT