In-Context Learning with Representations: Contextual Generalization of Trained Transformers

Authors: tong  yang
yu  cmu
  huang
upenn yingbin liang
osu  yuejie
chi § cmu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 02:08:38

Abstract: In-context learning (ICL) refers to a remarkable capability of pretrained large language models, which can learn a new task given a few examples during inference. However, theoretical understanding of ICL is largely under-explored, particularly whether transformers can be trained to generalize to unseen examples in a prompt, which will require the model to acquire contextual knowledge of the prompt for generalization. This paper investigates the training dynamics of  transformers by gradient descent through the lens of non-linear regression tasks. The contextual generalization here can be attained via  learning  the template function for each task in-context, where all template functions lie in a linear space with $m$ basis functions. We analyze the training dynamics of one-layer multi-head transformers to {in-contextly} predict unlabeled inputs given partially labeled prompts, where the labels contain Gaussian noise and the number of examples in each prompt  are not sufficient to determine the template. Under mild assumptions, we show that the training loss for a one-layer multi-head transformer converges linearly to a global minimum. Moreover, the transformer effectively learns to perform ridge regression over the basis functions. To our knowledge, this study is the first provable demonstration that transformers can learn contextual (i.e., template) information to generalize to both unseen examples and tasks when prompts contain only a small number of query-answer pairs.

Main paper: 
 introduction: 
 Transformers [Vaswani et al., 2017] have achieved tremendous successes in machine learning, particularly in natural language processing, by introducing self-attention mechanisms that enable models to capture long-range dependencies and contextualized representations. In particular, these self-attention mechanisms endow transformers with remarkable in-context learning (ICL) capabilities, allowing them to adapt to new tasks or domains by simply being prompted with a few examples that demonstrate the desired behavior, without any explicit fine-tuning or updating of the model's parameters [Brown et al., 2020] .A series of papers have empirically studied the underlying mechanisms behind in-context learning in transformer models [Garg et al., 2022 , Von Oswald et al., 2023 , Wei et al., 2023 , Olsson et al., 2022 , Xie et al., 2021 , Chen and Zou, 2024 , Agarwal et al., 2024] , which have shown that transformers can predict unseen examples after being prompted on a few examples. The pioneering work of Since our paper studies ICL of non-linear function regression, the function mapping (which we also term as "template") naturally serves as the "contextual information" that can be learned for generalization to unseen examples. When each prompt contains only a small number of (noisy) examples, the template that generates the labels may be underdetermined, i.e., multiple templates could generate the same labels in the prompt. Such an issue of underdetermination further raises a series of intriguing questions, such as:When the template that generates a prompt is underdetermined, what is the transformer's preference for choosing the template and how good is such a choice?In this paper, we answer the above questions by analyzing the training dynamics of a one-layer transformer with multi-head softmax attention through the lens of non-linear regression tasks. In our setting, the template function for each task lies in the linear space formed by m nearly-arbitrary basis functions that capture representation (i.e., features) of data. Our goal is to provide insights on how transformers trained by gradient descent (GD) acquire template information from more general function classes to generalize to unseen examples and tasks when each prompt contains only a small number of query-answer pairs. We summarize our contributions are as follows.• We first establish the convergence guarantee of a one-layer transformer with multi-head softmax attention trained with gradient descent on general non-linear regression in-context learning tasks. We assume each prompt contains only a few (i.e., partial) examples with their Gaussian noisy labels, which are not sufficient to determine the template. Under mild assumptions, we establish that the training loss of the transformer converges at a linear rate. Moreover, by analyzing the limit point of the transformer parameters, we are able to uncover what information about the basic tasks the transformer extracts and memorizes during training in order to perform in-context prediction.• We then analyze the transformer's behavior at inference time after training, and show that the transformer chooses its generating template by performing ridge regression over the basis functions. We also provide the iteration complexity for pretraining the transformer to reach ε-precision with respect to its choice of the template given an arbitrary prompt at inference time. We further compare the choice of the transformer and the best possible choice over the template class and characterize how the sequence length of each prompt influences the inference time performance of the model.• Under more realistic assumptions, our analysis framework allows us to overcome a handful of assumptions made in previous works such as large prompt length [Huang et al., 2023 , Chen et al., 2024 , Li et al., 2024 , Nichani et al., 2024] , orthogonality of data [Huang et al., 2023 , Chen et al., 2024 , Li et al., 2024 , Nichani et al., 2024] , restrictive initialization conditions [Chen et al., 2024] , special structure of the transformer [Nichani et al., 2024] , and mean-field models [Kim and Suzuki, 2024] . Further, the function classes we consider are a generalization of those considered in most theoretical works [Huang et al., 2023 , Chen et al., 2024 , Li et al., 2024 , Wu et al., 2023 , Zhang et al., 2023a] . We also highlight the importance of multi-head attention mechanism in this process.To our best knowledge, this is the first work that analyzes how transformers learn contextual (i.e., template) information to generalize to unseen examples and tasks when prompts contain only a small number of query-answer pairs. Table 1 ✗ ✗ ✓ ✗ ✗ ✓ Ours ✓ ✓ ✓ ✓ ✓ ✓Table 1 : Comparisons with existing theoretical works that study the learning dynamics of transformers in ICL. Here, the last column refers to the fact that the response in the regression task is generated by a linearly weighted unknown representation (feature) model.In-context learning. Recent research has investigated the theoretical underpinnings of transformers' ICL capabilities from diverse angles. For example, several works focus on explaining the in-context learning of transformers from a Bayesian perspective [Xie et al., 2021 , Ahuja et al., 2023 , Han et al., 2023 , Jiang, 2023 , Wang et al., 2023 , Wies et al., 2024 , Zhang et al., 2023b , Jeon et al., 2024 , Hahn and Goyal, 2023] . Li et al. [2023] analyzed the generalization and stability of transformers' in-context learning. Focusing on the representation theory, Akyürek et al. [2022] , Bai et al. [2024] studied the expressive power of transformers on the linear regression task. Akyürek et al. [2022] showed by construction that transformers can represent GD of ridge regression or the closed-form ridge regression solution. Bai et al. [2024] extended Akyürek et al. [2022] and showed that transformers can implement a broad class of standard machine learning algorithms in-context. Dai et al. [2022] , Von Oswald et al. [2023] showed transformers could in-context learn to perform GD.More pertinent to our work, Guo et al. [2023] considered an ICL setting very similar to ours, where the label depends on the input through a basis of possibly complex but fixed template functions, composed with a linear function that differs in each prompt. By construction, the optimal ICL algorithm first transforms the inputs by the representation function, and then performs linear ICL on top of the transformed dataset. Guo et al. [2023] showed the existence of transformers that approximately implement such algorithms, whereas our work is from a different perspective, showing that (pre)training the transformer loss by GD will naturally yield a solution with the aforementioned desirable property characterized in Guo et al. [2023] .Training dynamics of transformers performing ICL. A line of work initiated by Garg et al. [2022] aims to understand the ICL ability of transformers from an optimization perspective. [Zhang et al., 2023a, Kim and Suzuki, 2024] analyzed the training dynamics of transformers with linear attention. Huang et al. [2023] , Chen et al. [2024] , Li et al. [2024] studied the optimization dynamics of one-layer softmax attention transformers performing simple in-context learning tasks, such as linear regression [Huang et al., 2023 , Chen et al., 2024] and binary classification [Li et al., 2024] .Among them, Huang et al. [2023] was the first to study the training dynamics of softmax attention, where they gave the convergence results of a one-layer transformer with single-head attention on linear regression tasks, assuming context features come from an orthogonal dictionary and each token in the prompts is drawn from a multinomial distribution. In order to leverage the concentration property inherent to multinomial distributions, they require the sequence length to be much larger than the size of dictionary. Their analysis indicates that the prompt tokens that are the same as the query will have dominating attention weights, which allows the transformer to copy-paste the correct answer from those prompt tokens. Li et al. [2024] studied the training of a one-layer single-head transformer in ICL on binary classification tasks. Same as Huang et al. [2023] , they required the data to be pairwise orthogonal, and shared the same copy-paste mechanism as in Huang et al. [2023] . To be precise, a fraction of their context inputs needs to contain the same pattern as the query to guarantee that the total attention weights on contexts matching the query pattern outweigh those on other contexts. Chen et al. [2024] studied the dynamics of gradient flow for training a one-layer multi-head softmax attention model for ICL of multi-task linear regression, where the coefficient matrix has certain spectral properties. They required the sequence length to be sufficiently large [Chen et al., 2024, Assumption 2 .1], together with restrictive initialization conditions [Chen et al., 2024, Definition 3.1] . While using the copy-paste analysis framework as in Huang et al. [2023] , Li et al. [2024] , the attention probability vector in their work is delocalized, so that the attention is spread out to capture the information from similar tokens in regression tasks. Kim and Suzuki [2024] studied the dynamics of Wasserstein gradient flow for training a one-layer transformer with an infinite-dimensional fullyconnected layer followed by a linear attention layer for ICL of linear regression, assuming infinite prompt length. Nichani et al. [2024] analyzed the optimization dynamics of a simplified two-layer transformer with gradient descent on in-context learning a latent causal graph.Notation. Boldface small and capital letters denote vectors and matrices, respectively. Sets are denoted with curly capital letters, e.g., W. We let (R d , ∥•∥) denote the d-dimensional real coordinate space equipped with norm ∥•∥. I d is the identity matrix of dimension d. The ℓ p -norm of v is denoted by ∥v∥ p , where 1 ≤ p ≤ ∞, and the spectral norm and the Frobenius norm of a matrix M are denoted by ∥M ∥ 2 and ∥M ∥ F , respectively. M † stands for the Moore-Penrose pseudoinverse of matrix M , and M :,i stands for its i-th column vector. We let [N ] denote {1, . . . , N }, and denote 1 N to represent the all-one vector of length N , and by 0 a vector or a matrix consisting of all 0's. We allow the application of functions such as exp(•) to vectors or matrices, with the understanding that they are applied in an element-wise manner. We use e i to denote the one-hot vector whose i-th entry is 1 and the other entries are all 0. 
problem setup: 
 In-context learning with representation. We consider ICL of regression with unknown representation, similar to the setup introduced in Guo et al. [2023] . To begin, let f : R d → R m be a fixed representation map that f(x) = (f 1 (x), • • • , f m (x)) ⊤ for any x ∈ R d .The map f can be quite general, which can be regarded as a feature extractor that will be learned by the transformer. We assume that each ICL task corresponds to a map λ ⊤ f (•) that lies in the linear span of those m basis functions in f (•), where λ is generated according to the distribution D λ . Thus, for each ICL instance, the (noisy) label of an input v k (∀k ∈ [K]) is given asEQUATION) where τ > 0 is the noise level.The goal of ICL is to form predictions on query x query given in-context labels of the form (1) on a few inputs, known as prompts. In this paper, we use V to denote the dictionary set that contains all K unit-norm distinct tokens, i.e., V := {v 1 , • • • , v K } ⊂ R d with each token ∥v k ∥ 2 = 1. We assume that each prompt P = P λ provides the first N tokens (with N ≪ K) and their labels, and is embedded in the following matrixEQUATIONwhereV := (v 1 , • • • , v N ) ∈ R d×N (3) head H head 1 … … … … softmax × × Value Key Query T Figure 1:The structure of a one-layer transformer with multi-head softmax attention.is the collection of prompt tokens, and y := (y 1 , • • • , y N ) ⊤ is the prompt label. Given the prompt as the input, the transformer predicts the labels for all the K tokens y 1 , • • • , y K in the dictionary set.Transformer architecture. We adopt a one-layer transformer with multi-head softmax attention [Chen et al., 2024] -illustrated in Figure 1 -to predict the labels of all the tokens in the dictionary V, where H is the number of heads. Denote the query embedding asEQUATIONand denote the embedding of both the prompt and the query as E := (E P , E Q ) ∈ R (d+1)×K . We define the output of each transformer head asEQUATIONwhere d+1) , and W V h ∈ R K×(d+1) are the query, key, and value matrices, respectively, and the softmax is applied column-wisely, i.e., given a vector input x, the i-th entry of softmax(x) is given by e xi / j e xj . The attention map of the transformer T (E) is defined asW Q h ∈ R de×(d+1) , W K h ∈ R de×(EQUATIONwhere W O is the output matrix. Following recent theoretical literature to streamline analysis [Huang et al., 2023 , Nichani et al., 2024 , Deora et al., 2023 , Chen et al., 2024] , we assume that the embedding matrices take the following forms:EQUATIONEQUATIONwherew h = (w h,1 , • • • , w h,K ) ⊤ ∈ R K and Q h ∈ R d×d are trainable parameters for all h ∈ [H].The prediction of the labels is provided by the diagonal entries of T (E), which we denote byy = ( y 1 , • • • , y K ) ∈ R K .Note that y k takes the following form under our parameter specification:EQUATIONTraining via GD. Let θ = {Q h , w h } H h=1 denote all trainable parameters of T . Let ϵ := (ϵ 1 , • • • , ϵ K ) ∈ R m×K denote the noise matrix. Given training data over ICL instances, the goal of training is to predict labels y k for all v k ∈ V. Specifically, we train the transformer using gradient descent (GD) by optimizing the following mean-squared population loss:EQUATION)We apply different learning rates η Q , η w > 0 for updating {Q h } H h=1 and {w h } H h=1 , respectively, i.e., at the t-th (t ≥ 1) step, we have∀h ∈ [H] : Q (t) h = Q (t-1) h -η Q ∇ Q h L(θ (t-1) ), w (t) h = w (t-1) h -η w ∇ w h L(θ (t-1) ), (10) where θ (t) = {Q (t) h , w (t)h } H h=1 is the parameter at the t-th step.Inference time. At inference time, given a prompt P = P λ with N examples, where λ may not be in the support of the generation distribution D λ , the transformer applies the pretrained parameters and predicts the labels of all K tokens without further parameter updating.3 Theoretical Analysis 
training time convergence: 
 In this section, we show that the training loss L converges to its minimum value at a linear rate during training, i.e., the function gap∆ (t) := L(θ (t) ) -inf θ L → 0, t → ∞ (11)at a linear rate, under some appropriate assumptions.Key assumptions. We first state our technical assumptions. The first assumption is on the distribution D λ for generating the coefficient vector λ of the representation maps.Assumption 1 (Assumption on distribution D λ ). We assume that in (1) each entry λ i is drawn independently and satisfiesE[λ i ] = 0 and E[λ 2 i ] = 1 for all i ∈ [m].To proceed, we introduce the following notation:EQUATIONwhere zi is the i-th column vector of Z for i ∈ [N ]. We further defineC (t) k (k ∈ [K], t ∈ N + ) and B (t)k as follows:EQUATION)To guarantee the convergence, we require the initialization of the parameters θ (0) satisfies the following condition.Assumption 2 (Assumption on initialization). For all k ∈EQUATIONk has full row rank. Before stating our main theorem, let us examine when the initialization condition in Assumption 2 is met. Fortunately, we only require the following mild assumption on V to ensure our parameter initialization has good properties. Assumption 3 (Assumption on V ). There exists one row vectorx = (x 1 , • • • , x N ) ⊤ of the prompt token matrix V (cf. (3)) such that x i ̸ = x j , ∀i ̸ = j.Assumption 3 implies that V has distinct tokens, i.e., v j ̸ = v k when j ̸ = k. It is worth noting that Assumption 3 is the only assumption we have on the dictionary V. In comparison, all other theoretical works in Table 1 impose somewhat unrealistic assumptions on V. For example, Huang et al. [2023] , Li et al. [2023] , Nichani et al. [2024] assume that the tokens are pairwise orthogonal, which is restrictive since it implies that the dictionary size K should be no larger than the token dimension d, whereas in practice it is often the case that K ≫ d [Reid et al., 2024 , Touvron et al., 2023] . In addition, Chen et al. [2024] , Zhang et al. [2023a] , Wu et al. [2023] assume that each token is independently sampled from some Gaussian distribution, which also does not align with practical scenarios where tokens are from a fixed dictionary and there often exist (strong) correlations between different tokens.The following proposition states that when the number of heads exceeds the number of prompts, i.e. H ≥ N , we can guarantee that Assumption 2 holds with probability 1 by simply initializing {Q h } H h=1 using Gaussian distribution.Proposition 1 (Initialization of {Q h } H h=1 ). Suppose Assumptions 1, 3 hold and H ≥ N . For any fixed β > 0, letQ (0) h (i, j) i.i.d.∼ N (0, β 2 ), then Assumption 2 holds almost surely.Proof. See Appendix E.1.Choice of learning rates. DefineEQUATION)where ∆ (0) is the initial function gap (c.f. ( 11)). Assumption 2 indicates that ζ 0 > 0. Let γ be any positive constant that satisfiesEQUATION)We set the learning rates asEQUATION)where5 EQUATION)Theoretical guarantee. Now we are ready to state our first main result, regarding the training dynamic of the transformer.Theorem 1 (Training time convergence). Suppose Assumptions 1, 2 hold. We let w (0) k = 0 and set the learning rates as in (16). Then we haveEQUATION)Proof. See Appendix C.Theorem 1, together with Proposition 1, shows that the training loss converges to its minimum value at a linear rate, under mild assumptions of the task coefficients and token dictionary. This gives the first convergence result for transformers with multi-head softmax attention trained using GD to perform ICL tasks (see Table 1 ). Our convergence guarantee (18) also indicates that the convergence speed decreases as the size K of the dictionary or the number H of attention heads increases, which is intuitive because training with a larger vocabulary size or number of parameters is more challenging. However, a small H will limit the expressive power of the model (see Section 3.3 for detailed discussion), and we require H ≥ N to guarantee Assumption 2 holds, as stated in Proposition 1.We now move to examine the inference time performance, where the coefficient vector λ corresponding to the inference task may not drawn from D λ . In fact, we only assume that the coefficient vector λ at inference time is bounded as in the following assumption.Assumption 4 (Boundedness of λ at inference time). We assume that at inference time ∥λ∥ 2 ≤ B for some B > 0.For notational simplicity, letEQUATION)The following theorem characterizes the performance guarantee of the transformer's output y (after sufficient training) at the inference time.Theorem 2 (Inference time performance). Let λ be the solution to the following ridge regression problem:EQUATION)Under the assumptions in Theorem 1, for any ε > 0 and δ ∈ (0, 1), if the number of training iterates T satisfiesEQUATIONthen given any prompt P that satisfies Assumption 4 at the inference time, with probability at least 1 -δ, the output of the transformer y satisfiesEQUATION)Proof. See Appendix D.In Theorem 2, ( 22) shows that after training, the transformer learns to output the given labels of the first N tokens in each prompt, and more importantly, predicts the labels of the rest K -N tokens by implementing the ridge regression given in (20). Note that Akyürek et al. [2022] studied the expressive power of transformers on the linear regression task and showed by construction that transformers can represent the closed-form ridge regression solution. Interestingly, here we show from an optimization perspective that transformers can in fact be trained to do so.Generalization capabilities of the pretrained transformer. Theorem 2 captures two generalization capabilities that the pretrained transformer can have.i) Contextual generalization to unseen examples: Theorem 2 suggests that the transformer exploits the inherent contextual information (to be further discussed in Section 3.3) of the function template in the given prompt, and can further use such information to predict the unseen tokens. ii) Generalization to unseen tasks: Theorem 2 also suggests that the pretrained transformer can generalize to a function map corresponding to any λ ∈ R m at the inference time (albeit satisfying Assumption 4), which is not necessarily sampled from the support of its training distribution D λ .We note that the contextual generalization that the transformer has here is different in nature from the prediction ability shown in previous works on ICL [Huang et al., 2023 , Chen et al., 2024 , Li et al., 2024 , Nichani et al., 2024] . Those work focuses on a setting where each prompt contains a good portion of tokens similar to the query token, allowing the transformer to directly use the label of the corresponding answers from the prompt as the prediction. However, in practical scenarios, prompts often contain only partial information, and our analysis sheds lights on explaining how transformers generalize to unseen examples by leveraging ridge regression to infer the underlying template.How does the representation dimension affect the performance? Beyond the above discovery, several questions are yet to be explored. For instance, while we demonstrate that transformers can be trained to implement ridge regression, how good is the performance of the ridge regression itself? What is the best choice of ridge regression we could expect? How close is the transformer's choice to the best possible choice? We address these questions as follows.Given any prompt P at inference time, since there is no label information about the rest K -N tokens, the best prediction we could hope for from the transformer shall beEQUATIONwhere Z Q is defined in ( 19), and λ τ satisfies:EQUATION)In other words, we hope the transformer outputs the given N labels as they are. For the rest K -N labels, the best we could hope for is that the transformer estimates the coefficient vector λ by solving the above regression problem to obtain λ τ , and predict the k-th label by λ 24) is equivalent to the following ridge regression problem (see Lemma 4 in the appendix for its derivation):⊤ τ f (v k ) for k = N + 1, • • • , K. Note that (EQUATION)The only difference between the two ridge regression problems ( 20) and ( 25) is the coefficient of the regularization term. This indicates that at the training time, the transformer learns to implement ridge regression to predict the labels of the rest K -N tokens, assuming the noise level is given by m N τ . This observation also reflects how the sequence length N affects the transformer's preference for choosing templates and its performance at inference time:• The closer m is to N , the closer the transformer's choice of templates is to the best possible choice, and the better the transformer's prediction will be;• When N < m, the transformer tends to underfit by choosing a λ with small ℓ 2 -norm;• When N > m, the transformer tends to overfit since it underestimates the noise level and in turn captures noise in the prediction.We provide more interpretation on our results, which may lead to useful insights into the ICL ability of the transformer.How does the transformer gain ICL ability with representations? Intuitively speaking, our pretrained transformer gains in-context ability by extracting and memorizing some "inherent information" of all basic function maps f i (i ∈ [m]) during the training. Such information allows it to infer the coefficient vector λ from the provided labels in each prompt and calculate the inner product ⟨λ, f (v k )⟩ to compute y k given any token v k ∈ V at inference time. To be more specific, the "inherent information" of all basic tasks could be described by the N -by-K matrix A defined as follows (see also (34)):A := Z ⊤ Z + mτ I N -1 Z ⊤ Z + (mτ I N , 0) ∈ R N ×K ,whereZ := (f (v 1 ), • • • , f (v K )) = (Z, Z Q ) ∈ R m×K .During training, the transformer learns to approximate A :,k byH h=1 w h,k softmax(V ⊤ Q h v k ) for each k ∈ [K].To further elaborate, we take a closer look at the special case when the labels do not contain any noise, i.e., τ = 0, and N ≥ m. In this case, A becomes Z † Z, and given any prompt P = P λ , the coefficient vector λ could be uniquely determined from the provided token-label pairs in the prompt. It is straightforward to verify that the label of each token v k could be represented by the inner product of the given label vector y and the k-th column of Z † Z, i.e.,EQUATION)Comparing the above equation with (8), it can be seen that in order to gain the in-context ability, the transformer needs to learn an approximation of Z † Z :,k byH h=1 w h,k softmax(V ⊤ Q h v k ) for each k ∈ [K].More generally, in the proof of Theorem 2, we show thatEQUATIONcomparing which with (8) suggests that a small training error implies thatH h=1 w h,k softmax(V ⊤ Q h v k ) is close to A :,k .In fact, this is the necessary and sufficient condition for the training loss to be small. A rigorous argument is provided in Lemma 5.The necessity and trade-offs of multi-head attention mechanism. Multi-head attention mechanism is essential in our setting. In fact, it is generally impossible to train a shallow transformer with only one attention head to succeed in the ICL task considered in our paper. This is because, as we have discussed above, the key for the transformer is to approximate A :,k byH h=1 w h,k softmax(V ⊤ Q h v k ) for each k ∈ [K]. If H = 1, the transformer could not approximate each A :,k by w 1,k softmax(V ⊤ Q 1 v k ) ingeneral since the entries of the latter vector are either all positive or all negative. In addition, Proposition 1 indicates that when H ≥ N , the weights of the transformer with a simple initialization method satisfy our desired property that is crucial to guarantee the fast linear convergence. However, ( 18) implies that we should not set H to be too large, since larger H yields slower convergence rate. 
conclusion: 
 We analyze the training dynamics of a one-layer transformer with multi-head softmax attention trained by gradient descent to solve complex non-linear regression tasks using partially labeled prompts. In this setting, the labels contain Gaussian noise, and each prompt may include only a few examples, which are insufficient to determine the underlying template. Our work overcomes several restrictive assumptions made in previous studies and proves that the training loss converges linearly to its minimum value. Furthermore, we analyze the transformer's strategy for addressing the issue of underdetermination during inference and evaluate its performance by comparing it with the best possible strategy. Our study provides the first analysis of how transformers can acquire contextual (template) information to generalize to unseen examples when prompts contain a limited number of query-answer pairs.  
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Experiment details are included in Section A.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: stochasticity is not critical in our experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: the paper aims to provide a better understanding on existing algorithms and thus poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: the paper does not use existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: This section aims to provide some empirical validation to our theoretical findings and verify that some of our results could be generalized to deeper transformers.Setup. We conduct experiments on a synthetic dataset, where we randomly generate each token v k and their representation f (v k ) from standard Gaussian distribution. We employ both the 1-layer transformer described in Section 2 and a standard 4-layer transformer in Vaswani et al. [2017] with Training and inference performance. We set N = 30, K = 200, d = 100, m = 20, and set H to be 64 and 8 for 1-layer and 4-layer transformers, respectively. Figure 2 shows the training and inference losses of both 1-layer and 4-layer transformers, where we measure the inference loss byd1 K ∥ y -y ⋆ ∥ 22 to validate (22): after sufficient training, the output of the transformer y converges to y ⋆ . From Figure 2 we can see that for both 1-layer and 4-layer transformers, the three curves have the same descending trend, despite the inference loss on the ood dataset is higher than that on the in-domain dataset. This experiment also shows the transformer's contextual generalization to unseen examples and to unseen tasks, validating our claim in Section 3.2.Figure 3 plots the performance gap 1 K y ⋆ -y best 2 2 of the one-layer transformer with respect to different N ranging from 50 to 150, when we fix m = 100 and τ = 0.01. This verifies that the ridge regression implemented by the pretrained transformer has a better performance when m is close to N , again verifying our claim at the end of Section 2.Impact of the number of attention heads. We now turn to examine the impact of the number of attention heads. In this experiment, we use the population loss (9), and set the other configurations same as those in Figure 2 . Figure 4 shows the training loss curves for different H with respect the iteration number, which validates our claims. From Figure 4 , we can see that we need to set H large enough to guarantee the convergence of the training loss. However, setting H too large (H = 400) leads to instability and divergence of the loss. Recall that in Proposition 1, we require H ≥ N to guarantee our convergence results hold. Although this condition may not be necessary, Figure 4 shows that when H < N = 30, the loss stopped descending even when it is far from the minimal value. On the other side, the loss keeps descending when H = 30 (though slowly). We also explore how H affects the training of the 4-layer transformer, as displayed in Figure 5 , where we set K = 200 and the configurations other than H are the same as in Figure 3 . We fix the wall-clock time to be 100 seconds and plot the training loss curves with different H. to different H within 100s may provide some explanation to this phenomenon: (i) transformers with larger H could complete less iterations within a fixed amount of time (the curves corresponding to larger H are shorter); (ii) the training loss curves corresponding to large H (H = 32, 64) descend more slowly. This suggests our claim that larger H may yield slower convergence rate is still valid on deeper transformers. Note that unlike the 1-layer transformer, deeper transformers don't require a large H to guarantee convergence. This is because deeper transformers have better expressive power even when H is small.We summarize the frequently used notation in Table 2 for ease of reference.We provide some useful facts that will be repeatedly used later on. Let 12), we can rewritez k := f (v k ) = (f 1 (v k ), f 2 (v k ), • • • , f m (v k )) ⊤ ∈ R m , ∀k ∈ [K]. Recalling (Z := (z 1 , • • • , z N ) ∈ R m×N .We further define s h k ∈ R N as follows: This experiment shows that unlike 1-layer transformers, deeper transformers don't require H to be large to guarantee convergence of the loss.EQUATIONnotation meaningK ∈ N + total number of tokens d ∈ N + token dimension m ∈ N + number of basic tasks H ∈ N + number of attention heads N ∈ N + number of examples in each prompt v k ∈ R d , k ∈ [K] the k-th token f i : R d → R, i ∈ [m] the i-th basic task λ ∈ R m coefficient vector y k = λ ⊤ (f (v k ) + ϵ k ), k ∈ [K]the k-th label Table 2 : Notation for key parameters.Lemma 1 (Softmax gradient). For all j ∈ [N ], k ∈ [K] and h ∈ [H], we haveEQUATIONwhere s h jk is defined in (28).Proof. See the proof of Lemma A.1 in Huang et al. [2023] .Lemma 2 (Smoothness of softmax). For vectors ξ 1 , ξ 2 ∈ R l , we haveEQUATIONProof. See Corollary A.7 in Edelman et al. [2022] .We also need to make use of the following form of Young's inequality. Lemma 3. For anyEQUATION)The following lemma shows the equivalence between ( 24) and (25).Lemma 4 (Equivalence of the regression problems). Given any promptP λ := (v 1 , y 1 , • • • , v N , y N ),we have the following equivalence:EQUATION)Proof. See Appendix E.2.We first outline the proof. To prove Theorem 1, we first remove the expectation in the expression of the loss function L in ( 9) by reformulating it to a deterministic form (see Lemma 5). With this new form, we show by induction that the loss function L is smooth (Lemma 10) and satisfies the Polyak-Łojasiewicz (PL) condition (c.f. ( 49)). Provided with both smoothness and PL conditions, we are able to give the desired linear convergence rate [Karimi et al., 2016] .We defineEQUATION)We also define the following matrices:EQUATIONEQUATION)whereZ := (z 1 , • • • , z K ) ∈ R m×K .We first reformulate the loss function to remove the expectation in the population loss.Lemma 5 (Reformulation of the loss function). Under Assumption 1, the loss function L(θ) could be rewritten into the following equivalent form:EQUATIONwhereL ⋆ = 1 2K N k=1 -Z ⊤ z k + mτ e k ⊤ Z ⊤ Z + mτ I -1 Z ⊤ z k + mτ e k + ∥z k ∥ 2 2 + mτ + 1 2K K k=N +1 -Z ⊤ z k ⊤ Z ⊤ Z + mτ I -1 Z ⊤ z k + ∥z k ∥ 2 2is a constant that does not depend on θ, and Z is defined in (12).Proof. See Appendix E.3.Lemma 5 indicates that L ⋆ is a lower bound of L. We'll later show that L ⋆ is actually the infimum of L, i.e., L ⋆ = inf θ L(θ).Lemma 5 also indicates that, the necessary and sufficient condition for L(θEQUATIONwhich folllows immediately that (37) is equivalent toEQUATION)To simplify the analysis, we introduce the following reparameterization to unify the learning rates of all parameters, and we'll consider the losses after reparameterization in the subsequent proofs. Lemma 6 (Reparameterization). DefineEQUATIONand letξ := {Q h , α h } H h=1 , ℓ(ξ) := L(θ).(40) Then (10) is equivalent toξ (t) = ξ (t-1) -η Q ∇ ξ ℓ(ξ (t-1) ), ∀t ∈ [T ].(41)Proof. See Appendix E.4.We denote α as α := (α h,k ) h∈[H],k∈[K] ∈ R H×K .The following lemma bounds the gradient norms by the loss function, which is crucial to the proof of Theorem 1.Lemma 7 (Upper bound of the gradient norms). Suppose Assumption 1 holds and |αEQUATION)Proof. See Appendix E.5. Now we are ready to give the main proof.Proof of Theorem 1. To prove Theorem 1, it suffices to prove that under our assumptions, we have (Upper bound of the parameters:)EQUATION(Lower bound of eigenvalues:)EQUATION(Linear decay of the loss:EQUATIONwhereEQUATIONand γ, α h is defined in (39), ζ 0 is defined in ( 14). We shall prove (43a),( 43b) and (43c) by induction.Base case. It is apparent that (43a),( 43b) and (43c) all hold when t = 0.Induction. We make the following inductive hypothesis, i.e., when s ∈ [t -1], (43a), ( 43b) and (43c) hold. Below we prove that (43a),( 43b) and (43c) hold when s = t by the following steps.Step 1: verify (43b) and the Polyak-Łojasiewicz condition. We first compute the gradient of the loss w.r.t. α:EQUATION)where the first equality follows from Lemma 5, C k , B k is defined in (13).EQUATION)where the third line uses Lemma 2, and that ∀h ∈ [H] :EQUATION)where the second inequality follows from Lemma 7 (cf. ( 42)) and the third inequality follows from the inductive hypothesis and the fact that ℓ(ξ (s) ) = L(θ (s) ), ∀s. Combining ( 47) with ( 46), we haveEQUATION)where the last inequality follows from (15). The above inequality (48) indicates that∀x ∈ R K : x ⊤ B (t) k 2 ≥ x ⊤ B (0) k 2 -x ⊤ (B (t) k -B (0) k ) 2 ≥ ζ 0 /2,which gives (43b).Therefore, we obtain the following PL condition:EQUATION)where the equality comes from (45), and the last equality follows from (36).Step 2: verify the smoothness of the loss function. We first give the following lemma that bounds the Lipschitzness of b h k and δ θ k , which will be used later on. For notation simplicity, we let B, Q, α denote B(θ), Q(θ), α(θ), respectively, and letEQUATIONEQUATION)Proof. (50) follows from a similar argument in (46). Regarding the Lipschitzness of δ θ k , we haveδ θ k -δ θ ′ k 2 = γ H h=1 α h,k (s h k (θ) -s h k (θ ′ )) + H h=1 (α h,k -α ′ h,k )s h k (θ ′ ) 2 ≤ γ H h=1 |α h,k | s h k (θ) -s h k (θ ′ ) 2 + γ H h=1 |α h,k -α ′ h,k | s h k (θ ′ ) 2 ≤ 2γ √ Hα H h=1 ∥Q h -Q ′ h ∥ 2 F + γ √ H ∥α k -α ′ k ∥ 2 ,where we use ( 46) again to bound the first term in the second line, and use the fact that s h k (θ ′ ) 2 ≤ 1 and Cauchy-Schwarz inequality to bound the second term in the second line.We also need the following lemma which bounds the norm of B k and δ θ k .Lemma 9 (Upper bounds of b h k and δEQUATION)where A is defined in (34).Proof. (52) follows fromb h k 2 ≤ Z 2 s h k 2 ≤ Z 2 .(53) follows fromδ θ k 2 ≤ γ H h=1 |α h,k | s h k 2 + ∥Ae k ∥ 2 ≤ γHα + ∥A∥ 2 .As a consequence of Lemma 8 and Lemma 9, For all k ∈ [K], and all transformer parameters θ, θEQUATIONfrom which we obtain the smoothness of the ℓ w.r.t. α as follows:EQUATIONwhere the first inequality uses Young's inequality (c.f. Lemma 3).To obtain the smoothness of the loss function w.r.t. Q h , we first note that by (82) we haveEQUATIONwhere the third inequality uses Cauchy-Schwarz inequality. Combining the above inequality (57) with Lemma 8 and Lemma 9, we haveEQUATIONwhere the last line uses (46) to bound s h k (θ) -s h k (θ ′ ) 2 . The above inequality (58) further givesEQUATIONwhere the first inequality makes use of Young's inequality (c.f. Lemma 3).Combining the above two relations ( 55) and ( 59), we obtain the smoothness of ℓ w.r.t. ξ as follows:Lemma 10 (Smoothness of the loss function).EQUATIONL 2 = 2 1 K 2γ Z 2 2 (2γHα + ∥A∥ 2 ) 2 + γ 4 K 2 H 2 Z 4 2 + 8γ fmax Z 2 • max 1, (2 √ Kα) 2 (αγH) 2 + 4N (αγH + ∥A∥ 2 ) 2 .(61)Step 3: verify (43a). ( 45) implies∂ℓ(ξ) ∂α h,k = γ K (b h k ) ⊤ Zδ θ k ,which, combining with (52), gives∀k ∈ [K], h ∈ [H] : ∂ℓ(ξ) ∂α h,k 2 ≤ γ 2 K 2 Z 2 2 Zδ θ k 2 2 .Combining this with (36) we obtainEQUATIONTherefore, we haveα (t) h 2 = α (0) h -η Q t-1 i=0 ∂ℓ(ξ (i) ) ∂α h 2 ≤ α (0) h 2 + η Q t-1 i=0 ∂ℓ(ξ (i) ) ∂α h 2 ≤ α (0) h 2 + η Q Z 2 2γ 2 K t-1 i=0 ℓ(ξ (i) ) -L ⋆ ≤ α (0) h 2 + η Q Z 2 2γ 2 L(θ (0) ) -L ⋆ K t-1 i=0 1 - η Q σ 2 i ≤ α (0) h 2 + η Q Z 2 2γ 2 L(θ (0) ) -L ⋆ K • 4 η Q σ ,where the second inequality follows from (62) and the third inequality follows from the induction hypothesis (43c). (43a) follows from plugging σ defined in (44) into the above inequality and using the initializtion condition that α (0) = 1 γ w (0) = 0.Step 4: verify the linear convergence rate (43c). Combining (43a), ( 60) and Lemma 4.3 in Nguyen and Mondelli [2020] , we haveEQUATIONwhich indicates when η Q ≤ 1/L, we haveEQUATIONwhich, combined with the fact that L(θ (s) ) = ℓ(ξ (s) ) for all s (see Lemma 6), verifies (43c).Note that (36) implies that L ⋆ ≤ L(θ) holds for all θ. And from (43c) we know that L(θ (t) ) → L ⋆ as t → ∞. Therefore, it follows thatL ⋆ = inf θ L(θ).Consequently, (43c) is equivalent to (18).By (43c) we know that L(θ (t) ) → L ⋆ as t → ∞. Thus from (36) we know that (37) and ( 38) hold.By Sherman-Morrison-Woodbury formula, we haveEQUATION)Thus we haveEQUATION)where Z Q is defined in ( 19).On the other hand, it's straightforward to verify that λ defined in (20) admits the following closed form: λ = mτ I m + ZZ ⊤ -1 Zy.(67) Combining the above two equations, we obtainA ⊤ y = y Z Q ⊤ mτ I m + ZZ ⊤ -1 Zy = y Z Q ⊤ λ = y ⋆ ,where the last equality follows from ( 22). Now we give the iteration complexity for the mean-squared error between the prediction y and the limit point y ⋆ to be less than ε. Given any prompt P = P λ , where λ satisfies Assumption 4, we havey i = λ ⊤ (z i + ϵ i ) ∼ N (λ ⊤ z i , ∥λ∥ 2 2 τ ).Lettingx i = yi-λ ⊤ zi ∥λ∥ 2 √ τ , we have x i ∼ N (0, 1). Define Z = N i=1 ∥λ∥ 2 2 τ (x 2 i -1) = y -Z ⊤ λ 2 2 -N τ ∥λ∥ 2 2 .By Laurent and Massart [2000, Lemma 1], we have∀s > 0 : P Z ≥ 2 √ N ∥λ∥ 2 2 τ √ s + 2 ∥λ∥ 2 2 τ s ≤ exp (-s) .By letting s = log(1/δ) and using the definition of Z, we haveEQUATION)Thus with probability at least 1 -δ, we haveEQUATION)where we use (68) in the second inequality, and the third inequality follows from Assumption 4.On the other hand, by (36) we haveEQUATION)Thus we know that w.p. at least 1 -δ, we have1 2K ∥ y -y ⋆ ∥ 2 2 = 1 2K A -A ⊤ y 2 2 ≤ 1 2K A -A 2 2 ∥y∥ 2 2 ≤ ε,where the last relation follows from ( 69), ( 70) and ( 21).E.1 Proof of Proposition 1For notation simplicity we drop the superscript (0) in the subsequent proof.EQUATIONThis suggests the column vectors of D k are i.i.d. and the density of each column vector is positive at any point x ∈ R(V ), where R(V ) ⊂ R N is the row space of V .Since Z has full rank, to prove B k has full rank a.s., we only need to argue that C k (:, 1 : N ) has full rank w.p. 1. Below we prove this by contradiction (recall that by definition C k = softmax(D k ), and we assume H ≥ N ).Suppose w.p. larger than 0, there exists one of C k (:, 1 : N )'s column vector that could be linearly represented by its other N -1 column vectors. Without loss of generality, we assume this colomn vector is C k (:, 1) = softmax(D k (:, 1)). Let x = x(q 1 ) := exp(D k (:, 1)) = exp(V ⊤ q 1 ). Then x could be linearly represented by exp(D k (:, i)), i = 2, • • • , N .Let Ã := exp(D k (:, 2 : N )), then w.p. larger than 0, x ∈ C( Ã), where C( Ã) is the column vector space of Ã. i.e., we haveR N ×(m-1) P(x ∈ C( Ã)| Ã)dµ( Ã) > 0,which further indicates that there exists Ã ∈ R N ×(N -1) such that P(x ∈ C( Ã)) > 0. Since the dimension of C( Ã) is at most N -1, there exists y ∈ R N , y ̸ = 0 such that y⊥C( Ã). Therefore, we haveEQUATION)By Assumption 3, without loss of generality, we assume thatu 1 = (v 11 , v 12 , • • • , v 1N ) ⊤ has different entries. For any vector w = (w 1 , • • • , w d ) ⊤ ∈ R d , we let w = (w 2 , • • • , w d ) ⊤ ∈ R d-1denote the vector formed by deleting the first entry of w. Let q 1 = (q, q⊤ 1 ) ⊤ . For any fixed q1 ∈ R d-1 , the function g(•| q1 ) : R → R defined byg(q| q1 ) := N i=1 y i e qv1i+ q⊤ 1 ṽi = N i=1y i e q⊤ 1 ṽi e qv1i = y, exp(V ⊤ q 1 ) = ⟨y, x(q 1 )⟩ has finite zero points and thus {q ∈ R|g(q| q1 ) = 0} is a zero-measure set. Therefore, we haveP(⟨y, x⟩ = 0) = R d-1 P(g(q| q1 ) = 0| q1 )dµ( q1 ) = 0,which contradicts (72). Therefore, C k (:, 1 : N ) has full rank with probability 1.Lemma 4 can be verified by the following direct computation (recall that the noise in each label satisfies ϵi i.i.d ∼ N (0, τ I m ), ∀i ∈ [N ]): E ϵ 1 2N N i=1 (y i -λ ⊤ (f (v i ) + ϵ i )) 2 = E ϵ 1 2N N i=1 (y i -λ ⊤ f (v i )) 2 -2λ ⊤ ϵ i (y i -λ ⊤ f (v i )) + λ ⊤ ϵ i ϵ ⊤ i λ = 1 2N N i=1 (y i -λ ⊤ f (v i )) 2 + τ ∥λ∥ 2 2 = 1 2N N i=1 (y i -λ ⊤ f (v i )) 2 + τ 2 ∥λ∥ 2 2 .We let ϵEQUATIONandEQUATIONwhere a k denote the k-th column vector of matrix A(θ) defined in (35), and the fifth line uses Assumption 1.Note that for all k ∈ [K], we haveEQUATIONand thatEQUATIONwhere 1 {k ∈ [N ]} is the indicator function that equals 1 if k ∈ [N ] and 0 otherwise, and we have made use of the assumption that ϵ k i.i.d.∼ N (0, τ 2 I m ).Combining the above two equations with (75), we know that for k ∈ [N ], it holds thatL k (θ) = 1 2 ∥Z a k -z k ∥ 2 2 + mτ ∥ a k -e k ∥ 2 2 .Reorganizing the terms in the RHS of the above equation, we obtain thatEQUATIONwhereEQUATIONwhere 78), (79) together with (33) and the definition of L ⋆ give (36).c ′ k = -Z ⊤ z k ⊤ Z ⊤ Z + mτ I -1 Z ⊤ z k + ∥z k ∥ 2 2 . (First, it holds thatEQUATIONSecond, note thatw (t) h = w (t-1) h -η w ∇ w h L(θ (t-1) ) = γα (t-1) h -γ 2 • 1 γ η Q ∇ α h ℓ(ξ (t-1) ) = γ α (t-1) h -η Q ∇ α h ℓ(ξ (t-1) ) .Dividing both sides of the above equality by γ, we haveEQUATIONHence, (41) follows from combining (80) and ( 81).Throughout this proof, we omit the superscript (t) for simplicity. We first compute the gradient of L w.r.t. Q h . By (36) we know thatℓ(ξ) = L(θ) = 1 2K K k=1 Zδ k 2 2 ,and thus we haveEQUATIONNote thatEQUATIONwhere we use the fact that ( 82) and ( 83), we have the desired resultv j -v i )v ⊤ k 2 ≤ 2 (recall that we assume each v k has unit norm, k ∈ [K].) Combining (EQUATION)where fmax is defined in ( 12) and the third line follows from Cauchy-Schwarz inequality.Answer: [No] Justification: The experiments are very simple and can be easily reproduced by following the instructions in the paper.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.We leave a tighter, but more complicated, expression of L in the appendix (cf. (61)) in the appendix and present a simplified form in the main paper for readability. 


Decision: This paper provides a theoretical account of in-context learning in a single-layer transformer performing nonlinear regression tasks. This moves beyond prior theoretical accounts along some dimensions — though assumptions like a single layer are still quite limiting. The reviewers generally agreed that the paper offered a useful contribution despite these limitations, especially after clarification in the rebuttal period, but that the papers impact might be limited because of both its strong assumptions, and the relative lack of empirical validation or further practical insight provided.

Reviews: review: paper_summary: paper_summary: This paper presents a rigorous theoretical analysis of the training dynamics and generalization capabilities of a one-layer transformer with multi-head softmax attention for in-context learning (ICL) of non-linear regression tasks. The authors consider a more challenging and realistic setting where prompts contain only a small number of noisy labeled examples, insufficient to fully determine the underlying template function. This addresses limitations in previous theoretical work that often required unrealistically large numbers of examples per prompt.

A key theoretical contribution is proving convergence of the training loss to its global minimum at a linear rate when optimizing with gradient descent. This is the first such convergence result for transformers with multi-head softmax attention on ICL tasks. The analysis reveals that the transformer effectively learns to perform ridge regression during training. The paper provides precise bounds on the mean squared error between the transformer's predictions and those of the optimal ridge regression solution, as a function of the number of training iterations.

The authors demonstrate that multi-head attention is crucial for succeeding at this ICL task. They provide a lower bound on the number of heads needed (H ≥ N), while also noting that too many heads can slow down convergence. The analysis shows how the multi-head mechanism allows the transformer to approximate the required matrix operations for ridge regression.
strength_weakness: strength_weakness: - Limited model architecture (one-layer transformer) which is understandable

- Lack of empirical validation:  The paper is purely theoretical and does not include any experimental results to validate its predictions.

- Focus on regression : The analysis is limited to regression tasks, which represent only a subset of the problems typically addressed by transformers and in-context learning. Many real-world applications involve classification, sequence generation, or more complex structured prediction tasks. It's not immediately clear how the insights about ridge regression would translate to these other task types.

Lack of coomparison to other ICL approaches: The paper doesn't provide a comprehensive comparison to other theoretical approaches to in-context learning. While it does mention some previous work, a more in-depth discussion of how this approach relates to or improves upon other theoretical frameworks for ICL would have provided valuable context.
questions: questions: Have the authors considered adding any toy domain experiments that would illustrate and further validate their claims ?
limitations: limitations: - Simplified model: The analysis is limited to a one-layer transformer, which is much simpler than state-of-the-art models used in practice.
- Lack of empirical validation: While the theoretical results are impressive, the paper doesn't include experimental results to validate its predictions.
- Focus on regression: The paper only considers regression tasks, and it's not immediately clear how these results would generalize to classification or more complex tasks.
paper_summary: paper_summary: The paper investigates the theoretical understanding of in-context learning, focusing on whether transformers can generalize to unseen examples within a prompt by acquiring contextual knowledge. The paper analyzes the training dynamics of transformers using non-linear regression tasks, demonstrating that multi-head transformers can predict unlabeled inputs given partially labeled prompts. The study shows that the training loss for a shallow multi-head transformer converges linearly to a global minimum, effectively learning to perform ridge regression. Main contributions include: (1) The paper establishes the convergence guarantee of a shallow transformer with multi-head softmax attention trained with gradient descent on general non-linear regression in-context learning tasks. (2) The paper analyzes the transformer’s behavior at inference time after training, demonstrating that the transformer decides its generating template by performing ridge regression. (3) The analysis framework allows overcoming several assumptions made in previous works, such as the requirement for large prompt lengths, orthogonality of data, restrictive initialization conditions, special structure of the transformer, and super wide models.
strength_weakness: strength_weakness: While the paper excels in theoretical analysis, it lacks empirical validation of the proposed concepts and results. It is suggested to include empirical experiments that demonstrate the practical applicability of the theoretical findings would significantly strengthen the paper. For instance, running experiments on standard benchmarks to show how the proposed theoretical insights translate to improved performance in real-world tasks would validate the results more strongly.

The paper’s discussion on the practical implications of the theoretical findings is somewhat limited. This can make it challenging for practitioners to understand how to apply these insights to real-world problems.
questions: questions: 1. Do you have any empirical results or planned experiments that demonstrate the practical applicability of your theoretical findings?

2. What are the potential impacts of the remaining assumptions and simplifications (e.g., specific initialization conditions, structure of the transformer) on the generalizability of your results?
limitations: limitations: Please review the weaknesses.
paper_summary: paper_summary: This paper tackles the problem of task learning with representations using transformers. It presents a theoretical proof for convergence of training of a single-layer multi-head softmax attention transformer to the global minimum of the loss. Specifically, given $N$ example tokens and their corresponding task outputs, the transformer learns to predict the task output for unseen tokens $K-N$, where $K$ is the size of the dictionary. The paper shows that when trained for sufficiently long, the transformers can perform of ridge regression on the example tokens to label the unseen tokens.
strength_weakness: strength_weakness: (1) The optimal transformer configuration $(\boldsymbol{\theta})$ is not explicitly stated; only the functional form of the output is given, which makes it difficult to the interpret the result in the parameter $(\boldsymbol{\theta})$ space.

(2) The result in this paper is about the wide single-layer transformers (more heads), rather than the deep ones (more layers), and relies on Assumption 3, which holds when the number of heads, $H\ge N$ (Proposition 2). This is restrictive and there might be other ways in which the assumption can hold even if $N > H$. Some empirical demonstrations with shallow transformers on simple experiments might support this better.
questions: questions: (1) Proposition 1 provides a way to make (almost) sure that Assumption 3 holds, but requiring $H \ge N$ seems very demanding in terms of compute. Can you comment on how to obtain similar results by simpler assumption?

(2) I am curios about the motivation behind your problem formulation, i.e., task-based labelling of a finite dictionary? Is there a real-world problem where this might be of use?
limitations: limitations: (1) One minor limitation is that the performance of the trained transformer depends on the relation between the number $N$ of example tokens and the dimension $m$ of the representational space (lines 295-299).
paper_summary: paper_summary: This paper studies the training dynamics of multi-head Transformers by gradient descent on non-linear regression tasks via ICL. This work shows the linear convergence to the global minimum of the training and the in-context inference. An impressive contribution is that transformers are proven to learn contextual information to generalize to unseen examples when prompts contain a small number of input-label pairs. I am willing to update my review if my concerns are addressed.
strength_weakness: strength_weakness: 1. It seems that the assumption that $H\geq N$ in Proposition 1 is too strong. The single-head works in Table 1 do not need this condition for few-shot ICL (Obviously, for single-head, $1<N$).  I know the problem to solve is different, but this means that Transformers cannot handle long contexts without a large number of heads. 

2. No experiments are conducted to verify the theory. It is important to verify the theory by experiments, even with the most simplified Transformer model. For example, can you show if $H<N$, the learning will fail, while if $H\geq N$, the learning becomes successful by a one-layer Transformer?

3. Some statements related to the existing works are wrong. For example, in [Li et al., 2024], $N=\mathcal{O}(\epsilon^{-2}T)$ is the number of data for the pretraining to enable ICL for the Transformer. It is not the number of examples in each prompt as stated in line 129 of the submitted manuscript. The number of examples in each prompt in [Li et al., 2024] is $l\_{tr}$ and $l\_{ts}$ in Eqn 9 of Theorem 3.3.

4. I am unsure whether the theoretical proof is correct. For example, I don't know how you derive linear convergence of the loss via only PL condition and the smoothness of the loss function. Please see Question 1.
questions: questions: 1. Regarding Weakness 1, why do you have $\eta_Q$ in Eqn 63? I think by (41), (60), and Lemma 4.3 in [Nguyen and Mondelli, 2020], $\eta_Q$ should be replaced by $1$ in Eqn 63.

2. What is the requirement for the initialization of $w\_{h,k}$ defined in (7)?

3. It seems there is no requirement or assumption for the $f$ function to be learned. Is your bound derived for the worst case of all the possible $f$?
limitations: limitations: There is no potential negative societal impact of this work.

Comments: title: kind reminder as the rebuttal deadline approaches
Re: Re: Reviewer vG3D
Re: Reviewer vG3D
comment: Thanks for your rebuttal and the additional toy experiments, which I find sensible (as exemplified by discussion with Reviewer vG3D) and illustrative. I have therefore increased my score.
Many thanks for your feedback and increasing the score to support our paper. Regarding $H \ge N$, our numerical result somewhat indicates that $H$ cannot be much smaller than $N$ to guarantee a small loss. The main reason is because the one-layer transformer may not have substantial expressive power and hence could by nature requires a sufficiently large $H$. We suspect that such an assumption may not largely come from our technical tools, but may likely be due to the problem itself. It is great to hear you like the convergence analysis of this work, which we highly appreciate.
Dear Reviewer TyrV,

The author-reviewer discussion period will end in less than one day. We sincerely hope to receive your feedback and see if our responses have properly addressed your concerns. If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Best Regards, Authors
Dear Reviewer MzSs,

Thank you very much for your response!

In addition to what you mentioned as 'memorize' the basis vectors, our results contain two more points: (i) after 'memorizing/learning' basis vectors, the ICL implements **ridge regression** to use such basis information, which is a new characterization; (ii) such tranformers having those properties (memorizing basis vectors and implementing ridge regression) can be learned naturally by gradient descent.

Many thanks, 
Authors
Thank you for the valuable comments, and for the clarification on the limitation (the relation between $m$ and $N$ affecting the performance) to be a strength of the paper. I agree with you.

Despite the clean proof and the result of the fast convergence of the training, the finite dictionary still seems to be a limitation as the transformer can essentially 'memorize' the basis vectors into the value transforms in the heads. As such, the paper is more about how an attention mechanism can 'memorize' efficiently (possibly better than other architectures), which I agree to be novel. However, direct impact of the paper is unclear. Moreover, the construction using a large number of heads is not very appealing. 

Taking all of these into consideration, I shall keep my (positive) score.
Dear Reviewer TyrV,

As the author-reviewer discussion period will end soon, we would like to check whether our responses have properly addressed your concerns? If so, could you please kindly consider increasing your initial score accordingly? Certainly, we are more than happy to answer your further questions.

Thank you for your time and effort in reviewing our work!

Best Regards, 
Authors
Dear Reviewer MzSs,

We've taken your initial feedback into careful consideration in our response. Could you kindly confirm whether our responses have appropriately addressed your concerns?

If you find that we have properly addressed your concerns, could you please kindly consider increasing your initial score accordingly? Please let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors
Dear Reviewer dubr,

We've taken your initial feedback into careful consideration in our response. Could you kindly confirm whether our responses have appropriately addressed your concerns?

If you find that we have properly addressed your concerns, could you please kindly consider increasing your initial score accordingly? Please let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors
Dear Reviewer TyrV,

We've taken your initial feedback into careful consideration in our response. Could you kindly confirm whether our responses have appropriately addressed your concerns?

If you find that we have properly addressed your concerns, could you please kindly consider increasing your initial score accordingly? Please let us know if you have further comments.

Thank you for your time and effort in reviewing our work!

Many thanks, Authors
Thank you for your response. Although I am still concerned that $H\geq N$ is not tight, the overall trend makes sense. I also like the convergence analysis of this work. I have increased my rating to 5.
Thank you for your follow-up questions.

>**regarding inconsistency with Figure 3**

The configuration of the two settings are different. While both cases have $H=1$, in Figure 3, $N=30$ and clearly here $H=1$ cannot reach loss 0.5 (because our above experiment shows that $H$ should be at least 32 to reach loss 0.5 for $N=30$). But in the above experiment, $H=1$ is for $N=1$, where the loss can reach 0.5. We also mention that the parameters $K$ and $\tau$ are also set slightly different, but they are not the main reason to cause the difference.

>**If you use a different threshold, will the bound still be tight?**

We tried several configurations and found that the nature of the results are insensitive to the thresholds (and other hyperparameters): $N$ always increases along with $H$. In the following we present the results with loss threshold = 0.1 and 0.5 and iteration number = 200, and the other configurations remain the same as the above experiment:


| N   | 1   | 5   | 10  | 20  | 30  | 40  |
|-----|-----|-----|-----|-----|-----|-----|
| H (threshold=0.5)  | 1   | 5   | 8   | 24  | 29  | 35  |
| H (threshold=0.1)  | 1   | 7   | 21   | 36  | 51  | 57  |
Thank you for the response. I am wondering why you set the threshold of the training loss as 0.5. If you use a different threshold, will the bound still be tight? By the way, it seems that the result is not consistent with Figure 3 in the submitted PDF, where the case of $H=1$ cannot reach a loss of less than 0.5. Maybe the setting is different?
Thank you very much for your quick response! Below we change $N$ and find the smallest $H$ that achieves a training loss $\leq$ 0.5 within 100 iterations. We set $K=200, d=100, m=20, \tau=0.01$. Our result is as follows:

| N   | 1   | 5   | 10  | 20  | 30  | 40  |
|-----|-----|-----|-----|-----|-----|-----|
| H   | 1   | 5   | 8   | 29  | 32  | 41  |

This result verifies our theoretical finding that with a larger $N$, $H$ should be larger to ensure a small loss. 

If our response resolves your questions, we will highly appreciate your consideration in raising the score. Certainly, we are more than happy to answer your further questions.
Thank you for the reply. I am satisfied with the other answers except the one to W1. I still think the requirement that $H\geq N$ is too strong and not impractical. It is also not verified how loose the bound $H\geq N$ is by experiments. Can you at least verify by experiments that with a larger context length $N$, the number of heads $H$ should be larger to ensure a small loss?