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

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 01:11:55

Abstract: Large language models (LLMs) are known for their exceptional performance in natural language processing, making them highly effective in many human life-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 is one of the celebrated abilities of recent LLMs. 
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, in-context learning has been studied from a mathematical perspective with simplified linear self-attention without softmax unit. 
Based on a linear regression formulation $\min_x\| Ax  - b \|_2$, existing works show linear Transformers' capability of learning linear functions in context. The capability of Transformers with softmax unit approaching full Transformers, however, remains unexplored.

In this work, we study the in-context learning based on a softmax regression formulation $\min_{x} \| \langle \exp(Ax), {\bf 1}_n \rangle^{-1} \exp(Ax) - b \|_2$. We show the upper bounds of the data transformations induced by a single self-attention layer with softmax unit and by gradient-descent on a $\ell_2$ regression loss for softmax prediction function.
Our theoretical results imply that when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient-descent and Transformers show great similarity.

Main paper: 
 introduction: 
 In recent years, there has been a significant increase in research and development in the field of Artificial Intelligence, with large language models (LLMs) emerging as an effective way to tackle complex tasks. Transformers have achieved state-of-the-art results in various NLP tasks, such as machine translation (Prato et al., 2019; Gao et al., 2020 ) and text generation (Luo et al., 2022) . As a result, they have become the preferred architecture for NLP, where BERT (Devlin et al., 2018) , GPT-3 (Brown et al., 2020) , PaLM (Chowdhery et al., 2022) were proposed. They have demonstrated remarkable learning and reasoning capabilities and have proven to be more efficient than traditional models when processing natural language.Additionally, LLMs can be fine-tuned for multiple purposes without requiring a new build from scratch, making them a versatile tool for AI applications. Moreover, recent studies on the in-context learning abilities of LLMs have demonstrated that even without further fine-tuning, LLMs can generalize to new tasks with only a few demonstration examples in the prompt. To understand how LLMs become in-context learners, recent works have studied the problem and provided mathematical explanations from the Transformer architecture perspective, showing a simplified linear self-attention layer of Transformer can learn linear functions similarly as a step of gradient descent (Oswald et al., 2022; Akyürek et al., 2022; Garg et al., 2022) . While such linear approximation of full Transformers is overly simplistic, studies on more complex Transformer architecture are needed to further explain the in-context learning phenomenon.Transformers have a specific type of sequence-to-sequence neural network architecture. They utilize the attention mechanism (Vaswani et al., 2017; Radford et al., 2018; Devlin et al., 2018; Brown et al., 2020) that allows them to capture long-range dependencies and context from input data effectively. The core of the attention mechanism is the attention matrix which is comprised of rows and columns, corresponding to individual words or "tokens". The attention matrix represents the relationships within the given text. It measures the importance of each token in a sequence as it relates to the desired output. During the training process, the attention matrix is learned and optimized to improve the accuracy of the model's predictions. Through the attention mechanism, each input token is evaluated based on its relevance to the desired output by assigning a token score. This score is determined by a similarity function that compares the current output state with input states.Theoretically, the attention matrix is comprised of the query matrix Q ∈ R n×d , the key matrix K ∈ R n×d and the value matrix V ∈ R n×d . Following Zandieh et al. (2023) ; Alman & Song (2023) ; Brand et al. (2023) , the computation of the normalized attention function is defined as D -1 exp(QK ⊤ )V . Following the transformer literature, we apply exp to a matrix entry-wise way.Here D ∈ R n×n is diagonal matrix that defined as D = diag(exp(QK ⊤ )1 n ). Intuitively, D denotes the softmax normalization matrix. A more general computation formulation can be written asD -1 n×n diagonal matrix exp(XQK ⊤ X ⊤ ) n×n X n×d V d×d , D := diag(exp(XQK ⊤ X ⊤ )1 n )In the above setting, we treat Q, K, V ∈ R d×d as weights and X is the input sentence data that has length n and each word embedding size is d. In the remaining of the part, we will switch X to notation A and use A to denote sentence. Mathematically, the attention computation problem can be formulated as a regression problem in the following sense Definition 1.1. We consider the following problemmin X∈R d×d ∥D -1 exp(AXA ⊤ ) -B∥ Fwhere A ∈ R n×d can be treated as a length-n document and each word has length-d embedding size. Here D = diag(AXA ⊤ 1 n ). For any given A ∈ R n×d and B ∈ R n×n , the goal is to find some weight X to optimize the above objective function.In contrast to the formulation in Zandieh et al. (2023) ; Alman & Song (2023); Brand et al. (2023) , the parameter X in Definition 1.1 is equivalent to the QK ⊤ ∈ R d×d in the generalized version of Zandieh et al. (2023) ; Alman & Song (2023) ; Brand et al. ( 2023) (e.g. replacing Q ∈ R n×d by XQ where X ∈ R n×d and Q ∈ R d×d . Similarly for K and V . In such scenario, X can be viewed as a matrix representation of a length-n sentence.).A number of work (Akyürek et al., 2022; Garg et al., 2022; Oswald et al., 2022) study the in-context learning from mathematical perspective in a much simplified setting than Definition 1.1, which is linear regression formulation as in Definition 1.2. They show linear Transformer without softmax unit in its attention layer can mimic the ability of gradient descent in learning linear functions in context. While the softmax unit plays an important role in attention computations of full Transformers, their simplification of the softmax unit leaves a gap in explaining LLMs' in-context learning abilities. Definition 1.2. Given a matrix A ∈ R n×d and b ∈ R n , the goal is to solve minx ∥Ax -b∥ 2Several theoretical transformer work have studied either exponential regression (Gao et al., 2023; Li et al., 2023) or softmax regression problem (Deng et al., 2023b) . In this work, to take a step forward to understand the softmax unit in the attention scheme in LLMs. We consider the following softmax regression and study the in-context learning phenomena and its closeness to gradient descent from the data transformation perspective. Definition 1.3 (Softmax Regression). Given a A ∈ R n×d and a vector b ∈ R n , the goal is to solve minx∈R d ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2We remark that the Definition 1.3 of Softmax Regression is a formulation in between Definition 1.2 and Definition 1.1.We state our major result as follows:Theorem 1.4 (Bounded shift for in-context learning, informal version of the combination of Theorem 5.1 and Theorem 5.2). If the following conditions hold• Let A ∈ R n×d . • Let b ∈ R n .• ∥A∥ ≤ R.• Let ∥x∥ 2 ≤ R.• ∥A(x t+1 -x t )∥ ∞ < 0.01. • ∥(A t+1 -A t )x∥ ∞ < 0.01. • Let R ≥ 4. • Let M := n 1.5 exp(10R 2 ).We consider the softmax regression (Definition 1.3) problemmin x ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2 .• Part 1. If we move the x t to x t+1 , then we're solving a new softmax regression withmin x ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2 where ∥ b -b∥ 2 ≤ M • ∥x t+1 -x t ∥ 2• Part 2. If we move the A t to A t+1 , then we're solving a new softmax regression withmin x ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2 where ∥ b -b∥ 2 ≤ M • ∥A t+1 -A t ∥Recall that A ∈ R n×d denotes a length-n document and each word has the length-d embedding size and x denotes the simplified version of QK ⊤ . One-step gradient descent can be treated as an update to the model's weight x. Thus, part 1 of our result (Theorem 1.4) implies that the data transformation of b induced by gradient-descent on the ℓ 2 regression loss is bounded by M • ∥x t+1 -x t ∥ 2 .According to Oswald et al. (2022) , to do in-context learning, a self-attention layer update can be treated as an update to the tokenized document A. For detailed derivation, please refer to Oswald et al. (2022) . Thus, part 2 of our result (Theorem 1.4) implies that the data transformation of b induced by a single self-attention layer is bounded byM • ∥A t+1 -A t ∥.We remark that the data transformation of b induced by 1) a single self-attention layer and by 2) gradient-descent on the ℓ 2 regression loss are both bounded. The bounded transformation of b implies that when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient-descent and Transformers show great similarity.Roadmap. In Section 2, we introduce some related works. In Section 3, we give some preliminaries.In Section 4, we compute the gradient of the loss function with softmax function with respect to x.Those functions include α(x) -1 , α(x) and f (x). In Section 5, we give our formal theoretical results, validated by preliminary experiments present in Appendix D. In Section 6, we conclude our paper. Formally speaking, they consider the following problem min x ∥Ax -b∥ 2 defined in Definition 1.2. They show that one step of gradient descent carries out data transformation as follows: 
related work: 
 ∥A(x + δ x ) -b∥ 2 = ∥Ax -(b -δ b )∥ 2 = ∥Ax -b∥ 2where δ x denotes the one-step gradient descent on x and δ b denotes the corresponding data transformation on b. They also show that a self-attention layer is in principle capable of exploiting statistics in the current training data samples. Concretely, let Q, K, V ∈ R d×d denotes the weights for the query matrix, key matrix, and value matrix respectively. The linear self-attention layer updates an input sample by doing the following data transformation:b j = b j + P V K ⊤ Q jwhere b denotes the updated b and P denotes the projection matrix such that a Transformer step b j on every j is identical to the gradient-induced dynamics b j . This equivalence implies that when training linear-self-attention-only Transformers for fundamental regression tasks, the models learned by GD and Transformers show great similarity.Xie et al. ( 2021) explores the occurrence of in-context learning during pre-training when documents exhibit long-range coherence. The Language Model (LLM) develops the ability to generate coherent next tokens by deducing a latent document-level concept. During testing, in-context learning is observed when the LLM deduces a shared latent concept between examples in a prompt. They demonstrate that in-context learning happens even when there is a distribution mismatch between prompts and pretraining data, especially when the pretraining distribution is a mixture of Hidden Markov Models (Baum & Petrie, 1966) . Theoretically, they show that the error of the in-context predictor is optimal when a distinguishability condition holds. In cases where this condition does not hold, the expected error still reduces as the length of each example increases. This finding highlights the importance of both input and input-output mapping for in-context learning.The advancements of Transformers have been noteworthy, however, their learning mechanisms are not completely comprehensible yet. Although these models have performed remarkably well in structured and reasoning activities, our comprehension of their mathematical foundations lags significantly behind. Past research has indicated that the outstanding performance of Transformer-based models can be attributed to the information within their components, such as multi-head attention. Various studies (Tenney et al., 2019; Vig & Belinkov, 2019; Hewitt & Liang, 2019; Belinkov, 2022) have presented empirical proof that these components carry a substantial amount of information, which can help resolve different probing tasks.Recent research has investigated the potential of Transformers through both theoretical and experimental methods, including Turing completeness (Bhattamishra et al., 2020b) , function approximation (Yun et al., 2020; Chen et al., 2021) , formal language representation (Bhattamishra et al., 2020a; Ebrahimi et al., 2020; Yao et al., 2021) , and abstract algebraic operation learning (Zhang et al., 2022) . Some of these studies have indicated that Transformers may act as universal approximators for sequence-to-sequence operations and emulate Turing machines (Pérez et al., 2019; Bhattamishra et al., 2020b) . Liu et al. (2023) demonstrate the existence of contextual sparsity in LLM, which can be accurately predicted. They exploit the sparsity to speed up LLM inference without degrading the performance from both a theoretical perspective and an empirical perspective.  
preliminary: 
 In Section 3.1, we introduce the notations used in this paper. In Section 3.2, we give some facts about the basic algebra. In Section 3.3, we propose the lower bound on ⟨exp(Ax), 1 n ⟩.For a positive integer n, we use [n] to denote {1, 2, • • • , n}, for any positive integer n. We use E[•] to denote expectation. We use Pr[•] to denote probability. We use 1 n to denote the vector where all entries are one. We use 0 0 to denote the vector where all entries are zero. The identity matrix of size n × n is represented by I n for a positive integer n. The symbol R refers to real numbers and R ≥0 represents non-negative real numbers. For any vector x ∈ R n , exp(x) ∈ R n denotes a vector where the i-th entry is exp(x i ) and ∥x∥ 2 represents its ℓ 2 norm, that is, ∥x∥ 2 := ( n i=1 x 2 i ) 1/2 . We use ∥x∥ ∞ to denote max i∈[n] |x i |.Fact 3.1. For vectors x, y ∈ R n , we have• ∥x • y∥ 2 ≤ ∥x∥ ∞ • ∥y∥ 2 • ∥x∥ ∞ ≤ ∥x∥ 2 ≤ √ n∥x∥ ∞ • ∥ exp(x)∥ ∞ ≤ exp(∥x∥ 2 ) • For any ∥x -y∥ ∞ ≤ 0.01, we have ∥ exp(x) -exp(y)∥ 2 ≤ ∥ exp(x)∥ 2 • 2∥x -y∥ ∞ Fact 3.2.For matrices X, Y , we have• ∥X ⊤ ∥ = ∥X∥ • ∥X∥ ≥ ∥Y ∥ -∥X -Y ∥ • ∥X + Y ∥ ≤ ∥X∥ + ∥Y ∥ • ∥X • Y ∥ ≤ ∥X∥ • ∥Y ∥ • If X ⪯ α • Y , then ∥X∥ ≤ α • ∥Y ∥ 3.3 LOWER BOUND ON β Lemma 3.3. If the following conditions holds • ∥A∥ ≤ R • ∥x∥ 2 ≤ R • Let β be lower bound on ⟨exp(Ax), 1 n ⟩ Then we have β ≥ exp(-R 2 ) Proof. We have ⟨exp(Ax), 1 n ⟩ = n i=1 exp((Ax) i ) ≥ min i∈[n] exp((Ax) i ) ≥ min i∈[n] exp(-|(Ax) i |) = exp(-max i∈[n] |(Ax) i |) = exp(-∥Ax∥ ∞ ) ≥ exp(-∥Ax∥ 2 ) ≥ exp(-R 2 )the 1st step follows from simple algebra, the 2nd step comes from simple algebra, the 3rd step follows from the fact that exp(x) ≥ exp(-|x|), the 4th step follows from the fact that exp(-x) is monotonically decreasing, the 5th step comes from definition of ℓ ∞ norm, the 6th step follows from Fact 3.1, the 7th step follows from the assumption on A and x. 
softmax function with respect to x: 
 In Section 4.1, we give the definitions used in the computation. In Section 4.2, we compute the gradient of the loss function with softmax function with respect to x. Those functions includes α(x) -1 , α(x) and f (x).We define function softmax f as follows Definition 4.1 (Function f , Definition 5.1 in Deng et al. (2023b) ). Given a matrix A ∈ R n×d . Let 1 n denote a length-n vector that all entries are ones. We define prediction function f : R d → R n as follows f (x) := ⟨exp(Ax), 1 n ⟩ -1 • exp(Ax).L exp (x) := 0.5 • ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2 2 .For convenient, we define two helpful notations α and c Definition 4.3 (Normalized coefficients, Definition 5.4 in Deng et al. (2023b) ). We define α : R d → R as follows α(x) := ⟨exp(Ax), 1 n ⟩.Then, we can rewrite f (x) (see Definition 4.1) and L exp (x) (see Definition 4.2) as follows• f (x) = α(x) -1 • exp(Ax). • L exp (x) = 0.5 • ∥α(x) -1 • exp(Ax) -b∥ 2 2 .• L exp (x) = 0.5 • ∥f (x) -b∥ 2 2 . Definition 4.4 (Definition 5.5 in Deng et al. (2023b) ). We define function c : R d ∈ R n as followsc(x) := f (x) -b.Then we can rewrite L exp (x) (see Definition 4.2) as follows• L exp (x) = 0.5 • ∥c(x)∥ 2 2 .We state a lemma from previous work, Lemma 4.5 (Gradient, Lemma 5.6 in Deng et al. (2023b) ). If the following conditions hold• Given matrix A ∈ R n×d and a vector b ∈ R n .• Let α(x) be defined in Definition 4.3.• Let f (x) be defined in Definition 4.1.• Let c(x) be defined in Definition 4.4.• Let L exp (x) be defined in Definition 4.2.For each i ∈ [d], we have• Part 1. d exp(Ax) dx i = exp(Ax) • A * ,i • Part 2. d⟨exp(Ax), 1 n ⟩ dx i = ⟨exp(Ax), A * ,i ⟩ • Part 3. dα(x) -1 dx i = -α(x) -1 • ⟨f (x), A * ,i ⟩ • Part 4. df (x) dx i = dc(x) dx i = -⟨f (x), A * ,i ⟩ • f (x) + f (x) • A * ,i • Part 5. dL exp (x) dx i = A ⊤ * ,i 1×n • f (x) n×1 ⟨c(x), f (x)⟩ scalar + diag(f (x)) n×n c(x) n×1 
main results: 
 In Section 5.1, we show the lipschitz bound of function f . In Section 5.2, we show our upper bound result of δ b with respect to x. In Section 5.3, we show our upper bound result of δ b with respect to A. We also conduct preliminary experiments in Appendix D that provide empirical validation for our theoretical results.To bound the shift of b, we first show the Lipschitz property for the basic functions:• ∥ exp(Ax) -exp(Ay)∥ 2 ≤ 2 √ nR exp(R 2 ) • ∥x -y∥ 2 • |α(x) -α(y)| ≤ ∥ exp(Ax) -exp(Ay)∥ 2 • √ n • |α(x) -1 -α(y) -1 | ≤ β -2 • |α(x) -α(y)|Then, following the above Lipschitz property, we have ∥f(x t+1 ) -f (x t )∥ 2 ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2Similarly, we can show the Lipschtiz property of function f with respect to A as the following∥f (A t+1 ) -f (A t )∥ 2 ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥ 2Due to space limitation, we defer formal lemma and proof to Appendix A.2 and C.2.Theorem 5.1 (Bounded shift for shifting the weight parameter, formal of Theorem 1.4). If the following conditions hold• Let A ∈ R n×d • ∥A∥ ≤ R • ∥A(x t+1 -x t )∥ ∞ < 0.01 • Let R ≥ 4 • Let M := n 1.5 exp(10R 2 ).We consider the softmax regression problem minx ∥⟨exp(Ax),1 n ⟩ -1 exp(Ax) -b∥ 2If we move the x t to x t+1 , then we're solving a new softmax regression problem withmin x ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2 where ∥ b -b∥ 2 ≤ M • ∥x t+1 -x t ∥ 2Proof. We have∥ b -b∥ 2 ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2 ≤ 4n 1.5 R exp(2R 2 ) exp(2R 2 ) • ∥x t+1 -x t ∥ 2 ≤ n 1.5 (4R) exp(4R 2 ) • ∥x t+1 -x t ∥ 2 ≤ n 1.5 exp(6R 2 ) exp(4R 2 ) • ∥x t+1 -x t ∥ 2 ≤ n 1.5 exp(10R 2 ) • ∥x t+1 -x t ∥ 2 ≤ M • ∥x t+1 -x t ∥ 2where the 1st step follows from Lemma A.5, the 2nd step comes from Lemma 3.3, the 3rd step comes from simple algebra, the 4th step follows from simple algebra, the 5th step follows from simple algebra and the 6th step follows from the definition of M .Theorem 5.2 (Bounded shift for in-context learning, formal of Theorem 1.4). If the following conditions hold• Let A ∈ R n×d • ∥A∥ ≤ R • ∥(A t+1 -A t )x∥ ∞ < 0.01 • Let R ≥ 4 • Let M := n 1.5 exp(10R 2 ).We consider the softmax regression problemmin x ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2If we move the A t to A t+1 then we're solving a new softmax regression problem withmin x ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2 where ∥ b -b∥ 2 ≤ M • ∥A t+1 -A t ∥. Proof. We have ∥ b -b∥ 2 ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥ ≤ 4n 1.5 R exp(2R 2 ) exp(2R 2 ) • ∥A t+1 -A t ∥ ≤ n 1.5 (4R) exp(4R 2 ) • ∥A t+1 -A t ∥ ≤ n 1.5 exp(6R 2 ) exp(4R 2 ) • ∥A t+1 -A t ∥ ≤ n 1.5 exp(10R 2 ) • ∥A t+1 -A t ∥ ≤ M • ∥A t+1 -A t ∥where the 1st step follows from Lemma C.5, the 2nd step follows from Lemma 3.3, the 3rd step follows from simple algebra, the 4th step comes from simple algebra, the 5th step comes from simple algebra and the 6th step follows from the definition of M . 
conclusion: 
 The attention mechanism that incorporates the softmax unit is a crucial aspect of Large Language Models (LLMs) and significantly contributes to their extraordinary performance in various Natural Language Processing (NLP) tasks. The ability to learn in-context is highly valued in recent LLMs, and comprehending this concept is vital when querying LLMs. In this study, taking a step further from prior works' studies on linear Transformer's ability of learning linear functions, we examined the in-context learning process from a softmax regression perspective of Transformer's attention mechanism. We established the bound on the data transformations brought about by a single selfattention layer with softmax unit and gradient descent on an L2 regression loss. Our findings suggest that the update acquired through gradient descent and in-context learning are highly similar when training self-attention-only Transformers for softmax regression tasks, which is also validated through our preliminary experimental results. These results offer insights into the theoretical underpinnings of in-context learning in Transformers and can aid in improving the understanding and performance of LLMs in various NLP tasks. Our findings are restricted to small Transformer and simple regression problems. One interesting direction for further investigation is to acquire a comprehensive perception of in-context learning in larger models. To the best of our knowledge, we believe this work does not have any negative societal impact.• We can rewrite δ b,1 as followsδ b,1 = (α(x t+1 ) -1 -α(x t ) -1 ) • exp(Ax t+1 ),• We can rewrite δ b,2 as followsδ b,2 = α(x t ) -1 • (exp(Ax t+1 ) -exp(Ax t )).Proof. We haveδ b = δ b,1 + δ b,2 = α(x t+1 ) -1 exp(Ax t+1 ) -α(x t ) -1 exp(Ax t+1 )+ α(x t ) -1 exp(Ax t+1 ) -α(x t ) -1 exp(Ax t ) = α(x t+1 ) -1 exp(Ax t+1 ) -α(x t ) -1 exp(Ax t ) = ⟨exp(Ax t+1 ), 1 n ⟩ -1 exp(Ax t+1 ) -⟨exp(Ax t ), 1 n ⟩ -1 exp(Ax t ),where the 1st step follows from the definitions of δ b , the 2nd step follows from the definitions of δ b,1 and δ b,2 , the 3rd step follows from simple algebra, the 4th step comes from the definition of α. 
 
 Appendix: We can show that Lemma A.5. If the following conditions hold• Let β ∈ (0, 1).• Let δ b,1 ∈ R n be defined as Definition A.3.• Let δ b,2 ∈ R n be defined as Definition A.3.• Let δ b = δ b,1 + δ b,2 . • Let R ≥ 4.We have• Part 1. ∥δ b,1 ∥ 2 ≤ 2β -2 n 1.5 exp(2R 2 ) • ∥x t+1 -x t ∥ 2 • Part 2. ∥δ b,2 ∥ 2 ≤ 2β -1 √ nR exp(R 2 ) • ∥x t+1 -x t ∥ 2 • Part 3. ∥ f (x t+1 ) -f (x t ) δ b ∥ 2 ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2 Proof. Proof of Part 1. We have ∥δ b,1 ∥ 2 ≤ |α(x t+1 ) -1 -α(x t ) -1 | • ∥ exp(Ax t+1 )∥ 2 ≤ |α(x t+1 ) -1 -α(x t ) -1 | • √ n • exp(R 2 ) ≤ β -2 • |α(x t+1 ) -α(x t )| • √ n • exp(R 2 ) ≤ β -2 • √ n • ∥ exp(Ax t+1 ) -exp(Ax t )∥ 2 • √ n • exp(R 2 ) ≤ β -2 • √ n • 2 √ nR exp(R 2 )∥x t+1 -x t ∥ 2 • √ n • exp(R 2 ) = 2β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2where the first step follows from definition, the second step follows from assumption on A and x, the third step follows Lemma A.8, the forth step follows from Lemma A.7, the fifth step follows from Lemma A.6.Proof of Part 2.We have∥δ b,2 ∥ 2 ≤ |α(x t+1 ) -1 | • ∥ exp(Ax t+1 ) -exp(Ax t )∥ 2 ≤ β -1 • ∥ exp(Ax t+1 ) -exp(Ax t )∥ 2 ≤ β -1 • 2 √ nR exp(2R 2 ) • ∥x t+1 -x t ∥ 2where the first step follows from definition, the 2nd step comes from Lemma A.6.Proof of Part 3.We have∥δ b ∥ 2 = ∥δ b,1 + δ b,2 ∥ 2 ≤ ∥δ b,1 ∥ 2 + ∥δ b,2 ∥ 2 ≤ 2β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2 + 2β -1 n 0.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2 ≤ 2β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2 + 2β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2 ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥x t+1 -x t ∥ 2where the 1st step follows from the definition of δ b , the 2nd step follows from triangle inequality, the 3rd step follows from the results in Part 1 and Part 2, the 4th step follows from the fact that n ≥ 1 and β -1 ≥ 1, the 5th step follows from simple algebra.A.3 LIPSCHITZ FOR FUNCTION exp(Ax) WITH RESPECT TO x Lemma A.6. If the following conditions holds• Let A ∈ R n×d • Let ∥A(y -x)∥ ∞ < 0.01 • Let ∥A∥ ≤ R • Let x, y satisfy that ∥x∥ 2 ≤ R and ∥y∥ 2 ≤ R Then we have ∥ exp(Ax) -exp(Ay)∥ 2 ≤ 2 √ nR exp(R 2 ) • ∥x -y∥ 2 .Proof. We have∥ exp(Ax) -exp(Ay)∥ 2 ≤ ∥ exp(Ax)∥ 2 • 2∥A(x -y)∥ ∞ ≤ √ n • exp(∥Ax∥ 2 ) • 2∥A(x -y)∥ ∞ ≤ √ n exp(R 2 ) • 2∥A(x -y)∥ 2 ≤ √ n exp(R 2 ) • 2∥A∥ • ∥x -y∥ 2 ≤ 2 √ nR exp(R 2 ) • ∥x -y∥ 2where the 1st step follows from ∥A(y -x)∥ ∞ < 0.01 and Fact 3.1, the 2nd step comes from Fact 3.1, the 3rd step follows from Fact 3.2, the 4th step follows from Fact 3.2, the last step follows from ∥A∥ ≤ R.We state a tool from previous work Deng et al. (2023b) . Lemma A.7 (Lemma 7.2 in Deng et al. (2023b) ). If the following conditions hold• Let α(x) be defined as Definition 4.3Then we have|α(x) -α(y)| ≤ ∥ exp(Ax) -exp(Ay)∥ 2 • √ n. A.5 LIPSCHITZ FOR FUNCTION α(x) -1 WITH RESPECT TO xWe state a tool from previous work (Deng et al., 2023b) . Lemma A.8 (Lemma 7.2 in Deng et al. (2023b) ). If the following conditions hold• Let ⟨exp(Ax), 1 n ⟩ ≥ β • Let ⟨exp(Ay), 1 n ⟩ ≥ βThen, we have|α(x) -1 -α(y) -1 | ≤ β -2 • |α(x) -α(y)|.In this section, we consider the function with respect to A. We define function softmax f as follows Definition B.1 (Function f , Reparameterized x by A in Definition 4.1). Given a matrix A ∈ R n×d . Let 1 n denote a length-n vector that all entries are ones. We define prediction function f : R n×d → R n as followsf (A) := ⟨exp(Ax), 1 n ⟩ -1 • exp(Ax).Similarly, we reparameterized x by A for our loss function L. We define loss function L as follows Definition B.2 (Loss function L exp , Reparameterized x by A in Definition 4.2). Given a matrix A ∈ R n×d and a vector b ∈ R n×d . We define loss function L exp : R n×d → R as followsL exp (A) := 0.5 • ∥⟨exp(Ax), 1 n ⟩ -1 exp(Ax) -b∥ 2 2 .For convenience, we define two helpful notations α and c with respect to A as follows: Definition B.3 (Normalized coefficients, Reparameterized x by A in Definition 4.3). We define α : R n×d → R as follows α(A) := ⟨exp(Ax), 1 n ⟩.Then, we can rewrite f (A) (see Definition B.1) and L exp (A) (see Definition B.2) as follows Then we can rewrite L exp (A) (see Definition B.2) as follows• f (A) = α(A) -1 • exp(Ax). • L exp (A) = 0.5 • ∥α(A) -1 • exp(Ax) -b∥ 2 2 . • L exp (A) = 0.5 • ∥f (A) -b∥ 2 2 . Definition B.• L exp (A) = 0.5 • ∥c(A)∥ 2 2 .In Section C.1, we give the preliminary to compute the Lipschitz. In Section C.2, we show the upper bound of δ b with respect to A. In Section C.3, we compute the Lipschitiz of function exp(Ax) with respect to A. In Section C.4, we compute the Lipschitiz of the function α with respect to A. In Section C.5, we compute the Lipschitiz of function α -1 with respect to A.• Let R ≥ 4.We have• Part 1. ∥δ b,1 ∥ 2 ≤ 2β -2 n 1.5 exp(2R 2 ) • ∥A t+1 -A t ∥ 2 • Part 2. ∥δ b,2 ∥ 2 ≤ 2β -1 √ nR exp(R 2 ) • ∥A t+1 -A t ∥ 2 • Part 3. ∥ f (A t+1 ) -f (A t ) δ b ∥ 2 ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥ 2Proof. Proof of Part 1. We have∥δ b,1 ∥ 2 ≤ |α(A t+1 ) -1 -α(A t ) -1 | • ∥ exp(A t+1 x)∥ 2 ≤ |α(A t+1 ) -1 -α(A t ) -1 | • √ n • exp(R 2 ) ≤ β -2 • |α(A t+1 ) -α(A t )| • √ n • exp(R 2 ) ≤ β -2 • √ n • ∥ exp(A t+1 x) -exp(A t x)∥ 2 • √ n • exp(R 2 ) ≤ β -2 • √ n • 2 √ nR exp(R 2 )∥A t+1 -A t ∥ • √ n • exp(R 2 ) = 2β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥where the first step follows from definition, the second step follows from assumption on A and x, the third step follows Lemma C.8, the forth step follows from Lemma C.7, the fifth step follows from Lemma C.6.Proof of Part 2.We have∥δ b,2 ∥ 2 ≤ |α(A t+1 ) -1 | • ∥ exp(A t+1 x) -exp(A t x)∥ 2 ≤ β -1 • ∥ exp(A t+1 x) -exp(A t x)∥ 2 ≤ β -1 • 2 √ nR exp(2R 2 ) • ∥A t+1 -A t ∥ Proof of Part 3.We have∥δ b ∥ 2 = ∥δ b,1 + δ b,2 ∥ 2 ≤ ∥δ b,1 ∥ 2 + ∥δ b,2 ∥ 2 ≤ 2β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥ + 2β -1 n 0.5 R exp(2R 2 ) • ∥A t+1 -A t ∥ ≤ 2β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥ + 2β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥ ≤ 4β -2 n 1.5 R exp(2R 2 ) • ∥A t+1 -A t ∥where the 1st step follows from the definition of δ b , the 2nd step comes from triangle inequality, the 3rd step comes from the results in Part 1 and Part 2, the 4th step follows from the fact that n ≥ 1 and β -1 ≥ 1, the 5th step follows from simple algebra.C.3 LIPSCHITZ FOR FUNCTION exp(Ax) WITH RESPECT TO A Lemma C.6 (Reparameterized x by A in Lemma A.6). If the following conditions holds• Let A, B ∈ R n×d • Let ∥(A -B)x∥ ∞ < 0.01 • Let ∥A∥ ≤ R • Let x satisfy that ∥x∥ 2 ≤ R Then we have ∥ exp(Ax) -exp(Bx)∥ 2 ≤ 2 √ nR exp(R 2 ) • ∥A -B∥.Proof. We have∥ exp(Ax) -exp(Bx)∥ 2 ≤ ∥ exp(Ax)∥ 2 • 2∥(A -B)x∥ ∞ ≤ √ n • exp(∥Ax∥ 2 ) • 2∥(A -B)x∥ ∞ ≤ √ n exp(R 2 ) • 2∥(A -B)x∥ 2 ≤ √ n exp(R 2 ) • 2∥A -B∥ • ∥x∥ 2 ≤ 2 √ nR exp(R 2 ) • ∥A -B∥where the 1st step follows from ∥A(y -x)∥ ∞ < 0.01 and Fact 3.1, the 2nd step follows from Fact 3.1, the 3rd step follows from Fact 3.2, the 4th step comes from Fact 3.2, the last step follows from ∥A∥ ≤ R. Then we have|α(A) -α(B)| ≤ ∥ exp(Ax) -exp(Bx)∥ 2 • √ n.Proof. We have|α(A) -α(B)| = |⟨exp(Ax) -exp(Bx), 1 n ⟩| ≤ ∥ exp(Ax) -exp(Bx)∥ 2 • √ nwhere the 1st step comes from the definition of α(x), the 2nd step follows from Cauchy-Schwarz inequality (Fact 3.1).C.5 LIPSCHITZ FOR FUNCTION α(A) -1 WITH RESPECT TO A Lemma C.8 (Reparameterized x by A in Lemma A.8). If the following conditions hold• Let ⟨exp(Ax), 1 n ⟩ ≥ β • Let ⟨exp(Bx), 1 n ⟩ ≥ β Then, we have |α(A) -1 -α(B) -1 | ≤ β -2 • |α(A) -α(B)|.Proof. We can show that|α(A) -1 -α(B) -1 | = α(A) -1 α(B) -1 • |α(A) -α(B)| ≤ β -2 • |α(A) -α(B)|where the 1st step follows from simple algebra, the 2nd step follows from α(A) ≥ β, α(B) ≥ β.In this section, we show preliminary experimental findings supporting our theoretical results that when training self-attention-only Transformers for softmax regression tasks, the models learned by gradient-descent and Transformers show great similarity. {5, 10, 20, 35, 50}. Following Oswald et al. (2022) , we compare the two models in our experiment: a trained single self-attention (SA) layer with a softmax unit approximating the full Transformers, and a softmax regression model trained with one-step gradient descent. The training objective for both models is defined as in Definition 1.3. For the single self-attention layer with a softmax unit, we choose the learning rate η SA = 0.005. For the softmax regression model, we determine the optimal learning rate η GD by minimizing the ℓ 2 regression loss over a training set of 10 3 tasks through line search.To compare the trained single self-attention layer with a softmax unit and the softmax regression model trained with one-step gradient descent, we sample 10 3 tasks and record the losses of two models. In addition, we follow Oswald et al. ( 2022) to record• Pred Diff: the predictions difference measured with the ℓ 2 norm:∥ y SA (A) -y GD (x)∥ 2where y SA (A) is corresponding to the b in Theorem 5.1, and y GD (x) is corresponding to the b in Theorem 5.2.• Model Cos: the cosine similarity between the sensitivities of two models:CosSim( ∂ y GD (x) ∂x , ∂ y SA (A) ∂A )• Model Diff: the model sensitivity difference measured with the ℓ 2 norm:∥ ∂ y GD (x) ∂x - ∂ y SA (A) ∂A ∥ 2All experiments run on a single NVIDIA RTX2080Ti GPU with 10 independent repetitions.Results on tasks of different document lengths. The results of the comparisons between a trained single self-attention layer and one-step gradient descent on synthetic softmax regression tasks of document length n ∈ {25, 50, 100, 200, 400} and word embedding size d = 20 are shown in Figure 1 -5. We measure two models' losses and similarities over the training steps of the SA layer for each set of tasks. From the results, we observe identical performances of the two models measured in losses. We also observe considerable alignment of the two models across tasks of different document lengths, indicated by decreasing prediction and model difference and increasing cosine similarity between models. Besides, comparing all five figures, we observe that with larger document length n, which is common in practical language modeling tasks, the two models show greater similarity, validating the Transformer's in-context learning ability.Results on tasks of different word embedding sizes. The results of the comparisons between a trained single self-attention layer and one-step gradient descent on synthetic softmax regression tasks of document length n = 200 and word embedding size d ∈ {5, 10, 20, 35, 50} are shown in Figure 4 and 6-9. Similarly, we measure two models' losses and similarities over training steps of the SA layer for each set of tasks. We again observe similar performances and close alignment of the two models.In conclusion, our experimental results empirically validate our theoretical results in Section 5, showing that when training self-attention-only Transformers for softmax regression tasks, the models learned by gradient-descent and Transformers show great similarity. Due to the non-linearity of softmax regression, it is not expected for models to match exactly as implied in our theoretical results in Section 5, which is also observed in our experimental findings.  


Decision: This paper provides a theoretical analysis of in-context learning in Transformers, focusing on the role of the softmax unit within self-attention layers. By examining the upper bounds of data transformations and comparing these with gradient descent on a regression loss, the authors offer insights into the similarities between models learned by Transformers and traditional methods. While the study is primarily theoretical, it fills a critical gap in understanding how softmax contributes to in-context learning. Based on the reviews, I recommend to accept the paper.

Reviews: review: paper_summary: paper_summary: This research delves into enhancing the comprehension of in-context learning through theoretical analysis. It builds on prior studies showcasing how a single self-attention layer can learn gradient steps in linear regression contexts. The authors extend this concept to softmax regression and give the upper bounds of the data transformations driven by gradient descent for a single self-attention layer. Empirical results in the appendix validate these theoretical advancements.
strength_weakness: strength_weakness: - The presentation is hard to follow. It's better to organize the formulation in a question-driven format, and it's unclear why bounding the single step of data transformation relates to building a connection between in-context learning and softmax weight shift. Consider optimizing the presentation of certain mathematical proofs by relocating them to the appendix for conciseness, and leaving more place in the main paper for intuition illustration and experiment.
- As the theoretical analysis and experiment only consider a single attention layer, it's unclear whether other components in Transformers, such as MLP and Layer Norm, will affect the conclusion in the paper.
questions: questions: See above.
limitations: limitations: The authors adequately addressed the limitations
paper_summary: paper_summary: This paper explores the relationship between in-context learning and weight shifting in the context of softmax regression for large language models. The authors delve into the mathematical aspects and study the in-context learning based on a softmax regression formulation. They present theoretical results that imply when training self-attention-only Transformers for fundamental regression tasks, the models learned by gradient descent and Transformers show great similarity. The paper shows the upper bounds of the data transformations induced by a single self-attention layer with a softmax unit and by gradient descent on a ℓ2 regression loss for softmax prediction function. The authors provide a formal theoretical result and numerical experiments to validate the theoretical results.
strength_weakness: strength_weakness: 1. While the paper compares the performance of a single self-attention layer with a softmax unit and a softmax regression model trained with one-step gradient descent, it would be informative to include comparisons with state-of-the-art models and methods.

2. The paper could benefit from a deeper exploration of the explainability and interpretability of the models in the context of softmax regression.
questions: questions: Please refer to Weaknesses.
limitations: limitations: Please refer to Weaknesses.
paper_summary: paper_summary: This paper aims to investigate why transformers possess the capability of in-context learning from a theoretical perspective. Previous works have shown simplified self-attention layer’s capability of learning linear functions in context. This work conducts further research based on softmax regression, as softmax-based algorithms are more complex and closer to algorithms used in actual LLMs. Through mathematical analysis and experimental validation, the authors conclude that the update acquired through gradient descent and in-context learning are similar when training simplified transformers for softmax regression tasks.

I am not fully follow this paper and I wish to learn more intuitive understanding of this paper during rebuttal (from the authors and other reviewers). I may adjust my rating score after rebuttal period.
strength_weakness: strength_weakness: 1.This paper appears to build upon previous research that explored in-context learning capabilities of transformer using linear regression, but instead opts for a softmax regression approach, lacking significant innovation and contribution.

2.Although this paper represents a further advancement over previous studies based on linear regression, it still employs a highly simplified Transformer model, which is insufficient to fully explain the principles of LLMs’ in-context learning ability.

3.I suggest that the structure of the paper could be organized more clearly, allowing readers to grasp the overall framework first. The details of the two models compared in the experiments should be explained more thoroughly. (1) The section 2 and 3 are not so well-organized, where some formula lacks necessary comments and explanations. (2) The introduction section is not so intuitive to follow.  (3) There is almost no textual descriptions about the proposed method section (section 3). I believe most NeurIPS reviewers are hard to follow the section 3 without enough descriptions.
questions: questions: None
limitations: limitations: The limitations are fine with me.

Comments: title: Reply to Reviewer LX6R
Reply to Reviewer K6fK
Thanks for responding my comments
comment: Thank you for your response and your recognition of our contribution. As you suggested, we have improved the paper presentation as follows, which will be reflected in the revised version of our paper:
1. We have moved some facts in Section 2 that are not used in the main body to the Appendix to keep the main paper concise and easier to follow. 
2. We have included more intuitive explanations in Section 1 and 2 on the data transformation formulation of in-context learning.
3. For experiments part in Section 5, we have included additional details of the experimental setup from the Appendix and additional discussions of our experimental findings validating our theoretical results.
Thank you for your response and your recognition of our work. As you suggested, we have improved the paper presentation as follows, which will be reflected in the revised version of our paper:
1. For the introduction in Section 1, we have added more intuitive explanations for our studied in-context learning problem and its mathematical formulation.
2. We have added additional comments in Sections 2 and 3 between the definitions to help readers’ understanding, and moved some facts that are not used in the main body to the Appendix to ensure clarity.
3. For experiments part in Section 5, we have included additional details of the experimental setup from the Appendix and additional discussions of our experimental findings validating our theoretical results.
Thank you for your rebuttal. Given the complexity of understanding in-context learning in modern LLMs, which indeed poses a significant challenge and entails long-term commitment, this paper shows some advancement compared to the previous efforts (as listed in the rebuttal), notably from linear to softmax. I will adjust the score to favor acceptance. However, I still suggest the authors to improve the paper presentation.
Thank you for your rebuttal. I still suggest the authors should make the paper easy to follow. Considering other reviewer's comments, I tend to keep my score currently and am waiting for other reviewers' further ideas.