Title: One-Layer Transformer Provably Learns One-Nearest Neighbor In Context
Venue: NeurIPS 2024
Accepted: yes
Keywords: In-context learning theory

softmax attention

one-nearest neighbor
Abstract: Transformers have achieved great success in recent years. Interestingly, transformers have shown particularly strong in-context learning capability -- even without fine-tuning, they are still able to solve unseen tasks well purely based on task-specific prompts. In this paper, we study the capability of one-layer transformers in learning the one-nearest neighbor prediction rule. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we show that, although the loss function is nonconvex, when trained with gradient descent, a single softmax attention layer can successfully learn to behave like a one-nearest neighbor classifier. Our result gives a concrete example on how transformers can be trained to implement nonparametric machine learning algorithms, and sheds light on the role of softmax attention in transformer models.

One-Layer Transformer Provably Learns One-Nearest Neighbor In Context

Abstract

Transformers have achieved great success in recent years. Interestingly, transformers have shown particularly strong in-context learning capability - even without fine-tuning, they are still able to solve unseen tasks well purely based on taskspecific prompts. In this paper, we study the capability of one-layer transformers in learning one of the most classical nonparametric estimators, the one-nearest neighbor prediction rule. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we show that, although the loss function is nonconvex when trained with gradient descent, a single softmax attention layer can successfully learn to behave like a one-nearest neighbor classifier. Our result gives a concrete example of how transformers can be trained to implement nonparametric machine learning algorithms, and sheds light on the role of softmax attention in transformer models.

1 Introduction

Transformers have emerged as one of the most powerful machine learning models since its introduction in Vaswani et al. [2017], achieving remarkable success in various tasks, including natural language processing [Devlin et al., 2018, Achiam et al., 2023, Touvron et al., 2023], computer vision [Dosovitskiy et al., 2020, He et al., 2022, Saharia et al., 2022], reinforcement learning [Chen et al., 2021, Janner et al., 2021, Parisotto et al., 2020], and so on. One intriguing aspect of transformers is their exceptional In-Context Learning (ICL) capability [Garg et al., 2022, Min et al., 2022, Wei et al., 2023, V on Oswald et al., 2023, Xie et al., 2021, Akyürek et al., 2022]. It has been observed that transformers can effectively solve unseen tasks solely relying on task-specific prompts, without the need for fine-tuning. However, the underlying mechanisms and reasons behind the exceptional in-context learning capability of transformers remain largely unexplored, leaving a significant gap in our understanding of how and why transformers can be pretrained to exhibit such remarkable performance.

Several recent studies have attempted to understand in-context learning (ICL) through the lens of learning specific function classes. Notably, Garg et al. [2022] proposed a well-defined approach: the training data includes a demonstration prompt, consisting of a sequence of labeled data and a new unlabeled query. The in-context learning performance of a transformer is then evaluated based on its ability to successfully execute a machine-learning algorithm to predict the query data label using the prompt demonstration (i.e., the context). Based on such definition, several works such as Zhang et al. [2023], Huang et al. [2023], Chen et al. [2024] investigated ICL the optimization dynamics

of transformers under in-context learning from a theoretical lens, but their studies are limited to linear regression prediction rules, which is a significant simplification of the transformer in-context learning task. Another line of work including Bai et al. [2024], Akyürek et al. [2022] investigated the expressiveness of transformers in context, but no optimization result is guaranteed. Whether transformers can handle more complicated ICL tasks under regular gradient-based training is still, in general, unknown.

In this paper, we examine the ability of single-layer transformers to learn the one-nearest neighbor prediction rule. Our major contributions are as follows:

We establish convergence guarantees as well as prediction accuracy guarantees of a single-layer transformer in learning from examples of one-nearest neighbor classification. Utilizing the softmax attention layer, we demonstrate that the training loss can be minimized to zero despite the highly non-convex loss function landscapes. We further justify our results with numerical simulations.

Based on the optimization results, we further establish a behavior guarantee for the trained transformer, demonstrating its ability to act like a 1-NN predictor under data distribution shift. Our result thus serves as a concrete example of how transformers can learn nonparametric methods, surpassing the scope of previous literature focusing on linear regression.

In our technical analysis, we make the key observation that although the transformer loss is highly nonconvex when learning from one-nearest neighbor, its optimization process can be controlled by a two-dimensional dynamic system when choosing a proper initialization. By analyzing the behavior of such a system, we establish the convergence result despite the curse of nonconvexity.

To summarize, our result gives a concrete example of how transformers can be trained to implement nonparametric machine learning algorithms and sheds light on the role of softmax attention in transformer models. To our knowledge, this is the first paper that establishes a provable result in both optimization and consecutive behavior under distribution shift for a softmax attention layer beyond the scope of linear prediction tasks.

2 Preliminaries

In this section, we introduce the in-context learning data distribution based on the one-nearest neighbor data distribution and the setting of one-layer softmax attention transformers. Then, we discuss the training dynamics of transformers based on gradient descent.

2.1 In-Context Learning Framework: One-Nearest Neighbor

In an In-Context Learning (ICL) instance, the model is given a prompt { ( x i , y i ) } i ∈ [ N ] ∼ P prompt and a query input x N +1 ∼ P query from some data distributions P prompt and P query, where { x i } i ∈ [ N ] are the input vectors, { y i } i ∈ [ N ] ⊂ R are the corresponding labels (e.g. real-valued for regression, or { +1 , -1 } -valued for binary classification), and x N +1 is the query on which the model is required to make a prediction. Given a prompt { ( x i , y i ) } i ∈ [ N ] , the prediction task is to predict an ground truth model f ( x N +1 ; { ( x i , y i ) } i ∈ [ N ] ) that maps the query token x N +1 to a real number.

In this work, we consider using transformers as the model to perform in-context learning. For a prompt { ( x i , y i ) } i ∈ [ N ] of length N and a query token x N +1 , we consider use the following embedding:

We use the notation of h j = [ x j , y j , 0] for j ≤ N , and h N +1 = [ x N +1 , 0 , 1] . Here, { x i } i ∈ [ N ] represents the input vectors, each associated with a corresponding label { y i } i ∈ [ N ] , where y i ∈ R is the label. Throughout this paper, the sequence { ( x i , y i ) } i ∈ [ N ] are referred to as the context or prompt exchangeably. The ( d +2) -th row serves as the indicator for the training token, which equals to 0 value for i ∈ [ N ] and 1 for i = N +1 , analogous to a positional embedding vector. Such an indicator allows the model to distinguish the query token from the context. Similar models have been

SampledPointsXi

GridColorDecidedby1-NN

1.00

1.00

0.75

0.75

0.50

05'0

0.25

0.25

0.00

0.00

0.25

0.25

0.50

0.50

0.75

0.75

1.00

-1.00

1.00

1.00

1.00

0.75

0.50

0.25

0.50

0.75

1.00

0.75

0.50

0.25

0.50

0.75

0.00

0.25

0.00

0.00

0.25

0.50

0.75

1.00

1.00

0.75

0.50

0.25

0.00

0.25

0.50

0.75

1.00

1.00

0.75

0.50

0.25

studied in a line of recent works [Zhang et al., 2023, Huang et al., 2023, Chen et al., 2024, Bai et al., 2024, Akyürek et al., 2022] studying in-context learning of linear regression tasks.

Throughout this work, we focus on the case where the ground-truth prediction f ( x N +1 ; { x i , y i } i ∈ [ N ] ) of the training data is constructed based on a One-Nearest Neighbor (1NN) data distribution, defined by the following definition.

Definition 1 (One-Nearest Neighbor Predictor) . Given a prompt { ( x i , y i ) } i ∈ [ N ] and a query x N +1 , we define the one-nearest neighbor predictor by

We also define i = argmin ∥ x N +1 -x i ∥ 2

Without loss of generality, we assume that argmin j ∈ [ N ] ∥ x N +1 -x j ∥ 2 is unique. Such assumption holds almost surely whenever { x i } i ∈ [ N ] is sampled from a continuous distribution. Notably, for a fixed prompt { ( x i , y i ) } i ∈ [ N ] and query x N +1 , Definition 1 is identical to the nonparametric onenearest neighbor estimator [Peterson, 2009, Beyer et al., 1999], in which the algorithm outputs the label corresponding to the vector closest to the input, with the prompt { ( x i , y i ) } i ∈ [ N ] as the training data in 1-NN.

Next, we discuss the distribution of the training dataset { ( x i , y i ) } ∪ { x N +1 } . Throughout the training process, we focus on the case in which { x i } i ∈ [ N +1] are independently sampled from a uniform distribution on a d -1 -dimensional sphere S d -1 , { y i } i ∈ [ N ] is a zero-mean binary noise taken value in { +1 , -1 } , with { x i } i ∈ [ N +1] and { y i } i ∈ [ N ] being independent. Our data distribution assumption can be summarized formally by the following assumption:

̸

Assumption 1 (Training Distribution) . For an embedding H defined by Eq. (2.1) , we focus on the following underlying training distribution: (i) The sequence { x i } i ∈ [ N +1] are sampled independently from a uniform distribution on a d -1 dimensional sphere S d -1 ⊂ R d . (ii) The labels { y i } i ∈ [ N ] satisfies E [ y i y j | x 1: N ] = 0 and E [ y 2 i | x 1: N ] = 1 for all i = j, i, j ∈ [ N ] . (iii) We have P ( y 1: N | x 1: N ) = P ( y 1: N | -x 1: N ) .

Note that the case when { y i } i ∈ [ N ] and { x i } i ∈ [ N ] being independent when { x i } i ∈ [ N ] , with { x i } i ∈ [ N ] are uniformly sampled from the sphere and { y i } i ∈ [ N ] are randomly sampled from {± 1 } is an example of Assumption 1. We remark that by considering the training data distribution in Assumption 1, we aim to study the capability of transformers in learning one-nearest neighbor prediction rules starting from the cleanest possible setting. Despite the seemingly simple problem setting, we would like to point out that this data distribution is still challenging to study, especially because of the assumption that the second order moment of { y i } i ∈ [ N ] and { x i } i ∈ [ N +1] are uncorrelated. Due to

such uncorrelation, the classifier given by one-nearest neighbor models are rather complicated. For example, Fig. 1 illustrates the randomly generated context data and the corresponding one-nearest neighbor prediction regions with d = 3 and N = 500 , which clearly demonstrate the complexity of the training task. This further leads to a highly nonconvex and irregular objective function landscape, illustrated by Fig. 2.

2.2 One-Layer Softmax Attention Transformers

We consider a simplified version of the one-layer transformer architecture [Vaswani et al., 2017] that processes any input sequence H defined by Eq. (2.1) and outputs a scalar value:

where softmax( A ) applies softmax operator on each column of the matrix A , i.e. [softmax( A )] ij = exp( A ij ) / ∑ i exp( A ij ) . Our model is slightly different from the standard self-attention transformers, as we consider a frozen value matrix. However, we also claim that such practice is common in deep learning theory Fang et al. [2020], Lu et al. [2020], Mei et al. [2018]. We also merge the query and key matrices into one matrix denoted as W , which is often taken in recent theoretical frameworks [Zhang et al., 2023, Huang et al., 2023, Jelassi et al., 2022, Tian et al., 2023]. The output of the model is defined by the ( d +1) -th element of the last column of H W , with a closed form:

which is the weighted mean of y 1 , . . . , y N . Here and after, we may occasionally suppress dependence on { x i , y i } i ∈ [ N ] and write ̂ y W ( x N +1 ; { x i , y i } i ∈ [ N ] ) as ̂ y W ( x N +1 ) . Since the prediction takes only one entry of the token matrix output by the attention layer, actually only parts of W affect the prediction. To see this, we denote

with W 11 ∈ R d × d , W 21 ∈ R 1 × d , W 31 ∈ R 1 × d , W 12 ∈ R d × 1 , W 13 ∈ R 1 × d , W 22 , W 23 , W 32 and W 33 ∈ R . Then by Eq. 2.3, it is easy to see that W i 2 does not affect ̂ y W for i ∈ [3] , which means we can simply take all these entries as zero in the following sections. Notably, for a fixed prompt-query pair { ( x i , y i ) } i ∈ [ N ] and { x N +1 } , such an architecture allows an arbitrarily close approximation to the 1-NN model: consider W k 11 = ξ k 1 I d with ξ k 1 goes to positive infinity, W k 33 = ξ k 2 such that ξ k 2 -ξ k 1 converges to infinity, with the rest of W k ij bounded, then ̂ y W k ( x N +1 )) converges to y i ∗ as k goes to infinity.

2.3 Training Dynamics

To train the transformer model over the 1-NN task, we consider the Mean-Square Error (MSE) loss function. Specifically, the loss function is defined by

Above, the expectation is taken with respect to the sampled prompt { ( x i , y i ) } i ∈ [ N ] and the query x N +1 . Notably, when the underlying distribution for the prompt and query are defined by Assumption 1, this loss function is nonconvex with respect to W . Such nonconvexity makes the optimization hard to solve without further conditions such as PL condition or KL condition [Bierstone and Milman, 1988, Karimi et al., 2020]. We leave the proof of nonconvexity in Appendix E.

We shall consider the behavior of gradient descent on the single-layer attention architecture w.r.t. the loss function in Eq. (2.5). The parameters are updated as follows:

We shall consider the following initialization for the gradient descent:

Assumption 2 (Initialization) . Let σ > 0 be a parameter. We assume the following initialization:

Here the parameter σ is similar to masking, which is widely applied in self-attention training process, and prevents the model from focusing on the zero-label for the query x N +1 , e.g. Vaswani et al. [2017], Baade et al. [2022], Chang et al. [2022]. The reason we take the zero initialization for non-diagonal entries will be made clear when we describe the proof in Section 4. However, from a higher view, it is because we want to keep the model focusing on the inner product between different x i , which largely reduces the complexity of the dynamic system under gradient descent and makes it tractable. We leave the question of convergence under alternative random initialization schemes for future work.

3 Main Results

In this section, we summarize the convergence of training loss and testing error respectively. In Section 3.1, we discuss the convergence of training loss under gradient descent. Specifically, we prove that with a proper initialization constant σ , gradient descent is able to minimize the loss function L ( W ) despite the nonconvexity. In Section 3.2, we further discuss the testing error of the trained transformer under distribution shift. Specifically, we consider a distribution P test for the prompt { ( x i , y i ) } i ∈ [ N ] ∪{ x N +1 } , which is different from the training data distribution P prompt ⊗ P query, and discuss the difference between the trained transformer and 1-NN predictor under such distribution shift.

3.1 Convergence of Gradient Descent

First, we prove that under suitable initialization parameter σ , the loss function will converge to zero under gradient descent.

Theorem 1 (Convergence of Gradient Descent) . Consider performing gradient descent of the softmax-attention transformer model ̂ y W ( x N +1 ) . Suppose the initialization satisfies Assumption 2 with σ > 2(max { log( Nd ) , -log ( 1 -( N √ d ) 1 d ) , C d ( 1 -1 2 N ) } ) , where C d = poly( d ) , and the number of context N ≥ O ( √ d log d ) , then L ( W k ) converges to 0.

We leave the detailed proof in Appendix C. Theorem 1 shows that for the 1-NN data distribution, with a large enough initialization constant σ , the training loss of the transformer converges to zero under gradient descent. Here σ plays a role similar to the masking techniques in the self-attention training training process, in which σ is often set as infinity or an extremely large number. Such a technique has been widely accepted and shown to greatly accelerate the training process Vaswani et al. [2017], Devlin et al. [2018], Dosovitskiy et al. [2020]. We also compare our results to existing

works. Zhang et al. [2023] studied linear prediction tasks under gradient flow, however, their analysis is limited to linear attention layers. Huang et al. [2023] was the first to study softmax attention optimization under gradient descent, but their prediction is limited to linear prediction tasks under a finite orthogonal dictionary. Chen et al. [2024] established optimization convergence results for one-layer multi-head attention transformers under gradient flow. On the contrary, our work studies gradient descent convergence for transformer under a nonparametric estimator, setting it apart from all previous studies.

3.2 Results for New Task under Distribution shift

̸

In this section, we discuss the behavior of trained transformers under distribution shifts, i.e., how the model extrapolate beyond the training distribution. Following the definition in Garg et al. [2022], let us assume in the training process, the prompts { ( x i , y i ) } i ∈ ][ N ] ∼ P train prompt , and the query x N +1 ∼ P train query . During inference, the prompts and queries are sampled from a new distribution P test . We study the behavior of the trained transformers under possible prompt and query shift, i.e. P test = P train prompt ⊗ P train query . Our studies show that, under some mild conditions, the behavior of the trained model is still similar to a 1-NN predictor even under a distribution shift. Before formally stating our result, let us introduce the following assumption on the testing distribution:

Assumption 3 (Testing Distribution) . We make the following assumption on P test :

(i) There exists a R ≥ 0 such that | y i | ≤ R holds for all y i sampled from P test .

Note that Assumption 3 only requires the label y i is bounded and x i is supported on a sphere. We also remind the reader that we do not assume independence between different x i or { x i } i ∈ [ N +1] and { y i } i ∈ [ N +1] . Now we are ready to summarize our result in the following theorem.

Theorem 2 (Resemblance to 1-NN predictor under Distribution Shift) . Suppose Assumption 1 and 3 hold for P train prompt ⊗ P train query and P test . If we define

̸

̸

then, after K -iterations of gradient descent, we have

here the expectation is taken w.r.t { ( x i , y i ) } i ∈ [ N ] ∪ { x N +1 } ∼ P test . Recall that y i ∗ is the 1-NN predictor of x N +1 , which we defined in Definition 1.

We leave the detailed proof in Appendix D. Let us discuss the implication of Theorem 2. The event A δ describes the situation when the query x N +1 is located at an "inner point" away from its decision boundary, in which its distance to the nearest neighbor x i ∗ is strictly larger than all other points. Such a quantity is similar to the margin condition in classification theory in deep learning Bartlett et al. [2017] and k -NN literature Chaudhuri and Dasgupta [2014], where the optimal choice probability is strictly larger than all suboptimal choices. Specifically, if P test ( A δ ∗ ) = 1 for some δ ∗ > 0 , i.e., the query x N +1 is strictly bounded away from the decision boundary almost surely, then the L 2 distance between ̂ y W k and the 1-NN predictor will converge in a O ( R 2 K -poly( N,d ) δ ∗ ) even under a shifted distribution. We also introduce the following corollary, in which we show that when y i only takes value in a finite integer set, resembling a classification task, the trained transformer behaves like a 1-NN predictor under an additional rounding operation.

Corollary 1 (Classfication of Trained Transformer) . Suppose y i ∈ [ M ] for some integer M ≥ 0 under P test , then we have

̸

Here we define

i.e. the mapping from t ∈ R to its closest integer, and A δ is defined as in Theorem 2. Moreover, if there exists δ ∗ > 0 such that P test ( A δ ∗ ) = 0 , then we have

̸

We leave the detailed proof in Appendix D. Corollary 1 provides a convergence rate for the classification difference between 1-NN and the pretrained transformer. Notably, when x N +1 is well separated from the decision boundary in the testing distribution P test, the trained transformer will behave exactly the same as the 1-NN classifier in O ( log( MN ) poly( N,d ) δ ∗ ) gradient steps for the pretrained transformer. Theorem 2 and Corollary 1 show that the trained transformer under gradient descent is robust to both query and prompt distribution shift in the test distribution P test , in the sense that it will maintain its resemblance to a 1-NN predictor in both prediction and classification task, thus extended the results in Zhang et al. [2023], Huang et al. [2023], Chen et al. [2024] to a nonparametric estimator.

4 Sketch of Proof

In this section, we sketch the proof of Theorem 1 and highlight the techniques we used. The full proof is left to Appendix C.

Equivalence to a Two-Dimensional Dynamic System. Recall that { x i } i ∈ [ N +1] and the first and second moment of { y i } i ∈ [ N ] are uncorrelated. Utilizing this uncorrelation between { x i } i ∈ [ N +1] and { y i } i ∈ [ N ] , we can eliminate the reliance of the gradient on { y i } i ∈ [ N ] since we are considering a population loss. Moreover, utilizing the structure of the initialization, we can prove by induction that all W ij will remain zero except for W 11 and W 33 . This shows that with a suitable initialization, the transformer model will only focus on the relationship between different tokens x i throughout the whole training process. Our findings can be summarized by the following lemma.

Lemma 1 (Closed-Form Gradient) . With the initialization in Assumption 2, the gradient of L ( W k ) with respect to W 11 can be written in the following form for all k ≥ 0 :

where { g k i ( x ) } i ∈ [ N ] ∪ { g k i ∗ ( x ) } : R → R is a set of functions. Here the expectation is taken with respect to { x i } i ∈ [ N +1] , with x i ∗ = argmin x ∈{ x i } i ∈ [ N ] ∥ x -x N +1 ∥ 2 sampled i.i.d. from a uniform distribution on S d -1 . Moreover, we have ∇ W ij L ( W k ) = 0 for all ( i, j ) ∈ [3] × [3] and all k ≥ 0 except for W 11 and W 33 .

Lemma 1 shows that we only need to consider W 11 and W 33 in our update since all other entries will remain zero during the whole learning process. Note that in Eq. (4.1), all nonlinearity comes from the inner product between x ⊤ i x N +1 and x ⊤ i ∗ x N +1 . Recall that { x i } i ∈ [ N +1] are i.i.d. sampled from a uniform distribution supported on a d -1 -dimensional sphere S d -1 , therefore, the distribution of { x i } i ∈ [ N +1] is rotational invariance, which means ⊗ i ∈ [ N +1] P x i = ⊗ i ∈ [ N +1] P U x i for all orthogonal matrix U ∈ R d × d . Since the rotation of { x i } i ∈ [ N +1] does not change the inner products { x ⊤ i x i } i ∈ [ N ] and x ⊤ i ∗ x N +1 , from the structure of ∇ W 11 L ( W k ) illustrated by Eq. (4.1), we shall always have U ∇ W 11 L ( W k ) U ⊤ = ∇ W 11 L ( W k ) , which shows ∇ W 11 L ( W k ) = c k I d for some constant c k by simple algebra. We summarize our result in the following lemma.

Lemma 2 (Two-Dimensional System) . With the initialization in Assumption 2, there exists two sets of real numbers { ξ k 1 } k ≥ 0 and { ξ k 2 } k ≥ 0 , such that W k has the following form:

With Lemma 2, we reduce the dimension of the original dynamic system in Eq. (2.6) from ( d +2) 2 to 2 . We now only need to focus on the evolution of ξ k 1 and ξ k 2 for all k ≥ 0 .

Convergence of the Dynamic System. Lemma 2 helps us largely reduce the dimension of the training dynamics. However, this does not make our question a trivial one, as the loss function is still highly nonconvex even when we only need to consider a two-dimensional subspace of R ( d +2) × ( d +2) . To see this, we introduce the following lemma:

Lemma 3 (Nonconvexity of Transformer Optimization) . When W lie in a two-dimensional subspace of R ( d +2) × ( d +2) defined by W = diag { ξ 1 , . . . , ξ 1 ︸ ︷︷ ︸ d times , 0 , -ξ 2 } , the original loss function defined in

Eq. (2.5) is equivalent to the following:

Such loss function is still nonconvex.

We leave the detailed proof in Appendix E. Nonconexity shown by Lemma 3 implies attaining the global minimum could be hard. Previous works such as Zhang et al. [2023] utilize conditions such as Polyak-Lojasiewicz inequality to analyze such systems, however, those conditions are not applicable in our setting, and a more delicate analysis for the evolution of ξ k 1 and ξ k 2 is needed. We characterize their behavior by the following lemma.

Lemma 4. For ξ k 1 ≥ 0 , there exists constants c 1 , c 2 , c 3 , c 4 > 0 , such that

and

Lemma 4 shows that there exits a constant c b ∈ (0 , 1) , such that ξ k 1 will keep increasing with a scale of Ω( η log k ) until ξ k 1 ≤ c b ξ k 2 . With this ratio, we obtain the following lemma for the increment of ξ k 2 .

Lemma 5. For ξ k 1 ≥ 0 , there exits constant c ′ 1 , c ′ 2 , such that

Lemma 5 shows that ξ k 2 will monotonically increase with a scale of Ω( η exp( ξ k 2 )) , which implies ξ k 2 = Ω(log k ) . Combining Lemma 4 and 5, we show that both ξ k 1 and ξ k 2 converge to infinity, with ξ k 1 maintaining a slower speed, as its decreases when getting closer to ξ k 2 from the below. Recall that the loss function is equivalent to Eq. (4.2) under the initialization specified in Assumption 2, which shows that L ( ξ k 1 , ξ k 2 ) will converge to zero as long as ξ k 1 and ξ k 2 -ξ k 1 both converges to infinity. We thus conclude our proof of L ( W k ) eventually converges to it global minimum.

5 Numerical Results

In previous sections, we have shown that with the initialization specified in Assumption 2, a single softmax attention layer transformer is able to learn the 1-NN predictor under gradient descent and remain robust under distribution shift. We now conduct experiments in a less restrictive setting and show that even without specific initialization and full-batch gradient descent, simple stochastic gradient descent updates with random parameter initialization for the parameters are still sufficient for the model to learn the 1-NN predictor. First, we investigate the convergence of single-head single-layer transformers [Vaswani et al., 2017] trained on 1-NN tasks. The training data are sampled from P train prompt and P train query , defined in Assumption 1. We choose context length N ∈ { 16 , 32 , 64 } and input dimension d ∈ { 8 , 16 } . The model is trained on a dataset with a size of 10000, and an epoch

0

250

500

750

1000

1250

1500

1750

2000

Gradient Iterations

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Training Loss

Training Loss on Uniform Sphere Distribution

0

250

500

750

1000

1250

1500

1750

2000

Gradient Iterations

0.0

0.2

0.4

0.6

0.8

Testing MSE

Testing MSE under Distribution Shift

d = 8, N =16

d = 8, N =32

d = 8, N =64

d = 16, N =32

number of 2000. To ensure our training convergence result holds beyond the gradient descent scheme, we choose SGD as our optimizer, with a batch size of 128 and a learning rate of 0.1. We use the random Gaussian as our initialization. Our results for the training loss convergence are summarized in the left panel of Fig. 3. The results show that the model converges to 1-NN predictor on the training data even under SGD and random initialization. Moreover, as the dimension d and the length of contexts N becomes larger, the convergence speed becomes slower.

To verify our results on the distribution shift, we generate testing data sampled from a distribution difference from the training data, and report the mean square error between the model prediction and the 1-NN predictor. Furthermore, the testing data satisfies P ( A ∗ δ ) = 1 , with δ ∗ specified as 0 . 1 . Recall that we defined A δ in Theorem 2 as the event where x N +1 is separated from the decision boundary with a distance of at least δ . We leave the details of our data-generating process in Appendix B. We test the trained transformer model on this dataset once every epoch throughout the training process. Our results are summarized in the right panel of Fig. 3. The results show that the testing error decreases much faster than the training loss, due to the boundary separation condition, which the uniformly-sampled training data do not enjoy. Our result also coincides with our theoretical result in Theorem 2, showing that the trained transformers are robust under distribution shift, and benefits greatly from staying away from the decision boundary.

6 Conclusion

We investigate the ability of single-layer transformers to learn the one-nearest neighbor prediction rule, a classic nonparametric estimator. Under a theoretical framework where the prompt contains a sequence of labeled training data and unlabeled test data, we demonstrate that, despite the nonconvexity of the loss function during gradient descent training, a single softmax attention layer can successfully emulate a one-nearest neighbor predictor. We further show that the trained transformer is robust to the distribution shift of the testing data. As far as we know, this paper is the first to establish training convergence and behavior under distribution shifts for softmax attention transformers beyond the domain of linear predictors.

Acknowledgments and Disclosure of Funding

We thank the anonymous reviewers for their helpful comments. Yuan Cao is partially supported by NSFC 12301657 and Hong Kong RGC-ECS 27308624. Han Liu's research is partially supported by the NIH R01LM01372201. Jason M. Klusowski was supported in part by the National Science

Foundation through CAREER DMS-2239448, DMS-2054808 and HDR TRIPODS CCF-1934924. Jianqing Fan's research was partially supported by NSF grants DMS-2210833, DMS-2053832, and ONR grant N00014-22-1-2340. Mengdi Wang acknowledges the support by NSF IIS-2107304, NSF CPS-2312093, ONR 1006977 and Genmab.

References

Jacob Abernethy, Alekh Agarwal, Teodor Vanislavov Marinov, and Manfred K Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. In International Conference on Algorithmic Learning Theory , pages 3-46. PMLR, 2024.

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems , 36, 2024.

Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661 , 2022.

Alan Baade, Puyuan Peng, and David Harwath. Mae-ast: Masked autoencoding audio spectrogram transformer. arXiv preprint arXiv:2203.16691 , 2022.

Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems , 36, 2024.

Peter L Bartlett, Dylan J Foster, and Matus J Telgarsky. Spectrally-normalized margin bounds for neural networks. Advances in neural information processing systems , 30, 2017.

Kevin Beyer, Jonathan Goldstein, Raghu Ramakrishnan, and Uri Shaft. When is 'nearest neighbor' meaningful? In Database Theory-ICDT'99: 7th International Conference Jerusalem, Israel, January 10-12, 1999 Proceedings 7 , pages 217-235. Springer, 1999.

Edward Bierstone and Pierre D Milman. Semianalytic and subanalytic sets. Publications Mathématiques de l'IHÉS , 67:5-42, 1988.

Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11315-11325, 2022.

Kamalika Chaudhuri and Sanjoy Dasgupta. Rates of convergence for nearest neighbor classification. Advances in Neural Information Processing Systems , 27, 2014.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442 , 2024.

Herbert A David and Haikady N Nagaraja. Order statistics . John Wiley & Sons, 2004.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Cong Fang, Jason D. Lee, Pengkun Yang, and Tong Zhang. Modeling from features: a mean-field framework for over-parameterized deep neural networks, 2020.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D Lee. How well can transformers emulate in-context newton's method? arXiv preprint arXiv:2403.03183 , 2024.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 16000-16009, 2022.

Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249 , 2023.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:1273-1286, 2021.

Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems , 35:37822-37836, 2022.

Hong Jun Jeon, Jason D. Lee, Qi Lei, and Benjamin Van Roy. An information-theoretic analysis of in-context learning, 2024.

Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximalgradient methods under the polyak-łojasiewicz condition, 2020.

Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning , pages 19565-19594. PMLR, 2023a.

Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning , pages 19689-19729. PMLR, 2023b.

Yiping Lu, Chao Ma, Yulong Lu, Jianfeng Lu, and Lexing Ying. A mean-field analysis of deep resnet and beyond: Towards provable optimization via overparameterization from depth, 2020.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33), July 2018. ISSN 1091-6490. doi: 10.1073/pnas.1806579115. URL http://dx.doi.org/10.1073/pnas. 1806579115 .

Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work?, 2022.

Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference on machine learning , pages 7487-7498. PMLR, 2020.

Leif E Peterson. K-nearest neighbor. Scholarpedia , 4(2):1883, 2009.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding, 2022.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. Advances in Neural Information Processing Systems , 36:71911-71947, 2023.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning , pages 35151-35174. PMLR, 2023.

Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, et al. Larger language models do in-context learning differently. arXiv preprint arXiv:2303.03846 , 2023.

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 , 2021.

Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927 , 2023.

A Related Works

In-Context Learning. In-context learning refers to transformers' ability to solve unseen tasks without fine-tuning. Min et al. [2022] studied which aspects of the demonstrations contribute to end-task performance. Further, Garg et al. [2022] empirically investigated the ability of transformer architectures to learn a variety of function classes in context. From a theoretical perspective, Bai et al. [2024], Abernethy et al. [2024], Akyürek et al. [2022] studied the expressiveness of transformers to approximate statistical algorithms. Li et al. [2023a] studied the generation ability of transformers in ICL tasks. Jeon et al. [2024] studies the information-theoretical lower bound of in-context learning. However, none of these works studied the optimization process when training a transformer in context.

Optimization of Transformers. There are various works that studied optimization for transformers. Among all these studies, Huang et al. [2023], Zhang et al. [2023], Chen et al. [2024] studied the optimization dynamics of transformers of learning linear prediction tasks in context. Specifically, Zhang et al., 2023 studied the gradient flow dynamics of linear self-attention model and obtained convergence results utilizing PL condition. Huang et al. [2023] studied linear prediction tasks with a softmax-attention layer trained with gradient descent, with a finite support prompt/query distribution. Chen et al. [2024] studied the gradient flow dynamics in training multi-head softmax attention on linear prediction tasks. Ahn et al. [2024], Giannou et al. [2024] studied transformers' ability to learn optimization methods. Prior works also studied transformer optimization beyond in-context learning tasks. Tian et al. [2023] studied a single-layer transformer with one self-attention layer plus one decoder layer, and proved that the attention layer acts as a scanning algorithm. Li et al. [2023b] studied the optimization of transformers in learning semantic structures. Jelassi et al. [2022] studies the spatial localization property of vision transformer in optimization.

Notations. We write [ N ] = { 1 , . . . , N } . For two vector x 1 = ( x 1 1 , . . . , x 1 d ) and x 2 = ( x 2 1 , . . . , x 2 d ) , we write their inner product ∑ d i =1 x 1 i x 2 i as x 1 ⊤ x 2 . We use 0 n and 0 m × n to denote the zero vector and zero matrix of size n and m × n , respectively. We write { x ∈ R d : ∥ x ∥ 2 = 1 } as S d -1 . For two series { a k } k ≥ 0 and { b k } k ≥ 0 , we write a k = Ω( b k ) if there exists 0 < C 1 , C 2 such that C 1 · b k ≤ a k ≤ C 2 · b k . We write a k = O ( b k ) if there exists C > 0 such that a k ≤ C · b k . We use a k = poly( b k ) if there exits an n -degree polynomial P n ( x ) such that a k = O ( P n ( b k )) . For ( a k ) k ∈ [ N ] , we define its permutation ( a ( k ) ) k ∈ [ N ] such that a (1) ≥ a (2) ≥ . . . ≥ a ( n ) . We use I d to denote the d -dimensional identity matrix and sometimes we also use I when the dimension is clear from the context. Unless otherwise defined, we use lower case letters for scalars and vectors and use upper case letters for matrices. For a matrix A we denote its entry in the i -th row and j -th column by [ A ] ij .

B Data-Generating Process in Experiment

In our experiment, we generate the testing dataset with context length N and dimension d with separation parameter δ ∗ such that ∥ x j -x N +1 ∥ 2 2 ≥ ∥ x i ∗ -x N +1 ∥ 2 2 + δ for all j ∈ [ N ] , j = i ∗ by the following procedure:

̸

(i) We sample x i from the uniform distribution on S d -1 and y i ∼ N (0 , 1) for all i ∈ [ N ] ;

(ii) We random sample i ∗ ∈ [ N ] by uniform distribution and set x N +1 = x i ∗ with the 1-NN label being y i ∗ ;

(iii) If ∥ x j -x N +1 ∥ 2 2 ≤ δ , set x j = -x j .

C Proof for Theorem 1

In this section, we elaborate on the proof of Theorem 1. Our proof can be broken down into the following steps:

(i) With induction, we prove that the evolution dynamics of W k under gradient descent can be captured by a two-parameter dynamic system, parametrized by ξ k 1 and ξ k 2 ;

(ii) By estimating the update dynamics for ξ k 1 and ξ k 2 , we prove that with a proper initialization parameter σ , We will have ξ k 1 , ξ k 2 = Ω(log k ) , with ξ k 1 ≤ c · ξ k 2 for a constant c ∈ (0 , 1) .

(iii) With the non-asymptotic behavior of ξ 1 and ξ 2 determined, we further control the loss function L ( W k ) and establish a convergence for the loss function.

In the following sections, we will discuss how we prove those three items in Section C.1, C.2 and C.3, respectively.

C.1 Dynamic of Gradient Descent

In this section, we prove that we can characterize the evolution of W under gradient descent with a two-parameter dynamic system. We proceed by mathematical induction:

We prove that when W k is a diagonal matrix with [ W k ] [ d ] × [ d ] = c k I d for some constant c k , we can have the following break downs of our proof:

where ∆ ξ 0 1 and ∆ ξ 0 2 are two positive constants.

By W k being a diagonal matrix with [ W k ] [ d ] × [ d ] = c k I d , combined with Eq. (C.1), we prove that W k +1 is a diagonal matrix and [ W k +1 ] [ d ] × [ d ] = c k +1 I d . .

Now since (1) is naturally satisfied by the initialization in Assumption 2 for k = 0 , we can conclude the proof by simply proving (1) for k ≥ 1 .

Proof for Step (1) of Induction. Recall that our loss function can be written as

where

with W 11 ∈ R d × d , W 21 ∈ R 1 × d , W 31 ∈ R 1 × d , W 12 ∈ R d × 1 , W 13 ∈ R 1 × d , W 22 , W 23 , W 32 and W 33 ∈ R , and we make the additional definition of y † j = y j , W † j = W 23 , W ‡ j = W 21 for j ∈ [ N ] and y † N +1 = 1 , W † N +1 = W 33 , W ‡ N +1 = W 31 . Throughout the rest of this paper, we will also adopt the notation of

when there is no ambiguity, where q : W × { x i } i ∈ [ N +1] × { y i } i ∈ [ N ] . As we have discussed in Section 2.2, W ∗ 2 does not affect the outcome of the transformer, therefore the second column will remain 0 throughout the training procedure. Now we will calculate the gradient ∇ W ij L ( W 0 ) respectively. Note that for f j θ ( x ) = exp( g j θ ( x )) , we have

therefore, with some algebra, we have the following closed-form formula for ∇ W ij L ( W ) respectively,

and

By the induction assumption, we have

̸

̸

where the first equality comes from ∥ x N +1 ∥ 2 = 1 . Therefore, under our induction q j is only a function of { x i } i ∈ [ N +1] and W k , independent to { y i } i ∈ [ N ] . Now, with the closed form of ∇ W ij L ( W ) , we can make the following calculation. First, we calculate ∇ W 1 1 L ( W k ) :

since q ( x , W k ) is only a function of x by our induction assumption, we have

here

also for (2), we have

(1)+(2) gives us

For ∇ W 21 L ( W k ) , we have

note that it is the expectation of multiplication of three linear functions of y i , therefore ∇ W 21 L ( W k ) equals to 0 by symmetry. For ∇ W 31 L ( W k ) , we have

here the second equality comes from q j ( x , W k ) being an even function of x and x has a symmetric distribution. For ∇ W 13 L ( W k ) , we have

here the second equality comes from q j ( x, W k ) is a function of ∥ x ∥ 2 and x has a symmetric distribution. For ∇ W 23 L ( W k ) , we have

such that

here the expectation is taken with respect to { x i } i ∈ [ N +1] . By Lemma 6, we have P x i ⊗ P x j = P U x i ⊗ P U x j for any orthogonal matrix U ∈ O ( d ) , therefore

due to its symmetry in y . For ∇ W 33 L ( W k ) , we have the following calculation:

With all previous calculations, we know that ∇ W ij L ( W k ) equals 0 except for ∇ W 11 L ( W k ) and ∇ W 33 L ( W k ) with our induction assumption.

Now we only need to prove that ∇ W 11 L ( W ) is a diagonal matrix. Recall that { x i } i ∈ [ N +1] is identically uniformly distributed on S d -1 , thus it naturally satisfies the following Assumption:

Assumption 4 (Rotational Invariance) . We say the distribution of x is rotationally invariance, if x ∼ P x , and for every U ∈ O ( d ) , where O is the group of orthogonal matrices, we have P x = U # P x .

A straightforward lemma under Assumption 4 is the following.

Lemma 6. Suppose { x i } i ∈ [ N +1] are sampled i.i.d. from P x and Assumption 4 holds, then we have ⊗ N +1 i =1 P x i = ⊗ N +1 i =1 P U x i for all U ∈ O ( d ) .

Proof. Since x i ∗ are sampled i.i.d. from P x , we have

Lemma 7. When x i ∼ P x , P x satisfies Assumption 4, and W k is a diagonal matrix with [ W k ] [ d ] × [ d ] = c k I d , we have E [ ∇ W 11 L ( W k )] = a k · I d .

Proof. Note that by the induction assumption, we have

̸

therefore, q j ( x , W k ) = q j ( U x , W k ) for all U ∈ O ( d ) . By Eq. (C.10), there exits a set of functions

here the first inequality comes from P x i ⊗ P x j = P U x i ⊗ P U x j . Such identity holds for all orthogonal matrix U , therefore E [ g k i ( x ⊤ i x N +1 ) · x i x N +1 ] must be a multiplication of the identity matrix I d . Now, note that for all orthogonal matrix U ,

therefore

Again, since U is an arbitrary orthogonal matrix, we conclude that E [ g k i ∗ ( x ⊤ i ∗ x N +1 ) · x i ∗ x ⊤ N +1 ] is a multiplication of the identity matrix. Summing everything together, we conclude the proof that ∇ W k 11 L ( W k ) is a multiplication of an identity matrix.

With Lemma 7, Eq. (2.6), and the previous calculations, we conclude our induction.

C.2 Evolution of ξ 1 and ξ 2

In this section, we characterize the training dynamics of ξ 1 and ξ 2 under gradient descent. Our results can be broken down into the following steps:

We prove that when ξ 0 1 and ξ 0 2 are initialized as in Assumption 2 , with σ , N , d satisfy the condition in Theorem 1, we have ξ 1 1 -ξ 0 1 ≥ 0 .

We provide an upper bound and lower bound for ξ k +1 1 -ξ k 1 and ξ k +1 2 -ξ k 2 .

We prove that there exits constant c 1 ∈ (0 , 1) such that ξ k 1 ∈ (0 , c 1 ξ k 2 ) .

We prove that ξ k +1 2 -ξ k 2 = Ω( η · exp( -poly( N,d ) · ξ k 2 )) , we thus conclude that ξ k 2 = Ω( η · log(poly( N,d )) · log( k )) .

Combining (3) and (4), we conclude this section by proving that ξ k 1 , ξ k 2 and ξ k 2 -ξ k 1 are both of scale Ω(log k ) .

Step (1): Initial incrementation of ξ 0 1 . First, we prove that ξ 1 1 -ξ 0 1 ≥ 0 . Note that by Lemma 7 and Eq. (C.10), we have

To prove that ξ 0 1 -ξ 1 1 < 0 , we only need

However, by Lemma 18, this can be guaranteed by N ≥ O ( √ d log d ) .

Step (2): Scale of ξ k +1 1 -ξ k 1 and ξ k +1 2 -ξ k 2 Note that in every gradient update with stepsize η , we have

we also have the following estimation for the update of ξ 3 :

which shows that in each iteration, ξ 3 will at least increment by a scale of η · E [ q i ∗ ( x , W k ) · q 2 N +1 ( x , W k )] . Combining Eq. (C.12), (C.11), we have the following estimation:

Lemma 8. We have

for all k ≥ 0 .

Proof. First, note that we have the following relations:

we have

Therefore we have

here C d is a constant that only pertains to d . The first inequality comes from Eq. (C.16), and the third inequality comes from Lemma 17. Finally, note that

and we conclude the proof.

Next, we provide the following upper bound for the increment of ξ k 1 and ξ k 2 whenever ξ k 1 ≥ 0 .

here the inequality comes from x ⊤ i ∗ x N +1 ≥ x ⊤ j x N +1 for all j ∈ [ N ] . Moreover, note that when ξ 1 ≥ 0 , we have

which is equivalent to q i ∗ ( x , W ) ≥ 1 -q N +1 ( x , W ) N . Since

Lemma 9. When ξ k 1 ≥ 0 and N ≥ O ( √ d log d ) as in Theorem 1, we have the following inequalities:

where a n,d = ( 2 N √ d ) -2 d -3 , and

Proof. We first prove the first inequality. Recall that

when ξ k 1 ≥ 0 , we have q i ∗ ( x , W ) -∑ N j =1 q 2 j ( x , W ) ≥ 0 , therefore

We have the following bound for (iii) when ξ k 1 ≥ 0 :

where the first inequality comes from Lemma 19. Now we aim to bound (i) and (ii) separately. First, we have

here the first inequality comes from Jensen's inequality and the convexity of x exp( -x ) between [0 , 2] , and the second one comes from E [ x ⊤ i ∗ x N +1 ] ≥ 2 ( N +1) 2 , which is guaranteed by the condition of N in Theorem 1 and Lemma 18. Thus we conclude the proof for the first inequality. For (ii), denote the second largest order statistics in { x ⊤ i x N +1 } i ∈ [ N ] by x ⊤ (2) x N +1 , we have the following inequalities:

where we define a random variable ζ k by

the second inequality comes from

and the third inequality comes from Jensen's inequality. Combining We now aim to prove the second inequality. Recall that by Eq. (C.11), we have

therefore we have

Thus we conclude the proof of both inequalities.

The next lemma is a direct combination of Lemma 8 and Lemma 9.

Lemma 10. When ξ k 1 ≥ 0 , we have

The next lemma provides a lower bound for the update of ξ 3 . Notably, when ξ k 1 is close to ξ k 2 , ξ k 2 increases in a larger scale.

Lemma 11. When ξ k 1 ≥ 0 , we have the following inequality holds:

where a n,d is a constant only pertains to n and d that is defined in Lemma 9

Proof. With Eq. (C.11), we have

where the first inequality comes from Eq. (C.14), second inequality comes from Lemma 19, the third inequality comes from ξ k 1 ≥ 0 .

Step (3): Upper Bound for ξ k 1 /ξ k 2 . Now, combining Lemma 9 and 11, we immediately get the following result:

Lemma 12. If σ = ξ 0 2 ≥ 3 log( a n,d 2 N 4 d ) = O ( max { log( Nd ) , -log ( 1 -( N √ d ) 1 d ) } ) , and ξ k 1 ≥ 0 for all k ≥ 0 , then we will have ξ k 1 ≤ 7 15 ξ k 2 .

Proof. By Lemma 9, we know that whenever ξ k 1 ≥ 0 , we have

therefore, if

which is equivalent to

we will have ξ k +1 1 -ξ k 1 < 0 . By Lemma 11, we have ξ k 2 ≥ ξ 0 2 ≥ 3 log( a n,d 2 N 4 d ) . Therefore ξ k +1 1 will not increase as long as

finally, recall that and our result follows.

Step (4): Scale of ξ k 1 and ξ k 2 . Finally, we conclude the convergence of gradient descent with the following two Lemma:

Lemma 13. With σ and N satisfying the condition in Theorem 1, we have

with ξ k 1 ≤ 1 2 ξ k 2 holds for all k ≥ 0 .

Proof. We first establish a convergence rate for ξ 2 . Note that by Lemma 9, Lemma 11 and Lemma 12, we have and

holds for all k ≥ 0 . Therefore, we have ξ k 2 = Ω( η · poly( N,d ) · log k ) for all k ≥ 0 . Now we turn to ξ k 1 . By Lemma 11 , when we have

Since ξ 0 2 = σ ≥ 2 log ( C d ( 1 -1 2 N ) 4 ) , ξ k 2 monotonically increasing, and ξ 1 1 ≥ 0 , by induction, we have ξ k 1 ≥ 0 for all k , and ξ k 1 ≥ O (poly( N,d ) log k ) until ξ k 1 ≥ ξ k 2 8 . Meanwhile, Lemma 12 shows that

for all k ≥ 0 , which implies ξ k 1 = O (poly( N,d ) log k ) . Therefore, ξ k 1 = Ω(poly( N,d ) log k ) . The results of ξ k 1 ≤ 1 2 ξ k 2 follows by Lemma 12 and ξ k 1 ≥ 0 .

C.3 Convergece of Loss Function L ( W )

We also have the following bound for the loss function.

Lemma 14. When W defined in Eq. (2.4) satisfies W 11 = ξ 1 I d , W 33 = -ξ 2 , with ξ 1 ≥ 0 , and the rest of items are all zero matrices, we have

Proof. We have

note that

̸

̸

≤ 0 , Therefore E [( ∑ N +1 j =1 q ( x , W ) y j -y i ∗ ) 2 ] ≤ E [1 -q i ∗ ( x , W )] . However, we have E [ 1 -q i ∗ ( x , W ) ] ≤ E [ ∑ j = i ∗ ,j ∈ [ N +1] q j ( x , W ) ] = E [ q i ∗ ( x , W ) ∑ j = i ∗ ,j ∈ [ N +1] q j ( x , W ) q i ∗ ( x , W ) ] ≤ ( N -1) E [ exp ( ξ 1 · { x ⊤ j x N +1 -x ⊤ i ∗ x N +1 }) ] + E [exp( ξ 1 (1 -x ⊤ i ∗ x N +1 ) -ξ 3 )] , ( q i ∗ ≤ 1 ) ≤ ( N -1) E [ exp ( ξ 1 · { x ⊤ j x N +1 -x ⊤ i ∗ x N +1 }) ] +exp(2 ξ 1 -ξ 3 ) , By Lemma 20, we have E [( N +1 ∑ j =1 q ( x , W ) y j -y i ∗ ) 2 ] ≤ O ( N 3 k 2 d ξ 1 ) +exp(2 ξ 1 -ξ 3 ) , and we conclude the proof.

Now, combine Lemma 15 with our dynamic bounds for ξ k 1 and ξ k 2 developed in Section C.2, we have the following convergence result for the loss function.

Lemma 15. When W K defined in Eq. (2.4) is updated by gradient descent in Eq. (2.6) , we have

Proof. By Lemma 15, we have

By Lemma 13, we have ξ k 1 ≤ 7 15 ξ k 2 , with ξ k 1 = Ω(poly( N,d ) log k ) . Thus we have

D Proof for Theorem 2 and Corollary 1

In this section, we discuss the behavior of the pretrained transformer in tasks under distribution shift, and provide proof for Theorem 2 and Corollary 1 in Section 3.2.

D.1 Proof for Theorem 2

In this section, we provide a proof for Theorem 2. The result comes from the following observation. First, we condition our analysis on the A δ , i.e. assuming that there exists a constant δ such that

̸

Then we have

̸

uniformly for all { ( x i , y i ) } i ∈ [ N ] and x N +1 on A δ . Here the second inequality comes from Theorem 1, in which we show that ξ K 1 ≤ 7 15 ξ K 2 . Recall that we defined q j ( x , W k ) in Eq. (C.2) as the weight for the j -th label in { y j } j ∈ [ N +1] . Next, we have

where the inequality comes from Eq. (D.2). Here the expectation is taken over the testing distribution P test . By taking inferior on the inequality above for all δ > 0 , we conclude our proof for Theorem 2.

D.2 Proof for Corollary 1

In this section, we provide the proof for Corollary 1. The first statement of Corollary 1 comes from Markov's inequality:

̸

where the last inequality comes from Theorem 2. Next, we prove the second statement. Suppose P test ( A δ ∗ ) = 1 for some δ ∗ > 0 . Then similar to the argument in Eq. (D.2), we have

̸

̸

̸

holds for all { ( x i , y i ) } ∪ { x N +1 } ∼ P test almost surely. Note that

holds whenever therefore, it suffices to have 1 2 ≤ O ( MNK -poly( N,d ) δ ∗ ) , which is equivalent to

E Nonconvexity of Loss Function

In this section, we show that the loss function is defined by Eq. (2.5). We prove by a special subspace of W ∈ R ( d +2) × ( d +2) defined by two scalars ξ 1 , ξ 2 ,

By showing L ( W ) is nonconvex under such parametrization, we conclude our proof.

Lemma 16 (Nonconvexity of Transformer Optimization) . When W is a two-dimensional subspace of R ( d +2) × ( d +2) defined by Eq. (E.1) , the original loss function defined in Eq. (2.5) degenerates to the following:

where we use y N +1 = 0 . Such loss function is still, in general, nonconvex.

Proof. The degeneracy of the original Eq. (2.5) to L ( ξ 1 , ξ 2 ) can be shown by basic algebra. We only need to show the nonconvexity of L ( ξ 1 , ξ 2 ) in our proof. Note that by Assumption 1, the gradient of L ( ξ 1 , ξ 2 ) is defined by

where the third equation comes from

and E [( ∑ N j =1 y i ) 2 ] = E [ ∑ N j =1 y 2 i ] = N . Eq. (E.2) shows that when lim ξ 2 → + ∞ ∂ ξ 2 L (0 , ξ 2 ) = 0 and lim ξ 2 →-∞ ∂ ξ 2 L (0 , ξ 2 ) = 0 . If L ( ξ 1 , ξ 2 ) is convex, then ∂ ξ 2 L (0 , ξ 2 ) is monotonically increasing, which means ∂ ξ 2 L (0 , ξ 2 ) = 0 for all ξ 2 . However, this is clearly a contradiction.

F Auxiliary Lemma

Lemma 17 (Distribution of Sphere Inner Product) . With the assumption of x i sampled from a uniform distribution on the sphere, Let τ be the cosine of the angle between an arbitrary d -dimensional vector and a vector chosen uniformly at random from the unit sphere. Then the probability density function of random variable τ ∈ [ -1 , 1] is f τ ( t ) = 2Γ( d 2 ) √ π Γ( d -1 2 ) · (1 -t 2 ) d -3 2 .

Proof. For convenience, we assume that x N +1 = e 1 . Note that this does not change the distribution of x N +1 · x i due to rotation invariance. Let X i ∼ N (0 , 1) , define

Note that Y is distributed the same way as x i . Calculating the Jacobian:

Since J is of the form:

the determinant is easily evaluated:

Now to solve for the distribution of Y N .

where

Writing the distribution in terms of t :

Now we can make a substitution to remove y d from inside the integral. Let u n = ( 1 -t 2 ) -1 / 2 y n , we have

and

The integral can be seen to be a constant so:

or for notational convenience:

Since f τ ( y ) is a PDF and is defined for 0 ≤ t ≤ 1 :

Furthermore, for d > 1 , we have k d = 2Γ ( d 2 ) √ π Γ ( d -1 2 ) .

Lemma 18. With N = O ( log d · √ d ) , we have E [ x ⊤ i ∗ x N +1 ] ≥ 2 ( N +1) 2 .

Proof. Note that we have

To ensure that E [ x ⊤ i ∗ x N +1 ] ≥ 2 ( N +1) 2 , we only need

which is equivalent to

Therefore, we only need

which suffices by suffices.

Lemma 20. Suppose { x i } i ∈ [ N +1] are i.i.d. samples from a uniform distribution on a sphere in R d , with x ⊤ i ∗ x N +1 and x ⊤ (2) x N +1 being the largest and second largest order statistics among { x ⊤ i x N +1 } i ∈ [ N ] , respectively. Then we have

where Ω( · ) hides constant depends on N and d .

Moreover, we have

Note that

therefore we only need

Now, let α = 1 d -2 , since k d = Γ( d 2 ) Γ( d -1 2 ) = O ( √ d ) , we only need N ≥ O ( d e √ d log( d -2)) = O ( √ d log d ) .

Lemma 19 (Concentration upper bound for x ⊤ i ∗ x N +1 ) . When x ⊤ i ∗ x N +1 is defined by max i ∈ [ N ] { x ⊤ i x N +1 } , where x i are independently sampled from a uniform sphere distribution, we have

Note that

since (1 -1 N ) N is monotonically increasing, we only need k d (1 -α )(1 -α ) d -3 2 ≤ 1 N . Setting

Proof. We denote x ⊤ i x N +1 by Y i , x ⊤ i ∗ x N +1 by Y i ∗ and x ⊤ (2) x N +1 by Y (2) . By Lemma 17, we have the density function of Y i being f d ( t ) = k d · (1 -t 2 ) d -3 2 , where k d = 2Γ( d 2 ) √ π Γ( d -1 2 ) . Then by the joint distribution of order statistics [David and Nagaraja, 2004], the joint density function of Y i ∗ and Y (2) is

Thus we obtain the upper bound. Next, we establish a lower bound.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our main claims are made clear in the abstract and introduction.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Our paper is limited to the theoretical analysis of single-layer transformers under 1-NN contexts.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All assumptions and proofs are included in the main paper and the appendix. Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All technical details are provided in the paper and the appendix

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in

some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We only use simulated data, and provided enough technical details for the data and code we used in the main paper and appendix.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. 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: We specify all hyperparameters and optimizers in the main paper and appendix.

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.

7. 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: [Yes]

Justification: We include error bar obtained from 10 independent trials

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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All experiments are conducted on a CPU cluster

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our work conforms with the code of ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our work discusses the theoretical performance of a well-known architecture, thus the social impacts are insignificant.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

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).

11. 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: Our work does not 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: Our 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Our paper does not introduce any new asset.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the gradient descent dynamics of a softmax-activated self-attention unit trained on a population loss over in-context learning (ICL) tasks. Each task entails predicting the binary label of a query input, where the true label is the label of the 1-nearest neighbor (1-NN) in the context, and the context and query come from a particular distribution. The main result is that starting from a particular initialization, the key-times-query weight matrix converges to an infinite-valued matrix that exactly implements a 1-NN predictor. An additional result bounds the error of this predictor when the context and query come from any distribution satisfying a weaker assumption than the training distribution, and the query label is again the 1-NN. Empirical results verify these theoretical results in the analyzed setting, with slight modifications.\n\nSTRENGTHS:\n- The results develop understanding of three key areas that are understudied in the in-context learning (ICL) theory literature: (1) the behavior of *softmax*-activated attention (rather than linear attention), (2) ICL of tasks other than linear regression, and (3) gradient-based optimization dynamics.\n\n- Aside from typos, the results are rigorous, well-formulated and non-trivial.\n\n- Regarding non-triviality: Lemma 3 is especially helpful to show that the population loss is still nonconvex even when reduced to the loss over the two scalars.\n\n- The proof sketch is mostly well-written (see below).\n\n- The experiments are well-explained and suggest that some of the simplifications made in the analysis (specific initialization, full-gradient descent on population loss) do not fundamentally change the results.\n\nWEAKNESSES:\n1. The required conditions on the training data distribution (context inputs that are uniform on the sphere with labels that are independent Bernoulli random variables with parameter exactly 0.5 and query label that is generated exactly by a 1-NN classifier), especially the training label distribution, are very specific. It is not clear whether the message of the paper can generalize beyond this specific training distribution. Ideally, the paper would present convergence results for a more general class of training distributions, in the same vein as the class of test distributions it considers, and perhaps even when the query label is not generated by a 1-NN classifier. An even more general result may be that the attention unit behaves as a $k$-NN predictor where $k$ depends on some property of the training distribution.\n\n2. The training data distribution is not only specific, but also entails that the label of the query is not generated in the same way as the label of the context examples, which is inconsistent with practice. Specifically each example label must be independent of the corresponding input (as well as all other terms) while the query label does depend on the query and the context examples. The reasoning behind the statement that independence of $\\{\\mathbf{x}\\_{i}\\}\\_{i\\in[N+1]}$ and  $\\{\\mathbf{y}\\_{i}\\}\\_{i\\in[N]}$ “is also essential to properly study in-context learning of one-nearest neighbor” is incorrect; $\\mathbf{y}\\_i$ can depend on $\\mathbf{x}\\_i$ and predictors of both forms mentioned will achieve 50% accuracy.\n\n3. The distribution shift result is solid but not surprising since 1-NN should behave similarly on all distributions for which the true label of the query is in fact the label of the closest nearest neighbor in the context — it requires no knowledge of the data to achieve strong performance.\n\n4. From Lemma 4 and 5, it is not clear how $\\xi^k_1$ grows slower than $\\xi^k_2$, as is claimed and as is needed for the final result. For any $\\xi^k\\_1,\\xi\\_2^k$, the upper bound on $\\xi\\_1^{k+1}-\\xi^k\\_1$ can be dominated by $\\exp( \\text{poly}( N, d ) \\xi\\_1^k)$, which can be much larger than the lower bound on $\\xi\\_2^{k+1}-\\xi^k\\_2$ of $\\exp( -\\text{poly}( N, d ) \\xi\\_2^k)$ even when $\\xi_1^k = \\Omega(\\xi_2^k)$.\n\nMinor\n\n- The ICL tasks are binary classifications but the loss is the squared error.\n\n- In Theorem 1, the term $\\log(1 - (N\\sqrt{d}^{1/d}))$ should be $\\log(1 - (N\\sqrt{d}^{-1/d}))$?\n\n- Equation 2.6 describes gradient ascent, not gradient descent, and the step size is inverted\n\n- Lemma 4 clearly has multiple typos, one of which makes it not clear how to deduce the true statement.\n\n- The term “epoch” is improperly used in the experiments section.\n\n- The caption of Figure 1 says that error bars are plotted but this is false.\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nThe authors have addressed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the ability of single-layer transformers to learn the one-nearest neighbor (1-NN) prediction rule through in-context learning. It focuses on how transformers can handle nonparametric methods like 1-NN classification, moving beyond simpler tasks like linear regression that previous studies have focused on. The main contributions include establishing that a transformer with a softmax attention layer can minimize the training loss to zero in a highly non-convex landscape and behave like a 1-NN predictor under data distribution shifts. The authors establish rigorous theoretical results in line with their investigation.\n\nSTRENGTHS:\n**Originality.** To the best of my knowledge, this paper is the first to theoretically study the behavior of in context learning under distribution shifts for softmax attention transformers. In general, most theoretical works in related areas study either one layer model or linear attention models. \n\n**Quality.** The paper is well written and the claims are well-substantiated with detailed proofs and theorems\n\n**Clarity.** The paper is well-organized and clear. Notations and background are provided well for better understanding. Although I suggest the authors do a grammatical pass since there are minor mistakes throughout the paper.\n\nWEAKNESSES:\n1. I like the intuitions and reasoning provided for justifying assumption 1.  However, assuming that the data lies on a hypersphere and assuming no class imbalance seems rather strict to me. Further, it is far from practical. I would like to see more insights with respect to relaxing these assumptions. \n\n2. The results show that the model converges to 1-NN predictor on the training data even under SGD and random initialization. How well does this generalize to relaxing conditions on the input data lying on the hypersphere?\n\nQUESTIONS:\n1. Authors mention in the caption of Figure 5 about “error bars”. However, I don’t see any shaded error curve in their figure.\n\nLIMITATIONS:\nPlease refer to the weaknesses section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis submission considers learning to implement 1-NN in-context with a single self-attention layer. In particular, they consider training on in-context learning (ICL) sequences of form $(x1, y1), (x2, y2), … (x_N, y_N), x_{N+1}$, where $x_i$ are sampled from the $d$-dimensional unit sphere independently and with uniform probability, $y_i$ are i.i.d. $\\pm 1$ labels independent of $x_i$, and $x_{N+1}$ is a test point with a prediction target equal to the prediction of 1-NN. The authors show that with a specific weight initialization gradient descent on the population loss converges to a global minimum that corresponds to implementing 1-NN in-context. The proof relies on the observation that under the proposed weight initialization most parameters stay zero during optimization and it becomes possible to describe the dynamics with only 2 variables. While the loss function written as a function of these 2 variables is nonconvex, they show that these two variables converge to infinity, with their difference converging to infinity too. This limit corresponds to the 1-NN algorithm.\n\nThe authors also prove that the learned algorithm is robust to distribution shifts, with increasing robustness with the number of gradient descent iterations. Finally, they conduct experiments with random weight initialization and show that a single self-attention layer can be trained to implement 1-NN in-context and be robust to distribution shifts.\n\nSTRENGTHS:\n1. Overall the paper is well-written. The related work is properly referenced.\n2. Understanding what learning algorithms transformers can implement in-context and what in-context learning algorithms are learned during training is highly important. A large body of work shows that transformers can implement many of the standard supervised learning algorithms. The main findings of this submission are a good contribution to this body of work and show that transformers can implement 1-NN and under certain conditions learn to implement 1-NN when trained on ICL instances.\n\nWEAKNESSES:\nMy only concern with this submission is the generality of findings.\n* While the technique is interesting, it depends critically on the initialization scheme. As the experimental results hint, a single self-attention might be able to learn to implement 1-NN in-context even with standard initialization. It would be great to see a discussion on how the technique employed in this work can be useful for proving convergence under standard initialization.\n* As I understand the employed technique is also tied to the *k=1* case of k-NN. It is unclear whether the technique is general enough to be useful for $k>1$.\n\nQUESTIONS:\n* Lines 127-133: I recommend expanding this part a bit. Also, $W_{3,3}$ should be $-\\xi_2$ so that the softmax attention peaks on the example with highest dot product (i.e., the closest point as all points are on the unit sphere).\n* In the first equation of Lemma 4, the rightmost term should be $\\exp(2\\xi_1^k - \\xi_2^k)$.\n* Denoting sequences with index $k$ as $\\xi^k_1$ and $\\xi^k_2$ is confusing. I recommend using either $\\xi^{(k)}_1$ and $\\xi^{(k)}_2$ notation or even better, $\\xi_k$ and $\\zeta_k$.\n\nLIMITATIONS:\nThis submission would benefit form a discussion on the generality of the technique (see the \"weaknesses\" section).", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the theoretical ability of attention layers to implement a 1-nearest-neighbor (1-NN) classifier via in-context learning (ICL). While prior work has studied the ability for transformers to implement algorithms such as linear regression, this paper is the first to establish that attention layers can also implement the 1-NN prediction rule in-context -- a non-parametric learning algorithm, for which in-context learning with attention seems particularly well-suited. The paper contributes the following: a convergence analysis of 1-NN via ICL, a characterization of how ICL 1-NN performs under distribution shift, and under careful initialization, the non-convexity of transformer optimization becomes tractable.\n\nSTRENGTHS:\n- The analysis seems quite complete -- the authors provide theoretical claims surrounding training and initialization, test-time ICL, and out-of-distribution ICL at test time. \n- The fact that the convergence analysis is done for 1-NN, a non-parametric learner, is conceptually interesting. To me, this idea makes a lot of sense because attention seems to do some form of non-parametric learning in-context at test time anyway. Overall, this gives me some hope that the analysis could be a useful tool for understanding the in-context learning ability of transformers, more generally. \n- The fact that the convergence analysis, which occurs in a non-convex setting, is solvable using more careful analysis is also interesting.\n\nWEAKNESSES:\n- It would be great for the authors to contextualize the work a bit more in terms of understanding transformers as language models, more generally. I understand that this work is primarily theoretical, but I feel that it hints at a key point that isn't coming through very strongly in the text: in-context learning with attention, seems to do a form of non-parametric learning at test time. Understanding how attention implements basic non-parametric learning methods such as 1-NN is a great first step toward understanding how attention does this, and it would be nice to include some commentary (or even speculation) on how this paper could fit into this broader narrative. \n- Can the analysis be extended in any trivial way to k-NN (e.g. using multiple heads)? There seems to be a lack of commentary on this, and if this is doable in some simple way, this result should be included. If it turns out to be non-trivial, the paper could benefit from commentary on this as well. \n- Throughout the paper, the authors refer to the input examples (xs and ys) as being either independent or not independent. It was somewhat unclear to me whether independence was being used throughout, or not.\n\nQUESTIONS:\n- In equation 2.2, the softmax output is directly multiplied with the embedding matrix, meaning that the W_v matrix is set to the identity. While this is spelled out in the text, including this in the equation would improve clarity. \n- In Assumption 2, I think $\\sigma_1$ should just be $\\sigma$. \n- Some of the equations run off of the right-hand side of the page in the appendix.\n\nLIMITATIONS:\nAs a theoretical work, social impact is non-applicable. However, the authors do not seem to discuss any limitations of their analysis, more broadly.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}