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

Authors: zihao  li
yuan  cao
cheng  gao
yihan  he
han  liu
jason m klusowski
jianqing  fan
mengdi  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:07:37

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.

Main paper: 
 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 , Von 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.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. 
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.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: +1) .H = [h 1 , h 2 , . . . , h N +1 ] =   x 1 x 2 . . . x N x N +1 y 1 y 2 . . . y N 0 0 0 . . . 0 1   ∈ R (d+2)×(N(2.1)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 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 byy i * := N i=1 1(i = argmin j∈[N ]∥x N +1 -x j ∥ 2 )y i .We also define i* = argmin i∈[N ] ∥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 } 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 ).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 .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:H W = H • softmax(H ⊤ W ⊤ K W Q H), (2.2)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:y W (x N +1 ; {x i , y i } i∈[N ] ) := [H W ] (d+1,N +1) = N j=1 y j exp(h ⊤ j Wh N +1 ) N +1 j=1 exp(h ⊤ j Wh N +1 ) . (2.3)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 denoteW =   W 11 W 12 W 13 W 21 W 22 W 23 W 31 W 32 W 33   , (2.4) 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 32and W 33 ∈ R. Then by Eq. 2.3, it is easy to see that W i2 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: considerW 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.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 byL(W) = 1 2 E {xi,yi} i∈[N ] ,x N +1 y W (x N +1 ) -y i * 2 . (2.5)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:W k+1 -W k = 1 η ∇ W L(W k ).(2.6)We shall consider the following initialization for the gradient descent:Assumption 2 (Initialization). Let σ > 0 be a parameter. We assume the following initialization:W 0 = 0 (d+1)×(d+1) 0 d+1 0 d+1 -σ ,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.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. 
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.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 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.(N d), -log 1 -(N √ d) 1 d , C d 1 -1 2 N }),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 .(ii) For all {(x i ,y i )} i∈[N ] ∪ {x N +1 } ∼ P test , we have x i ∈ S d-1 for all i ∈ [N + 1].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 defineA δ := {∥x j -x N +1 ∥ 2 2 ≥ ∥x i * -x N +1 ∥ 2 2+ δ for all j ̸ = i * such that y j ̸ = y i * }, then, after K-iterations of gradient descent, we haveE {(xi,yi)} i∈[N ] ,x N +1 y W K (x N +1 ) -y i * 2 ≤ O inf δ R 2 N 2 K -poly(N,d)δ + R 2 P test (A c δ ) ,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 haveP test Round y W k (x N +1 ) ̸ = y i * ≤ O inf δ M 2 N 2 K -poly(N,d)δ + M 2 P test (A c δ ) .Here we define Round(t) := 1[t]< 1 2 ⌊t⌋ + 1 [t]≥ 1 2 ⌈t⌉,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 haveP test Round y W k (x N +1 ) ̸ = y i * = 0 whenever K ≥ O log(M N ) poly(N,d)δ * .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(M N ) 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.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. 
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:∇ W11 L(W k ) = E N i=1 g k i (x ⊤ i x N +1 ) • x i x ⊤ N +1 + g k i * (x ⊤ i * x N +1 ) • x i * x ⊤ N +1 (4.1)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∈{xi} i∈[N ] ∥x -x N +1 ∥ 2 sampled i.i.d. from a uniform distribution on S d-1 . Moreover, we have ∇ Wij 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 betweenx ⊤ 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 xi = i∈[N +1] P U xi 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 ∇ W11 L(W k ) illustrated by Eq. (4.1), we shall always have U ∇ W11 L(W k )U ⊤ = ∇ W11 L(W k ), which shows ∇ W11 L(W k ) = c k I d forsome 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:W k = diag{ξ k 1 , . . . , ξ k 1 d times , 0, -ξ k 2 }.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:L(ξ 1 , ξ 2 ) := E N j=1 exp(ξ 1 ⟨x j , x N +1 ⟩)y j N i=1 exp(ξ 1 ⟨x i , x N +1 ⟩) + exp(ξ 1 -ξ 2 ) -y i * 2 . (4.2)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 d η (ξ k+1 1 -ξ k 1 ) ≥ c 1 • exp(-6ξ k 1 ) -c 2 • exp(2ξ k 1 -ξ k 2 ),and d η (ξ k+1 1 -ξ k 1 ) ≤ c 3 • exp poly(N, d) • ξ k 1 -c 4 • exp 2(ξ k 1 -ξ k 2 )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 c ′ 1 • exp(-poly(N, d) • ξ k 2 ) ≤ 1 η (ξ k+1 2 -ξ k 2 ) ≤ c ′ 2 • exp(-poly(N, d) • ξ k 2 ).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. 
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 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.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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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.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. 
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. 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). 
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). 
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. 
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. 
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. 
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]  
 
 Appendix: 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.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 productd 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 writea 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 .In our experiment, we generate the testing dataset with context length N and dimension d with separation parameter δ * such that ∥xj -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 1In 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.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:(1) We prove that when W k is a diagonal matrix with[W k ] [d]×[d] = c k I dfor some constant c k , we can have the following break downs of our proof:1 η (W k+1 -W k ) = diag{∆ξ 0 1 , . . . , ∆ξ 0 1 d times , 0, -∆ξ 0 2 }, (C.1)where ∆ξ 0 1 and ∆ξ 0 2 are two positive constants.(2) By W k being a diagonal matrix with 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.[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 (Step (1) of Induction. Recall that our loss function can be written asL(W) = 1 2 E {xi,yi} i∈[N ] ;x N +1 y W (x N +1 ) -f (x N +1 ; {x i , y i } i∈[N ] ) 2 = 1 2 E {xi,yi} i∈[N ] ;x N +1 N j=1 exp x ⊤ j W 11 x N +1 + y † j W ‡ j x N +1 + x j ⊤W 13 + y † j W † j y j N +1 j=1 exp x ⊤ j W 11 x N +1 + y † j W ‡ j x N +1 + x j ⊤W 13 + y † j W † j -y i * 2 where W =   W 11 W 12 W 13 W 21 W 22 W 23 W 31 W 32 W 33   , 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 32and 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 ofq j (x, W) := exp x ⊤ j W 11 x N +1 + y † j W ‡ j x N +1 + x j ⊤W 13 + y † j W † j i exp x ⊤ i W 11 x N +1 + y † i W ‡ j x N +1 + x i ⊤W 13 + y † j W † j (C.2)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 ∇ Wij L(W 0 ) respectively. Note that for f j θ (x) = exp(g j θ (x)), we have∇ θ j f j θ (x) • y j j f j θ (x) -y i * 2 = 2 j f j θ (x) • y j j f j θ (x) -y i * • j f j θ (x)y j ∇ θ g j θ j f j θ (x) - j f j θ (x)∇g j θ (x) N +1 j=1 f j θ (x) • j f j θ (x)y j j f j θ (x), therefore, with some algebra, we have the following closed-form formula for ∇ Wij L(W) respectively,∇ W11 L(W) = E N +1 j=1 q j (x, W)(x j x ⊤ N +1 )y j - N +1 j=1 q j (x, W)(x j x ⊤ N +1 ) N +1 j=1 q j (x, W)y j • N +1 j=1 q j (x, W)(y j -y i * ) , (C.3) ∇ W21 L(W) = E N j=1 q j (x, W)(y j x ⊤ N +1 )(y j ) - N j=1 q j (x, W)(y j x ⊤ N +1 ) N +1 j=1 q j (x, W)(y j ) • N +1 j=1 q j (x, W)(y j -y i * ) , (C.4) ∇ W13 L(W) = E N +1 j=1 q j (x, W)(y j )x j - N +1 j=1 q j (x, W)(x j ) N +1 j=1 q j (x, W)(y j ) • N +1 j=1 q j (x, W)(y j -y i * ) , (C.5) ∇ W23 L(W) = E N j=1 q j (x, W)(y j )y j - N j=1 q j (x, W)y j N +1 j=1 q j (x, W)(y j ) • N +1 j=1 q j (x, W)(y j -y i * ) , (C.6) ∇ W31 L(W) = E - N +1 j=1 q j (x, W)(y j ) • q N +1 (W)x ⊤ N +1 • N +1 j=1 q j (x, W)(y j -y i * ) , (C.7)and∇ W33 L(W) = E - N +1 j=1 q j (x, W)y j q N +1 (W) • N +1 j=1 q j (x, W)(y j -y i * ) . (C.8)By the induction assumption, we haveq j (x, W k ) = 1 j̸ =N +1 exp(ξ k 1 ⟨x j , x N +1 ⟩) + 1 j=N +1 exp(ξ k 1 ∥x N +1 ∥ 2 2 -ξ k 2 ) N i=1 exp(ξ k 1 ⟨x i , x N +1 ⟩) + exp(ξ k 1 ∥x N +1 ∥ 2 2 -ξ k 2 ) = 1 j̸ =N +1 exp(ξ k 1 ⟨x j , x N +1 ⟩) + 1 j=N +1 exp(ξ k 1 -ξ k 2 ) N i=1 exp(ξ k 1 ⟨x i , x N +1 ⟩) + exp(ξ k 1 -ξ k 2 ) (C.9)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 ∇ Wij L(W), we can make the following calculation. First, we calculate ∇ W11 L(W k ):∇ W11 L(W k ) = E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 )(y j -y i * ) - N +1 j=1 q j (x, W k )(x j x ⊤ N +1 ) N +1 j=1 q j (x, W k )(y j -y i * ) • N +1 j=1 q j (x, W k )(y j -y i * ) = E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 )(y j -y * i ) N +1 j=1 q j (x, W k )(y j -y * i ) (1) -E N +1 j=1 q j (x, W k )x j x ⊤ N +1 N +1 j=1 q j (x, W k )(y j -y * i ) 2 (2), since q(x, W k ) is only a function of x by our induction assumption, we have(1) = E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 )y j N +1 j=1 q j (x, W k )y j (i) -E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 )y j y i * (ii) -E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 )y i * N +1 j=1 q j (x, W k )y j (iii) + E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 )(y i * ) 2 (iv), here(i) = E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 ) y j N +1 j ′ =1 q j (x, W k )y j ′ = E N j=1 q j (x, W k ) 2 (x j x ⊤ N +1 ) , (ii) = -E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 )E y j y i * x 1:N +1 = -E N j=1 q j (x, W k )(x j x ⊤ N +1 ) N +1 j ′ =1 1 j ′ =argmax i∈[N ] ⟨x N +1 ,xi⟩ E[y j ′ y j |x 1:N +1 ] = -E q i * (x, W k )x i * x ⊤ N +1 , (iii) = -E N +1 j,j ′ =1 q j (x, W k )q j ′ (x, W k )(x j x ⊤ N +1 )y i * y j ′ = -E N +1 j=1 q j (x, W k )q i * (x, W k )(x j x ⊤ N +1 ) (iv) = E N +1 j=1 q j (x, W k )(x j x ⊤ N +1 ) ,also for (2), we have(2) = -E N +1 j=1 q j (x, W k )x j x ⊤ N +1 E N +1 j=1 q j (x, W k )(y j -y i * ) 2 x 1:N +1 = -E N +1 j=1 q j (x, W k )x j x ⊤ N +1 N +1 j,j ′ =1 q j (x, W k )q j ′ (x, W k )E y j y j ′ -(y j + y j ′ )y i * + y 2 i * x 1:N +1 = -E N +1 j=1 q j (x, W k )x j x ⊤ N +1 1 + N j=1 q 2 j (x, W k ) - N +1 j,j ′ =1 q j (x, W k )q j ′ (x, W k )E (y j + y j ′ )y i * x 1:N +1 = -E N +1 j=1 q j (x, W k )x j x ⊤ N +1 1 + N j=1 q 2 j (x, W k ) -2q i * (x, W k )(1)+(2) gives us∇ W11 L(W k ) = E N j=1 q 2 j (x, W k )(x j x ⊤ N +1 ) -E q i * (x, W k )(x i * x ⊤ N +1 ) + E q i * (x, W k ) N +1 j=1 q j (x, W k )(x j x ⊤ N +1 ) -E N j=1 q 2 j (x, W k ) N +1 j=1 q j (x, W k )(x j x ⊤ N +1 ) . (C.10) For ∇ W21 L(W k ), we have ∇ W21 L(W k ) = E N j=1 q j (x, W k )(y j x ⊤ N +1 )(y j ) - N j=1 q j (x, W k )(y j x ⊤ N +1 ) N +1 j=1 q j (x, W k )y j • N +1 j=1 q j (x, W k )(y j -y i * ) ,note that it is the expectation of multiplication of three linear functions of y i , therefore ∇ W21 L(W k ) equals to 0 by symmetry. For ∇ W31 L(W k ), we have∇ W31 L(W k ) = E - N +1 j=1 q j (x, W k )(y j ) • q N +1 (x, W k )x ⊤ N +1 • N +1 j=1 q j (W k )(y j -y i * ) = 0,here the second equality comes from q j (x, W k ) being an even function of x and x has a symmetric distribution. For ∇ W13 L(W k ), we have∇ W13 L(W k ) = E N +1 j=1 q j (W KQ )(y j )x j - N +1 j=1 q j (W KQ )(x j ) N +1 j=1 q j (W KQ )(y j ) • N +1 j=1 q j (W KQ )(y j -y i * ) = 0here the second equality comes from q j (x, W k ) is a function of ∥x∥ 2 and x has a symmetric distribution. For ∇ W23 L(W k ), we have∇ W23 L(W k ) = N j=1 q j (W KQ )(y j )y j - N j=1 q j (W KQ )y j N +1 j=1 q j (W KQ )(y j ) • N +1 j=1 q j (W KQ )(y j -y bi * ) = 0,due to its symmetry in y. For ∇ W33 L(W k ), we have the following calculation:∇ W33 L(W k ) = E - N +1 j=1 q j (x, W k )y j q N +1 (x, W k ) • N +1 j=1 q j (x, W k )(y j -y i * ) = E q N +1 (x, W k ) y i * N +1 j=1 q j (x, W k )y j - N +1 j=1 q j (x, W k )y j 2 = E q N +1 (x, W k ) q i * (x, W k ) - N j=1 q 2 j (x, W k ) (C.11)With all previous calculations, we know that ∇ Wij L(W k ) equals 0 except for ∇ W11 L(W k ) and ∇ W33 L(W k ) with our induction assumption. Now we only need to prove that ∇ W11 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 haveN +1 i=1 P xi = N +1 i=1 P U xi for all U ∈ O(d).Proof. Since x i * are sampled i.i.d. from P x , we haveN +1 i=1 P xi = N +1 i=1 P U xi .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[∇ W11 L(W k )] = a k • I d .Proof. Note that by the induction assumption, we haveq j (x, W k ) = 1 j̸ =N +1 exp(ξ k 1 ⟨x j , x N +1 ⟩) + 1 j=N +1 exp(ξ k 1 -ξ k 2 ) N i=1 exp(ξ k 1 ⟨x i , x N +1 ⟩) + exp(ξ k 1 -ξ k 2 ), 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{g k i (x)} i∈[N ] ∪ {g k i * (x)} : R → R such that ∇ W k 11 L(W) = N i=1 g k i (x ⊤ i x N +1 ) • x i x ⊤ N +1 + g k i * (x ⊤ i * x N +1 ) • x i * x ⊤ N +1 ,here the expectation is taken with respect to {x i } i∈[N +1] . By Lemma 6, we haveP xi ⊗ P xj = P U xi ⊗ P U xj for any orthogonal matrix U ∈ O(d), therefore E[g k i (x ⊤ i x N +1 ) • x i x N +1 ] = E[g k i ((U x i ) ⊤ (U x N +1 )) • (U x i )(U x N +1 ) ⊤ ] = U E[g k i (x ⊤ i x N +1 ) • x i x ⊤ N +1 ]U ⊤ ,here the first inequality comes from P xi ⊗P xj = P U xi ⊗P U xj . Such identity holds for all orthogonal matrix U , thereforeE[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 ,x i * = argmin xi∈{xj } j∈[N ] ∥x i -x N +1 ∥ 2 , U x i * = U argmin xi∈{xj } j∈[N ] ∥U x i -U x N +1 ∥ 2 , therefore P(x ⊤ i * x N +1 |x 1:(N +1) ) = P(U x i * x ⊤ N +1 U ⊤ |U x 1:N +1 ). Since P(x 1:N +1 ) = P(U x 1:N +1 ), we have P(x ⊤ i * x N +1 ) = P(U x i * x ⊤ N +1 U ⊤ ), therefore E[g k i * (x ⊤ i * x N +1 ) • x i * x ⊤ N +1 ] = E[g k i * ((U x i * ) ⊤ (U x N +1 )) • (U x i * )(U x N +1 ) ⊤ ] = U E[g k i * (x ⊤ i * x N +1 ) • x i * x ⊤ N +1 ]U ⊤ .Again, since U is an arbitrary orthogonal matrix, we conclude thatE[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.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:(1) 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.(2) We provide an upper bound and lower bound for ξ k+11 -ξ k 1 and ξ k+1 2 -ξ k 2 .(3) We prove that there exits constant c 1 ∈ (0, 1)such that ξ k 1 ∈ (0, c 1 ξ k 2 ). (4) 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 haved η (ξ 0 1 -ξ 1 1 ) = ∇ W11 L(W 0 ) = E N j=1 q 2 j (x, W 0 )(x ⊤ j x N +1 ) -E q i * (x, W 0 )(x ⊤ i * x N +1 ) + E q i * (x, W 0 ) N +1 j=1 q j (x, W 0 )(x ⊤ j x N +1 ) -E N j=1 q 2 j (x, W 0 ) N +1 j=1 q j (x, W 0 )(x ⊤ j x N +1 ) ≤ - 1 N + 1 E[x i * x ⊤ N +1 ] + 1 (N + 1) 3 E[xx ⊤ ], (ξ 0 2 ≥ 0)To prove that ξ 0 1 -ξ 1 1 < 0, we only need- 1 N + 1 E[x i * x ⊤ N +1 ] + 1 (N + 1) 3 E[xx ⊤ ] ≤ 0.However, by Lemma 18, this can be guaranteed by N ≥ O(√ d log d).Step (2): Scale of ξ k+11 -ξ k 1 and ξ k+1 2 -ξ k 2Note that in every gradient update with stepsize η, we haved η (ξ k+1 1 -ξ k 1 ) = E q i * (x, W)(x ⊤ i * x N +1 ) -E N j=1 q 2 j (x, W)(x ⊤ j x N +1 ) -E q i * (x, W) - N j=1 q 2 j (x, W) N +1 j=1 q j (x, W)(x ⊤ j x N +1 ) (C.12) = E q i * (x, W) - N j=1 q 2 j (x, W) x ⊤ i * x N +1 - N +1 j=1 q j (x, W)(x ⊤ j x N +1 ) + E N j=1 q 2 j (x, W) x ⊤ i * x N +1 -x ⊤ j x N +1 (C.13)we also have the following estimation for the update of ξ 3 :1 η (ξ k+1 2 -ξ k 2 ) = E[q N +1 (x, W k ) q i * (x, W k ) - N j=1 q 2 j (x, W k ) ] ≥ E q N +1 (x, W k ) • q i * (x, W k ) -q i * (x, W k ) • N j=1 q j (x, W k ) = E[q i * (x, W k ) • q 2 N +1 (x, W k )], (C.14)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 have1 η • d(ξ k+1 1 -ξ k 1 ) + 2(ξ k+1 2 -ξ k 2 )} ≥ 1 - 1 2 N C d exp(-6ξ k 1 )for all k ≥ 0.Proof. First, note that we have the following relations:1 η • d(ξ k+1 1 -ξ k 1 ) + 2(ξ k+1 2 -ξ k 2 ) = E q i * (x, W)(x ⊤ i * x N +1 ) -E N j=1 q 2 j (x, W)(x ⊤ j x N +1 ) + E[q N +1 (x, W k ) q i * (x, W k ) - N j=1 q 2 j (x, W k ) ] -E q i * (x, W) - N j=1 q 2 j (x, W) N j=1 q j (x, W)(x ⊤ j x N +1 ) = E q i * (x, W) - N j=1 q 2 j (x, W) (x ⊤ i * x N +1 + 1) - N j=1 q j (x, W)(x ⊤ j x N +1 ) + E N j=1 q 2 j (x, W) x ⊤ i * x N +1 -x ⊤ j x N +1 ≥ E q i * (x, W) - N j=1 q 2 j (x, W) • q N +1 (x, W)(x ⊤ i * x N +1 + 1) + E N j=1 q 2 j (x, W) x ⊤ i * x N +1 -x ⊤ j x N +1 , (C.15)here the inequality comes fromx ⊤ i * x N +1 ≥ x ⊤ j x N +1 for all j ∈ [N ]. Moreover, note that when ξ 1 ≥ 0, we have1 = N j=1 q j (x, W) + q N +1 (x, W) ≤ N q i * (x, W) + q N +1 (x, W), which is equivalent to q i * (x, W) ≥ 1-q N +1 (x,W) N . Since q i * (x, W) q N +1 (x, W) = exp ξ 1 (x ⊤ i * x N +1 -1) + ξ 3 ,we haveq i * (x, W) ≥ 1 N + exp(ξ 1 (1 -x ⊤ i * x N +1 ) -ξ 3 ) . (C.16)Therefore we haveE N j=1 q 2 j (x, W) x ⊤ i * x N +1 -x ⊤ j x N +1 =E q 2 i * (x, W k ) N j=1 q 2 j (x, W k ) q 2 i * (x, W k ) • x ⊤ i * x N +1 -x ⊤ j x N +1 ≥E 1 N + exp(ξ 1 (1 -x ⊤ i * x N +1 ) -ξ 3 ) 2 N j=1 exp 2ξ k 1 x ⊤ j x N +1 -x ⊤ i * x N +1 x ⊤ i * x N +1 -x ⊤ j x N +1 ≥ exp(-4ξ k 1 ) (N + exp(ξ k 1 -ξ k 2 )) 2 E N j=1 x ⊤ i * x N +1 -x ⊤ j x N +1 |x ⊤ i * x N +1 ≥ 0 P(x ⊤ i * x N +1 ≥ 0) ≥ 1 - 1 2 N C d exp(-4ξ k 1 ) (N + exp(ξ k 1 -ξ k 2 )) 2 ≥ 1 - 1 2 N C d exp(-6ξ k 1 ) (C.17)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 thatE q i * (x, W) - N j=1 q 2 j (x, W) • q N +1 (x, W)(x ⊤ i * x N +1 + 1) ≥ 0,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.Lemma 9. When ξ k 1 ≥ 0 and N ≥ O( √ d log d) as in Theorem 1, we have the following inequalities:1 η (ξ k+1 1 -ξ k 1 ) ≤ 2N d exp(- 4 (N + 1) 2 ξ k 1 ) - a n,d dN 3 e exp(2(ξ k 1 -ξ k 2 )),wherea n,d = 2N √ d) -2 d-3, and1 η (ξ k+1 2 -ξ k 2 ) ≤ exp(2ξ k 1 -ξ k 2 ).Proof. We first prove the first inequality. Recall thatd η (ξ k+1 1 -ξ k 1 ) = E q i * (x, W) - N j=1 q 2 j (x, W) x ⊤ i * x N +1 - N +1 j=1 q j (x, W)(x ⊤ j x N +1 ) + E N j=1 q 2 j (x, W) x ⊤ i * x N +1 -x ⊤ j x N +1 , when ξ k 1 ≥ 0, we have q i * (x, W) - N j=1 q 2 j (x, W) ≥ 0, therefore d η (ξ k+1 1 -ξ k 1 ) ≤ E N j=1 q 2 j (x, W) x ⊤ i * x N +1 -x ⊤ j x N +1 (i) -E q i * (x, W k ) - N j=1 q 2 j (x, W k ) q N +1 (x, W k )(1 -x ⊤ i * x N +1 ) (iii) + E q i * (x, W) - N j=1 q 2 j (x, W) N j=1 q j (x, W)(x ⊤ i * x N +1 -x ⊤ j x N +1 ) (ii)We have the following bound for (iii) when ξ k 1 ≥ 0:-(iii) ≥ 1 e 1 2N √ d) 2 d-3 • E q i * (x, W k ) - N j=1 q 2 j (x, W k ) q N +1 (x, W k ) x ⊤ i * x N +1 ≥ 1 -a n,d ≥ 1 e 1 2N √ d) 2 d-3 • E q i * (x, W k )q 2 N +1 (x, W k ) ≥ 1 e 1 2N √ d) 2 d-3 • 1 N • E[ 1 -q N +1 (x, W k ) q 2 N +1 (x, W k )] ≥ 1 N 3 1 e 1 2N √ d) 2 d-3 exp(2(ξ k 1 -ξ k 2 ))where the first inequality comes from Lemma 19. Now we aim to bound (i) and (ii) separately. First, we have(i) = E q 2 i * (x, W k ) N j=1 q 2 j (x, W k ) q 2 i * (x, W k ) x ⊤ i * x N +1 -x ⊤ j x N +1 = E N j=1 exp -2ξ k 1 x ⊤ i * x N +1 -x ⊤ j x N +1 x ⊤ i * x N +1 -x ⊤ j x N +1 ≤ N exp -2ξ k 1 E[x ⊤ i * x N +1 -x ⊤ j x N +1 ] E[x ⊤ i * x N +1 -x ⊤ j x N +1 ] = N exp(-2ξ k 1 E[x ⊤ i * x N +1 ])E[x ⊤ i * x N +1 ] ≤ N exp - 4 (N + 1) 2 ξ k 1 (C.18)here the first inequality comes from Jensen's inequality and the convexity of x exp(-x) between [0, 2], and the second one comes fromE[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:(ii) ≤ E q i * (x, W) -q 2 i * (x, W) N j=1 q j (x, W)(x ⊤ i * x N +1 -x ⊤ j x N +1 ) = E 1 -q i * (x, W k ) N j=1 q j (x, W k ) q i * (x, W k ) (x ⊤ i * x N +1 -x ⊤ j x N +1 ) ≤ E min{1, ζ k } N j=1 exp -ξ k 1 x ⊤ i * x N +1 -x ⊤ j x N +1 x ⊤ i * x N +1 -x ⊤ j x N +1 ≤ N exp E[-ξ k 1 x ⊤ i * x N +1 -x ⊤ j x N +1 ] E[x ⊤ i * x N +1 -x ⊤ j x N +1 ] ≤ N exp(-ξ k 1 /2) where we define a random variable ζ k by ζ k = (N -1) exp ξ k 1 • {x ⊤ (2) x N +1 -x ⊤ i * x N +1 } + exp(ξ k 1 • {1 -x ⊤ i * x N +1 } -ξ k 2 ), the second inequality comes from 1 -q i * (x, W k ) ≤ 1 - exp(ξ k 1 • x ⊤ i * x N +1 ) exp(ξ k 1 -ξ k 2 ) + (N -1) exp(ξ k 1 • x ⊤ (2) x N +1 ) + exp(ξ k 1 • x ⊤ i * x N +1 ) ≤ (N -1) exp ξ k 1 • {x ⊤ (2) x N +1 -x ⊤ i * x N +1 } + exp(ξ k 1 • {1 -x ⊤ i * x N +1 } -ξ k2 ) 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 have1 η (ξ k+1 2 -ξ k 2 ) = E[q N +1 (x, W k ) q i * (x, W k ) - N j=1 q 2 j (x, W k ) ],therefore we have1 η (ξ k+1 2 -ξ k 2 ) ≤ E[q N +1 (x, W k ) q i * (x, W k ) -q 2 i * (x, W k ) ] ≤ E min{1, ζ k } q N +1 (x, W k ) q i * (x, W k ) ≤ E 1 1 + exp (x ⊤ i * x N +1 -1)ξ k 1 + ξ k 2 ≤ E[exp (1 -x ⊤ i * x N +1 )ξ k 1 -ξ k 2 ] ≤ exp(2ξ k 1 -ξ k 2). 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 d η (ξ k+1 1 -ξ k 1 ) ≥ 1 - 1 2 N C d exp(-6ξ k 1 ) -2 exp(2ξ k 1 -2ξ k 2 )The next lemma provides a lower bound for the update of ξ 3 . Notably, whenξ k 1 is close to ξ k 2 , ξ k 2increases in a larger scale.Lemma 11. When ξ k 1 ≥ 0, we have the following inequality holds:1 η (ξ k+1 2 -ξ k 2 ) ≥ 1 (N + 1) 3 e exp 2a n,d • ξ k 1 -2ξ k 2 ,where a n,d is a constant only pertains to n and d that is defined in Lemma 9Proof. With Eq. (C.11), we have1 η (ξ k+1 2 -ξ k 2 ) ≥ E[q i * (x, W k ) • q 2 N +1 (x, W k )] = E q 3 i * (x, W k ) q 2 N +1 (x, W k ) q 2 i * (x, W k ) = E q 3 i * (x, W k ) exp 2ξ k 1 (1 -x ⊤ i * x N +1 ) -2ξ k 2 ≥ 1 e E q 3 i * (x, W k ) exp 2a n,d • ξ k 1 -2ξ k 2 |x ⊤ i * x N +1 ≥ a n,d ≥ 1 (N + 1) 3 e exp 2a n,d • ξ k 1 -2ξ k 2where 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 2N 4 d ) = O max{log(N d), -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 have1 η (ξ k+1 1 -ξ k 1 ) ≤ 2N d exp(-ξ k 1 /2) - a n,d dN 3 e exp(2(ξ k 1 -ξ k 2 )), therefore, if 2N d exp(-ξ k 1 /2) < a n,d dN 3 e exp(2(ξ k 1 -ξ k 2 )), which is equivalent to 2.5ξ k 1 > log( a n,d 2N 4 d ) + ξ k 2 ,we will have ξ k+11 -ξ k 1 < 0. By Lemma 11, we have ξ k 2 ≥ ξ 0 2 ≥ 3 log( a n,d 2N 4 d ).Therefore ξ k+1 1 will not increase as long as15 7 ξ k 1 ≥ ξ k 2 , finally, recall that a n.d = 1 - 1 2N k d ) 2 d-3, 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ξ k 1 , ξ k 2 = Ω(η • poly(N, d) • log k), 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ξ k+1 2 -ξ k 2 ≤ η • exp(- 1 15 ξ k 2 ), and ξ k+1 2 -ξ k 2 ≥ η • 1 (N + 1 3 e) exp(-2ξ k 2 )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 8ξ k 1 ≤ ξ k 2 -log C d 1 -1 2 N 4 ,we haveξ k+1 1 -ξ k 1 ≥ η d 1 - 1 2 N C d exp(-6ξ k 1 ) ≥ 0. 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 ξ k+1 1 -ξ k 1 ≤ 2N d exp(-ξ k 1 /2)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.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 haveE N +1 j=1 q(x, W)y j -y i * 2 ≤ O N 3 k 2 d ξ 1 + exp(2ξ 1 -ξ 3 ) Proof.We haveE N +1 j=1 q(x, W)y j -y i * 2 = 1 + E N j=1 q 2 j (x, W) -2E q i * (x, W) = E[1 -q i * (x, W)] + E N j=1 q 2 j (x, W) -q i * (x, W) , note that E N j=1 q 2 j (x, W) -q i * (x, W) ≤ E[ N j=1 q i * (x, W)q j (x, W) -q i * (x, W)] ≤ E -q N +1 (W, x)q i * (x, W) ≤ 0, Therefore E N +1 j=1 q(x, W)y j -y i * 2 ≤ E[1 -q i * (x, W)]. However, we haveE 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 haveE 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 haveE N +1 j=1 q(x, W)y j -y i * 2 ≤ O poly(N, d) log K .Proof. By Lemma 15, we haveE N +1 j=1 q(x, W k )y j -y i * 2 ≤ O N 3 k 2 d ξ k 1 + exp(2ξ k 1 -ξ k 3 ).By Lemma 13, we have ξ k 1 ≤ 7 15 ξ k 2 , with ξ k 1 = Ω(poly(N, d) log k). Thus we haveE N +1 j=1 q(x, W k )y j -y i * 2 ≤ O N 3 k 2 d ξ k 1 + exp(- 1 15 ξ k 3 ) = poly(N, d) log k + 1 k 1/15 = poly(N, d) log k .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.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∥x j -x N +1 ∥ 2 2 ≥ δ + ∥x i * -x N +1 ∥ 2 2 , ∀j ̸ = i * , (D.1)Then we have| y(W k ) -y i * | = R N j=1 q j (x, W K )y j -y i * = j∈[N ],yj ̸ =y i * q j (x, W K )(y j -y i * ) + R|q N +1 (x, W K )| = 2R j∈[N ],yj ̸ =y i * q j (x, W K ) + R • q N +1 (x, W K ) ≤ 2R j∈[N ],yj ̸ =y i * q j (x, W K ) q i * (x, W K ) + R • q N +1 (x, W K ) q i * (x, W K ) = 2R j∈[N ],yj ̸ =y i * exp ξ K 1 x ⊤ j x N +1 -x ⊤ i * x N +1 + R • exp ξ K 1 (1 -x ⊤ i * x N +1 ) -ξ K 3 ) ≤ 2R • N exp(-ξ K 1 • δ) + R • exp - 1 15 ξ k 3 = 2RN exp -poly(N, d) • δ log K = O RN K -poly(N,d)δ (D.2)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 haveE P test y(W k ) -y i * 2 = E P test y(W k ) -y i * 2 1 A δ + E P test y(W k ) -y i * 2 1 A c δ ≤ O R 2 N 2 K -poly(N,d)δ + 4R 2 • P test (A c δ ), (D.3)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.In this section, we provide the proof for Corollary 1. The first statement of Corollary 1 comes from Markov's inequality:P test Round y W (x N +1 ) ̸ = y i * = P test | y W (x N +1 ) -y i * | ≥ 1 2 ≤ 4 • E[| y W (x N +1 ) -y i * | 2 ] ≤ O M N • inf δ≥0 K -poly(N,d)δ + P test (A c δ ) ,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| y(W k ) -y i * | ≤ O M N K -poly(N,d)δ * holds for all {(x i , y i )} ∪ {x N +1 } ∼ P test almost surely. Note that Round y W (x N +1 ) = y i * holds whenever | y(W k ) -y i * | < 1 2 , therefore, it suffices to have 1 2 ≤ O M N K -poly(N,d)δ * , which is equivalent to K = O log(M N ) poly(N, d)δ * .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 ,W = diag{ξ 1 , . . . , ξ 1 d times , 0, ξ 2 }. (E.1)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:L(ξ 1 , ξ 2 ) := E N j=1 exp(ξ 1 ⟨x j , x N +1 ⟩)y j N i=1 exp(ξ 1 ⟨x i , x N +1 ⟩) + exp(ξ 1 -ξ 2 ) -y i * 2 ,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∂ ξ2 L(0, ξ 2 ) = ∂ ξ2 E N j=1 y j N + exp(-ξ 2 ) -y i * 2 = ∂ ξ2 E 1 -2 N j=1 y j y i * N + exp(-ξ 2 ) + ( N j=1 y j ) 2 (N + exp(-ξ 2 )) 2 = ∂ ξ2 N N + exp(-ξ 2 ) 2 - 2 N + exp(-ξ 2 ) = -exp(-2ξ 2 ) N + exp(-ξ 2 ) 3 , (E.2)where the third equation comes fromE N j=1 y j y i * = E E N j=1 y j y i * {x i } i∈[N ] = E N j=1 1 j=i * E y j y i * {x i } i∈[N ] = E N j=1 1 j=i * E y 2 i * {x i } i∈[N ] = 1 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.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), defineY 1 = X 1 , . . . Y d-1 = X d-1 , Y d = X d d i=1 X 2 iNote that Y is distributed the same way as x i . Calculating the Jacobian:J =          1 0 • • • 0 0 0 1 . . . . . . . . . . . . 0 0 0 • • • 0 1 0 - X1X d [ d i=1 X 2 i ] 3/2 • • • -X1X d-1 [ N i=1 X 2 i ] 3/2 d i=1 X 2 i [ d i=1 X 2 i ] 3/2         Since J is of the form:J = I 0 a bthe determinant is easily evaluated:|J| = d i=1 X 2 i d i=1 X 2 i 3/2Now to solve for the distribution of Y N .f τ (y) = ∞ -∞ • • • ∞ -∞ 1 |J| f x (x)dxwheref x (x) = d I=1 1 √ 2π e -x 2 i 2 = 1 (2π) d/2 e -1 2 d i=1 x 2 iWriting the distribution in terms of t :f τ (t) = ∞ -∞ • • • ∞ -∞ d-1 i=1 y 2 i + Y 2 d 1-t 2 d-1 i=1 y 2 i 3/2 N -1 i=1 y 2 i 1 (2π) d/2 e -1 2 1 1-t 2 d-1 i=1 y 2 i dy 1 • • • dy d-1 f τ (t) = ∞ -∞ • • • ∞ -∞ 1 (1 -t 2 ) 3/2 d-1 i=1 y 2 i 1 (2π) d/2 e -1 2 1 1-t 2 d-1 i=1 y 2 i dy 1 • • • dy d-1 f τ (t) = 1 (1 -t 2 ) 3/2 1 (2π) N/2 ∞ -∞ • • • ∞ -∞ d-1 i=1 y 2 i e -1 2 1 1 -t 2 d-1 i=1 y 2 i dy 1 • • • dy d-1Now we can make a substitution to remove y d from inside the integral. Let u n = 1 -t 2 -1/2 y n , we haveEQUATIONf τ (t) = 1 -t 2 (d-3)/2 (2π) -d/2 ∞ -∞ • • • ∞ -∞ d-1 i=1 u 2 i e -1 2 d-1 i=1 u 2 i du 1 • • • du d-1The integral can be seen to be a constant so:f τ (t) = k d 1 -t 2 d-3 2or for notational convenience:EQUATIONSince f τ (y) is a PDF and is defined for 0 ≤ t ≤ 1 :1 0 k d 1 -t 2 d-3 2 dt = 1, where k d = 1 1 0 (1 -y 2 ) d-3 2 dy . 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 haveE[x ⊤ i * x N +1 ] ≥ P max i∈[N ] x ⊤ i x N +1 ≥ α • α + 1 -P max i∈[N ] x ⊤ i x N +1 ≥ α • (-1). To ensure that E[x ⊤ i * x N +1 ] ≥ 2 (N +1) 2 , we only need αP max i∈[N ] x ⊤ i x N +1 ≥ α -1 -P max i∈[N ] x ⊤ i x N +1 ≥ α ≥ 2 (N + 1) 2 ,which is equivalent toP max i∈[N ] x ⊤ i x N +1 ≥ α ≥ 1 1 + α 1 + 2 (N + 1) 2 .By Lemma 17, we haveP max i∈[N ] x ⊤ i x N +1 ≥ α = 1 -k d α -1 (1 -t 2 ) d-3 2 dt N .Therefore, we only need1 -k d α -1 (1 -t 2 ) d-3 2 dt N = P max i∈[N ] x ⊤ i x N +1 ≥ α ≥ 1 + 2 (N + 1) 2 • 1 1 + α ,which suffices byN ≥ log 1 -(1 + 2 (N +1) 2 ) • 1 1+α log(1 -k d 1 α (1 -t 2 ) d-3 2 dt).Note that1 α (1 -t 2 ) d-3 2 dt ≥ 1 α (1 -t) d-3 dt = 1 d -2 (1 -α) d-2 ,therefore we only needN ≥ log( α 1+α - α (1+α)(N +1) 2 ) k d d-2 (1 -α) d-2 . 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 forx ⊤ 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 haveP x ⊤ i * x N +1 ≤ 1 - 1 2N k d ) 2 d-3 ≥ 1 e .Proof. By Lemma 17,P x ⊤ i * x N +1 ≤ α = 1 -k d 1 α (1 -t 2 ) d-3 2 dt N .Note that1 α (1 -t 2 ) d-3 2 dt ≤ (1 -α)(1 -α 2 ) d-3 2 , since (1 -1 N ) N is monotonically increasing, we only need k d (1 -α)(1 -α) d-3 2 ≤ 1 N . Setting α = 1 - 1 2N k d ) 2 d-3 ≤ 1 - 1 N k d ) 2 d-3 1/2 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 haveE exp ξ x ⊤ (2) x N +1 -x ⊤ i * x N +1 ≤ O N 2 k 2 d ξMoreover, we haveE exp ξ x ⊤ (2) x N +1 -x ⊤ i * x N +1 = Ω 1 ξ ,where Ω(•) hides constant depends on N and d.Proof. We denotex ⊤ 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 ofY i being f d (t) = k d • (1 -t 2 ) d-3 2 , where k d = 2Γ( d 2 ) √ πΓ( d-12 ) . Then by the joint distribution of order statistics [David and Nagaraja, 2004] , the joint density function ofY i * and Y (2) is f Y i * ,Y2 (t)(y 1 , y 2 ) = N (N -1) • f d (y 2 )f d (y 1 ) • F N -2 d (y 2 )for -1 ≤ y 2 ≤ y 1 ≤ 1. Therefore, we haveE exp ξ x ⊤ (2) x N +1 -x ⊤ i * x N +1 ≤ N (N -1)k 2 d 1 -1 y1 -1 (1 -y 2 2 ) d-3 2 (1 -y 2 1 ) d-3 2 exp ξ(y 2 -y 1 ) dy 2 dy 1 ≤ N (N -1) 2 k 2 d 1 -1 y1 -1 (2 -y 2 1 -y 2 2 ) exp ξ(y 2 -y 1 ) dy 2 dy 1 ≤ N (N -1)k 2 d 1 -1 0 -1-y1 exp(ξy 2 )dy 2 dy 1 ≤ N 2 k 2 d 2 0 -2 exp(ξy 2 )dy 2 = O N 2 k 2 d ξ ,Thus we obtain the upper bound. Next, we establish a lower bound.E exp ξ x ⊤ (2) x N +1 -x ⊤ i * x N +1 ≥ N (N -1) 2 N k 2 d 1 0 (1 -y 2 1 ) d-3 y1 0 exp ξ(y 2 -y 1 ) dy 2 dy 1 = N (N -1) 2 N ξ k 2 d 1 0 (1 -y 2 1 ) d-3 1 -exp(-ξy 1 ) dy 1 ≥ N (N -1) e 5 2 N ξ k 2 d 2/ √ d 1/ √ d (1 -exp(-ξ/ √ d))dy 1 ≥ N (N -1) e 5 2 N √ dξ k 2 d (1 -exp(-ξ/ √ d))NeurIPS Paper Checklist 


Decision: The authors look at the ability of a one-layer softmax transformer to learn one-nearest neighbor following training by GD.  The authors assume the labels are independent of the context and show that GD-trained nets do succeed in learning the 1NN function. 

There was a decent level of support for this work among the reviewers, although some expressed reticence. I'm inclined to agree with the positive reviewers, as I think it's quite a nice contribution to know that GD over a transformer can provably learn a 1NN, at least when the labels are independent. Some other works in the DL theory literature also make rather strong assumptions on the generative mechanism (e.g. https://arxiv.org/abs/2305.15141) but I think the key is to provide insight into whether certain learning algorithms are possible to be implemented through training a transformer, and this work succeeds at that.  I recommend acceptance.

Reviews: review: paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: - 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. 
- 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. 
- 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.
questions: questions: - 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. 
- In Assumption 2, I think $\sigma_1$ should just be $\sigma$. 
- Some of the equations run off of the right-hand side of the page in the appendix.
limitations: limitations: As a theoretical work, social impact is non-applicable. However, the authors do not seem to discuss any limitations of their analysis, more broadly.
paper_summary: paper_summary: This 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.

The 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.
strength_weakness: strength_weakness: My only concern with this submission is the generality of findings.
* 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.
* 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$.
questions: questions: * 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).
* In the first equation of Lemma 4, the rightmost term should be $\exp(2\xi_1^k - \xi_2^k)$.
* 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$.
limitations: limitations: This submission would benefit form a discussion on the generality of the technique (see the "weaknesses" section).
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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. 

2. 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?
questions: questions: 1. Authors mention in the caption of Figure 5 about “error bars”. However, I don’t see any shaded error curve in their figure.
limitations: limitations: Please refer to the weaknesses section.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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.

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

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

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

Minor

- The ICL tasks are binary classifications but the loss is the squared error.

- In Theorem 1, the term $\log(1 - (N\sqrt{d}^{1/d}))$ should be $\log(1 - (N\sqrt{d}^{-1/d}))$?

- Equation 2.6 describes gradient ascent, not gradient descent, and the step size is inverted

- Lemma 4 clearly has multiple typos, one of which makes it not clear how to deduce the true statement.

- The term “epoch” is improperly used in the experiments section.

- The caption of Figure 1 says that error bars are plotted but this is false.
questions: questions: N/A
limitations: limitations: The authors have addressed limitations.

Comments: comment: Thank you for your positive review. We will carefully revise the paper, and make sure that all clarifications made in the rebuttal are included in the revised version.
Thank you for your positive feedback. In our revision, we will make sure to add clarifications and discussions about the points you have mentioned.
Thank you for your constructive comments, and for clarifying that you are not opposed to accepting our work. We agree that generalizing our current data distribution would greatly strengthen our results. After a careful examination of our proof, we believe that the data assumptions on $y_i$ in Assumption 1 could be extended to the following form:


(1) $\mathbb{E}(y_i y_j | \mathbf{x}\_{1:N}] = 0$ and $\mathbb{E}[y_i^2 | \mathbf{x}\_{1:N}] = 1$ for all $i \neq j, i,j \in[N]$.


(2) $\mathbb{P}(y\_{1:N} | \mathbf{x}\_{1:N}) = \mathbb{P}(y\_{1:N} | -\mathbf{x}\_{1:N}) $.


Note that such an assumption holds for a wide range of data distributions beyond the case where $\mathbf{x}_i$ and $y_i$ are independent. For example, the following data generating process gives $y_i$ that depends on $\mathbf{x}_i$, but the conditions (1) and (2) above still hold:


Consider an arbitrary fixed vector $\mathbf{a} \in R^d$ with $\|a\|_2 >2$. Suppose that $\mathbf{x}_i$, $i=1,\ldots,N$ are independently generated from the uniform distribution on the unit sphere, and supposed that given $\mathbf{x}_i$, $y_i$ is generated as follows:


- $y_i = 0 $ with probability $1- \frac{1}{ \max \\{\langle \mathbf{a}, \mathbf{x}_i \rangle^2, 1\\}} $,
- $y_i = \max \\{|\langle \mathbf{a}, \mathbf{x}_i \rangle|, 1\\}$ with probability $\frac{1}{2 \max \\{\langle \mathbf{a}, \mathbf{x}_i \rangle^2, 1\\}} $,
- $y_i = -\max \\{|\langle \mathbf{a}, \mathbf{x}_i \rangle|, 1\\}$ with probability $\frac{1}{2 \max \\{\langle \mathbf{a}, \mathbf{x}_i \rangle^2, 1\\}} $. 


It is easy to verify that $\mathbb{E}[y_i^2|\mathbf{x}\_{1:N}] =  \mathbb{E}[y_i^2|\mathbf{x}_i] = 1$, $\mathbb{E}[y_i y_j|\mathbf{x}\_{1:N}] = \mathbb{E}[y_i y_j|\mathbf{x}\_{i},\mathbf{x}\_{j}] = 0$ and $\mathbb{P}(y\_{1:N} | \mathbf{x}\_{1:N}) = \mathbb{P}(y\_{1:N} | -\mathbf{x}\_{1:N}) $. Moreover, $\mathbf{x}_i$ and $y_i$ are not independent, since $\mathbb{E}[y_i^4| \mathbf{x}_i ] =  \max\\{\langle \mathbf{a}, \mathbf{x}_i \rangle^2, 1 \\}$ is a function of $\mathbf{x}_i$.


We will update the paper to include the more general setting under conditions (1),(2) above. We assure that such an extension only requires minor modifications in the paper, and the proofs do not need any significant change. We believe that including such an extension can significantly strengthen our paper, and we hope that it can address your concerns on the limitation of our data models.
Thank you for the rebuttal. I will maintain my score.
I thank the authors for their thorough response, and for amending the typo in Lemma 4. I would not be opposed to seeing this paper accepted, since it provides novel, sound and non-trivial results towards addressing an important issue. Also, I am convinced that extensions to input data that is not uniform on the hypersphere and k-NN are highly non-trivial and probably worthy of separate papers. However, my main concern remains in that the results only apply to tasks generated by a very specific 1-NN-based data distribution in which the labels of the context examples are (unrealistically) independent of the input, whereas it is reasonable to expect that one softmax attention unit learns to behave like a 1-NN regressor in other settings as well, e.g. if the tasks are sinusoid regressions with high frequency, the best that softmax attention should be able to do is predict the label of the nearest input example.
Many thanks for the rebuttal that addresses many of the weaknesses identified and questions raised. I emphasize that all clarifications made during this rebuttal should be made in any revised manuscript to improve clarity of the work. Given my already positive review, I maintain my score.