The Closeness of In-Context Learning and Weight Shifting for Softmax Regression

2023/04/26 by Shuai Li, Li, Shuai, Zhao Song +7 · 3 citations
Computer Science · #Computation and Language (cs.CL) #Explainable Artificial Intelligence (XAI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning and Data Classification #Topic Modeling

paper · pdf · doi:10.48550/arxiv.2304.13276

openalex publication_date 2023/04/26 · openalex created_date 2023/04/28 · openalex updated_date 2026/07/28

Abstract

Large language models (LLMs) are known for their exceptional performance in natural language processing, making them highly effective in many human life-related or even job-related tasks. The attention mechanism in the Transformer architecture is a critical component of LLMs, as it allows the model to selectively focus on specific input parts. The softmax unit, which is a key part of the attention mechanism, normalizes the attention scores. Hence, the performance of LLMs in various NLP tasks depends significantly on the crucial role played by the attention mechanism with the softmax unit. In-context learning, as one of the celebrated abilities of recent LLMs, is an important concept in querying LLMs such as ChatGPT. Without further parameter updates, Transformers can learn to predict based on few in-context examples. However, the reason why Transformers becomes in-context learners is not well understood. Recently, several works [ASA+22,GTLV22,ONR+22] have studied the in-context learning from a mathematical perspective based on a linear regression formulation min_x‖ Ax - b ‖₂, which show Transformers' capability of learning linear functions in context. In this work, we study the in-context learning based on a softmax regression formulation min_x ‖ ⟨ exp(Ax), \bf 1_n ⟩^-1 exp(Ax) - b ‖₂ of Transformer's attention mechanism. We show the upper bounds of the data transformations induced by a single self-attention layer and by gradient-descent on a ℓ₂ regression loss for softmax prediction function, which imply that when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient-descent and Transformers show great similarity.

Cited by

• Genomic Next-Token Predictors are In-Context Learners
• On Theoretical Interpretations of Concept-Based In-Context Learning
• DICE: Dynamic In-Context Example Selection in LLM Agents via Efficient Knowledge Transfer

Related

• Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descent
• In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness
• Transformers are Minimax Optimal Nonparametric In-Context Learners
• Self-Adjust Softmax
• Provable In-Context Learning of Nonlinear Regression with Transformers