One-Layer Transformer Provably Learns One-Nearest Neighbor In Context
Authors: Zihao Li
Yuan Cao
Cheng Gao
Yihan He
Han Liu
Jason Matthew Klusowski
Jianqing Fan
Mengdi Wang
Venue: NeurIPS 2024 poster (2024)
Keywords: In-context learning theory
softmax attention
one-nearest neighbor
Primary area: learning_theory

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
Zihao Li Yuan Cao Cheng Gao'
Yihan He'
Han Liu
Jason M. Klusowski' Jianqing Fan't Mengdi Wang"t
Princeeon University "The University of Hong Kong Northwestem University edu
{zihaoli, chenggao, yihan. he, jason. jasja klusowski jqf af,n, yuancao@hku. hk
hanliiuhorthwwsterrnn edu
Abstract
Transformers have achieved great success in recent years. Interestingly, transform- ers 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 one the most classical nonparametric estimators, the one-nearest neighbor prediction rule. Under theoretical framework where the prompt contains 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 soft- max attention layer can successfully learn to behave like 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.
Introduction
Transformers have emerged as one of the most powerful machine learning models since its intro- duction in Vaswani al. [017], achieving remarkable success in various tasks, including natural
language processing [Devlin al., 2018 Achiam et al., 2023 Touvron al. 2023 computer vision
Dosovitskiy et al.)
2020 He et al. [022 Saharia et al. 2022]. reinforcement learning [Chen et al.)
Janner et al.
2021 Parisotto et al.
2020]. and so on. One intriguing aspect transformers
2022 Min et al.) 2022 Wei their exceptional In-Context Learning (ICL) capability [Garg al.
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 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 demonstration prompt, consisting of sequence of labeled data and new unlabeled query. The in-context learning performance of 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. (2044) investigated ICL the optimization dynamics
"Equal Contribution.
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

of transformers under in in-context learning from theoretical lens, but their studies are limited to linear regression prediction rules, which is a sigmificant simplification of the transformer in-context learning task. Another line of work including Bai 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- basden 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 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 trans- former, demonstrating its ability to act like 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 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 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 provable result in both optimization and consecutive behavior under distribution shift for 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 dynamies of transformers based on gradient descent.
2.1 In-Context Learning Framework: One-Nearest Neighbor
In an In-Context Learning (ICL) instance, the model given a prompt (x))eEM) Prompe and query input X P query from some data distributions IP prompt and IP queryy. where {x; he(v] the input vectors, {y: hveee C R are the corresponding labels (e.g. real-valued for regression, or 1+11.11)-valuce for binary classification), and XN+1 is the query on which the model required to make a prediction. Given prompt y.))ee the prediction task is predict an ground truth model fxxN+1; [xx,, h:))y:::rn.ee )EE that maps the query token XN+1 to a real number.
In this work, we consider using transformers as the model to perform in-context learning. For a prompt Jt))e of length N and query token XN+1. we consider use the following embedding:
X1
X2
XN
XN+1
H h1, h2 bN+1]
Y1
Y2
YN
0
E R ((++2))((++1)
(2.1)
0
0
0
1
We use the notation of h, = Y,, for j < N, and hy+1 [xN+1,0, 1]. Here, (x.)seN) represents the input vectors, each associated with a corresponding label {yi)ie[N) where y; E TR. is the label. Throughout this paper, the sequence prompt exchangeably. The (d + 2)-th row serves as the indicator for the training token, which equals 0 value for E [N] and for N + analogous to a positional embedding vector. Such an indicator allows the model to distinguish the query token from the context. Similar models have been
are referred to as the context or

---

Sampled Points
Grid Color Decided by 1-NN
Figure 1: Illustration of data distribution in Assumption 1 S and the corresponding ground-truth division of 3 generated by one-ncarest neighbor. (1) In the left panel, the red and bluc points correspond to the x; with y, and -1 for i [N]. respectively, with N 500. (2) In the right panel, the color of every point on the sphere is the same its closest neighbor in (x: kE[]- The sphere thus split into divisions by the one-nearest- neighbo decision rule.
studied in line of recent works (Zhang et al. 2023 Huang et al. 2023 Chen et al. 2024 Bai et al. 2024 Akyürek 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 yi VeleeN)) of the training data constructed based on One-Nearest Neighbor (INN) data distribution, defined by the following definition.
Definition (One-Nearest Neighbor Predictor). Given prompt (i)eeN) and a query XN+1 we define the one- neighbor predictor by
y
argmin ||x -x,=11xxl12)y; JE[N]
We also define argminge(N] [x++1 x,112-
Without loss of generality, we assume that argmin ||x--1 -xylle is unique. Such assumption holds almost surely whenever {x, hie(N) is sampled from continuous distribution. Notably, for fixed prompt yx and query XN+1. Definition is identical to the nonparametric one- nearest neighbor estimator [Peterson] [009 Beyer et al.) 1999]. in which the algorithm outputs the label corresponding to the vector closest to the input, with the prompt {x,, hi) [N] as the training data in I-NN.
Next, we discuss the distribution of the training dataset y(); U (xx=I). Throughout the training process, we focus on the case in which {x))<e[N+1] are independently sampled from uniform distribution on a 1-dimensional sphere gd-1, {y: (yi)e is a zero-mean binary noise taken value in -11,-1}. with (x and {y: [ys]eE 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. (.11p we focus on the following underlying training distribution: (i) The sequence (xs]ie(N+)) are sampled inde- pendently from a uniform distribution on a d - dimensional sphere gd-1 C Rd. (ii) The labels {yE]VENN) satisfies Ely.y; [x:N]] and Ely?[xiN]] for all i + j,i.j [N]. (iii) We have P(Y1.N (1:N) Ply X1:N)-
Note that the case when {y: hej] and {x: JiENN being independent when (xiee with {x, (xi]
are uniformly sampled from the sphere and (y:heiN) are randomly sampled from (=1} an example
of Assumption [ We remark that by considering the training data distribution in Assumption [ 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 point out that this data distribution is still challenging to study, especially because of the assumption that the second order moment of {y: HEEN and {xI]e((N+1) uncorrelated. Due to

---

Figure 2: Hcatmap and landscape of loss function of single layer transformer when learning from one me--neaees ge by sampling 100 training sequences according to Assumption neighbor. The loss is defined in Eq. (2.5), generated with d N = We parametrize W (i, (1,6),,
such uncorrelation, the classifier given by one-nearest neighbor models are rather complicated. For example, Fig. |illustrates the randomly generated context data and the corresponding one-nearest neighbor prediction regions with 3 and N 500, which clearly demonstrate the complexity of the training task. This further leads to highly nonconvex and irregular objective function landscape, illustrated by Fig.
2.2 One-Layer Softmax Attention Transformers
We consider 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 scalar value:
Hw H softmax H W K Wo
(2.2)
where softmax A) applies softmax operator on each column of the matrix A, i.e. [softmax (attmaxx(A);: exp(A:))/{ (xp(())).. Our model is slightly different from the standard self-attentio transformers, as we consider a frozen value matrix. However, we also claim that such practice is common in deep learning theory Fang al. [2020). Lu et [2020), Mei 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 defined by the (d 1)-th element of the last column of Hw, with a closed form:
Σ - exp(h) Whx+1) N+ E exp(h) Why41)
(x. = [Wwk0+1N4+ =
(2.3)
which is the weighted mean of yi.... YN. Here and after, we may occasionally suppress dependence on {x,, yi kejN and write Sw (x (x, y.)eNN))) as (xN+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
W
W
12
W
13
W =
W
W W22
W
(2.4)
Wa
W 32
W 33
with Wi E Rddd W21 R Wai glxd W 12 E [ddx1 W Rixd W22, W23, Wa2
2.3 is easy to see that W does not affect yw for E [3]. which means we can simply take all these entries as zero in the following sections. Notably, for a fixed prompt-query pair (xx, and (xN+1), such an architecture allows an arbitrarily close approximation
and Wa E R. Then by Eq.
to the I-NN model: consider W = gf with gf goes to positive infinity, W = E such that ek -t converges to infinity, with the rest of W bounded, then y (XN+1)) converges y,- as A goes to infinity.

---

2.3 Training Dynamies
To train the transformer model over the I-NN task, we consider the Suan--quree Error (MSE) loss function. Specifically, the loss function is defined by
1 - - (xeen [(
L(W)=
(owx()x - -y-)) ).
(2.5)
Above, the expectation is taken with respect to the sampled prompt y:) and the query XN+1- Notably, when the underlying distribution for the prompt and query are defined by Assumption
I this loss function is nonconvex with respect to W. Such nonconvexity makes the optimization hard
to solve without further conditions such as PL condition KL condition [Bierstone and Milman 1988 Karimi et al., 2020]. We leave the proof of nonconvexity in Appendix
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 W = 1 V
(2.6)
17
We shall consider the following initialization for the gradient descent: Assumption (Initialization). Let a > O be parameter. We assume the following initialization:
w = (Ol++Dx(a+D) -"V-(aat
0643
0a+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 XN+1. e.g. Vaswani et al. [2017). Baade et al. 2022 Chang et al. [002]]. 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, because we want to keep the model focusing on the inner product between different x,, 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 a,, gradient descent is able to minimize the loss function L(W) despite the nonconvexity. Section 3.2 we further discuss the testing error of the trained transformer under distribution shift. Specifically, we consider distribution Piest for the prompt yx NE)] U [XN+1), which is different from the training data distribution PP prompt & I Puueey. 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 o, the loss function will converge to zero under gradient descent.
Theorem 1 (Convergence of Gradient Descent). Consider performing gradient descent of the sofimux-attention transformer model Y (x)+)) Suppose the initialization sutisfies Assumption 2 with 2(max (log(N Nd), log (1 (N (N ())). Ca(1 - where Ca =doly(d), and the number of context N 2 0 OVVdlog d), then I. k(W converges to 0.
We leave the detailed proof in Appendix C
Theorem T shows that the 1-NN data distribution, with large enough initialization constant a, the training loss of the transformer converges under gradient descent. Here a plays 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 greatly accelerate the training process Vaswani al. [2017] Devlin et al. (2018 Dosovitskiy et al. 2020). We also compare our results to existing
5

---

works. Zhang et al. [023] studied linear prediction tasks under gradient flow, however, their analysis 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 finite orthogonal dictionary. [hen al. [2024] established optimization convergence results for one-layer multi-heaz attention transformers under gradient flow. On the contrary, our work studies gradient descent convergence for transformer under 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.c., how the model extrapolate beyond the training distribution. Following the definition in Garg et al. [022], let us assume in the training process, the prompts {(x,, Y:) ~ .ptraia prompt and the query query During inference, the prompts and queries are sampled from a new distribution prea, We study the behavior of the trained transformers under possible prompt and query shift, i.e. postt A prompt & qaary- Our studies show that, under some mild conditions, the behavior of the trained model is still similar to 1-NN predictor even under distribution shift. Before formally stating our result, let us introduce the following assumption on the testing distribution:
XN+1 Iprrii
Assumption (Testing Distribution). We make the following assumption on piest
(i) There exists a R 2 0 such that lysl S R holds for all y: sampled from press
(ii) For all
U {xN+1} ~ ppe", we have x, E gd-l for all E [NV + 1].
Note that Assumption 3 only requires the label y; is bounded and x; is supported on sphere. We also remind the reader that we do not assume independence between different x; or (xi and (yi] Now we are ready to summarize our result in the following theorem.
Theorem 2 (Resemblance to I-NN predictor under Distribution Shift). Suppose Assumption and hold for prrinps
2 query and giest, If we define
As = - - X ||x,- xN+1 [12 + ô for all j + i such that Yj y..h
then, after K- K-iterations of gradient descent, we have
E ((xx)))))e 8041 [ (XN+1) -ox-1)--..)>) v.. < O(nn
+ R peeee (Ag)))
here the expectation is taken w.t ((x,, U {xN+1} ~ piest, Recall that y:- is the 1-NN
predictor of XN+1, which we defined in Definition
We leave the detailed proof in Appendix D Let us discuss the implication of Theorem 2 The event As describes the situation when the query XN+1 1 is located at an "inner point" away from its decision boundary, in which its distance to the nearest neighbor x,- is strictly larger than all other points. Such a quantity is similar to the margin condition in classification theory in deep learning [Battettt et al. 2017) and k-NN literature Chaudhuri strictly larger than all suboptimal choices. Specifically. if (Ptss((As)) 1 for some 8' > 0, i.e., the query XN+1 is strictly bounded away from the decision boundary almost surely, then the L2 distance between yw and the 1-NN predictor will converge in a ORR² (ok-poly(NNo)) even under a shifted distribution. We also introduce the following corollary, in which show that when y, only takes value in finite integer set, resembling a classification task, the trained transformer behaves like 1-NN predictor under an additional rounding operation.
CheuuhureCChauuuuri and Dasgupta (2014), where the optimal choice probability
Corollary 1 (Classfication of Trained Transformer). Suppose y; E [M] some integer M under Paestt then we have
P Prea( Round (fw= x..) S O(nf
2 {
Here we define
Round(t) == [0] [t] + 1 1022 [t],

---

i.e. the mapping from t E R to its closest integer, and As is defined as in Theorem 2 Moreover, if there exists 8 >> such that Prn (As-) = 0, then we have
pies
[Round (F (XW+)))))))))))))) + yz-) =0
whenever K 2 (
log(MWN) poly(N
We leave the detailed proof in Appendix D Corollary provides convergence rate for the clas- sification difference between 1-NN and the pretrained transformer. Notably, when XN+1 is well separated from the decision boundary in the testing distribution Peett the trained transformer will behave exactly the same the 1-NN classifier in O ( [og[[M) gradient steps for the pretrained transformer. Theorem 2 and Corollary [ show that the trained transformer under gradient descent is robust to both query and prompt distribution shift in the test distribution puestt in the sense that will maintain its resemblance to a I-NN predictor in both prediction and classification task, thus
extended the results in Zhang al. 2023 estimator.
Huang al. [2023 Chen et al. [2024) to nonparametric
4 Sketch of Proof
In this section, we sketch the proof of Theorem [nd highlight the techniques we used. The full proof is left to Appendix C
Swo-Dimensional Dynamic System. Recall that {x, Ne and the first and second moment of (yr)ietM are uncorrelated. Utilizing this uncorrelation between (x, Ne([++] and {v: (yike we can eliminate the reliance of the gradient on (yikeENI since we are considering population loss. Moreover, utilizing the structure of the initialization, we can prove by induction that all Wij will remain zero except for Wn and W 33. This shows that with suitable initialization, the transformer model will only focus on the relationship between different tokens throughout the whole training process. Our findings can be summarized by the following lemma.
Equivalence
Lemma (Closed-Form Gradient). With the initialization in Assumption the gradient of L ((W)) with respect to W11 can be written in the following form for all ls 2 0:
N
Vw.. = E
Σ (x.
xN+1) X + o (x](xx+1) x
(4.1)
where {6) (x) fgth (x)) R R is set of finctions. Here the expectation is taken with
respect to (E(N with X,. x ||x XN+1ll2 1l2 sampled i.i.d. from uniform
distribution on gd-1, Moreover, we have V Vw, L(W = 0 for all (i,j) E (3] x [3] and all 2 0 except for Wn and W 33-
Lemma D shows that we only need to consider W11 i1 and W33 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. XN+1 and x]_ x Recall that [x are i.i.d. sampled from uniform distribution supported
d - 1-dimensional sphere gd-1, therefore, the distribution of {xi hie[N+1] is rotational invariance, which means X {EE[N+1] Px, = +1E Pux. for all orthogonal matrix U E Raxdd Since the rotation of (x, [eEEN++1 does not change the inner products {x] x, [E[N and x7. x from the structure of V L(W*) illustrated by Eq. (4.I), we shall always have U w.,, L = Vw,, L(W*), which shows Vw,, Vw., ((W*) = csla for some constant C by simple algebra. We summarize our result the following lemma.
Lemma 2 -rw-Diimnsiioal System). With the initialization in Assumption 2 there exists two sets of real numbers {Et 1200 and (sh 1 k2o, such that w has the following form:
W diag ck ..... ct. 0, }.
tomes
7

---

With Lemma 2
we reduce the dimension of the original dynamic system in Eq. (2.6) from (d + 2) to 2. We now only need to focus on the evolution of č and E for all & 2 0.
Convergence of the Dynamic System. Lemma 2 training dynamics. However, this does not make our question trivial one, as the loss function is still highly nonconvex even when we only need to consider a two-dimensional subspace of (24+2)x(d+2), To see this, we introduce the following lemma:
helps us largely reduce the dimension of the
Lemma 3 (Nonconvexity of Transformer Optimization). When W lie in two-dimensional subspace of IR (d)+)))))+)) defined by W diag (1,...55 -2),, the original loss function defined in
times
Eg. 2.5 is equivalent to the following:
Σ exp((i
yi+ )
L2) :-
(4.2)
ii= exp(Ei (x,, XN+1)) + exp(E1 +xp(E: - E2)
Such loss function is still nonconvex.
We leave the detailed proof in Appendix Nonconexity shown by Lemma implies attaining the global minimum could be hard. Previous works such Zhang al. [2023] utilize conditions such as Polyak-Lojja Lojaseowjsiiiicz inequality to analyze such systems, however, those conditions are not applicable in our setting, and a more delicate analysis for the evolution of eh and & is needed. We characterize their behavior by the following lemma.
Lemma 4, For Et 2 there exists constants C1, C22 C32 C4 > 0, such that
a
d gett
<t) 2 ci exp(-66)) c2 exp(2ct +
and
d [ct ck+1 Et) S c3 exp (poly(N, <t) c4 exp (2 -
Lemma shows that there exits a constant e E (0, 1), such that c will keep increasing with a scale of (n log k) until c cugh. With this ratio, we obtain the following lemma for the increment of et
Lemma 5. For et 0, there exits constant ci, Cz. such that
1 cht+ - + S exp((poolyl (
& (pp(-pol((N,d) d) +
17
Lemma 5 shows that & will monotonically increase with a scale of (2 exp()))) which implies Ioo k). Combining Lemma 5nd we show that both & and & converge to infinity, with SI maintaining a slower speed, as its decreases when getting closer to & 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(gt; { will converge to zero as long as et and & - E both converges to infinity. We thus conclude our proof of I(W)) eventually converges global minimum.
Numerical Results
In previous sections, we have shown that with the initialization specified 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 I-NN predictor. First, we investigate the convergence of single-head single-layer transformers Vasswani et al. 2017 trained on I-NN tasks. The training data are sampled from pirain and pirain promps query. defined in Assumption I We choose context length N E {16, 32, 64) and input dimension d E (8, 16). The model is trained dataset with size of 10000, and an epoch

---

Training Loss on Uniform Sphere Distribution
Testing under Distribution
R.N
6.N
B.N
0.7
16, N =32
Testing MSE
0.5
0.4
0.3
0.2
0.1
200 500 750 1000
1250 1500 1750 2000 Gradient Iterations
250
500
750 1000 1250 1500 1750 2000
Gradient Iterations
Figure 3: Prediction error single softmax attention layer as a function of gradient iteration number. (1) The left panel shows the convergence of loss function during the training process. (2) The right pannel shows the MSE between the trained model and a 1-NN predictor on well separated testing dataset under distribution shift, we discuss in Section[5] Curves and error bars in both panels computed twice the standard deviation based 10 independent trials.
number of 2000. To ensure our training convergence result holds beyond the gradient descent scheme, we choose SGD as our optimizer, with batch size of 128 and learning rate of 0.1. We use the random Gaussian as our initialization. Our results for the training loss convergence are summarized the left panel of Fig. The results show that the model converges to I-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 the distribution shift, we generate testing data sampled from distribution difference from the training data, and report the mean square error between the model prediction and the I-NN predictor. Furthermore, the testing data satisfies P(AA)) = 1, with 8* specified as 0.1. Recall that defined A, in Theorem as the event where X is separated from the decision boundary with distance of at least ö. We leave the details of our data-generating process in Appendix 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 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 sequence of labeled training data and unlabeled test data, we demonstrate that, despite the non- convexity of the loss function during gradient descent training, 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 partially supported by the NIH ROILMO1372201 Jason M. Klusowski was supported in part by the National Science

---

Foundation through CAREER D 2239448, DMS-2054800 and HDR TRIPODS CCF-1934924. Jianqing Fan's research was partially supported by NSF grants DMM-2210833,, MMS-2003832, 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 Abemethy, Alekh Agarwal, Teodor Vanislavov Marinov, and Manfred Warmuth. mecha- nism 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 Gpt-4 technical report. arXiv preprini
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 algo- rithm is in-context learning? investigations with linear models. arXiv preprint 2rXiv:2211.15600 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-normalize margin bounds for neural networks. Advances neural information processing systems, 30, 2017.
Kevin Beyer, Jonathan Goldstein, Raghu Ramakrishnan, and Uri Shaft. When is "nearest neighbor" meaningful? In Database Theory- Cheorr-ICDTTPR 999: 7th International Conference Jerusalem, Israel, January 10-12, 1999 Proceedings 7, pages 217-235. Springer, 1999.
Edward Bierstone and Pierre Milman. Semianalytic and subanalytic sets. Publications Mathéma- tiques de I'THÉ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- 1311--1325,, 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 newral information processing systems, 34:1 04: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.
Ming-Wei
10

---

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 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 parameterizzee deep neural networks, 2020.
Shivam Garg, Dimitris Tsipras, Percy 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 IEEF/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 big sequence modeling problem. Advances in neural information processing systems, 34:1273-122 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-theorede analysis of in-context learning, 2024.
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-toji tojas--tojssiewic 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 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 Sciences, 115(33), July 2018. ISSN 1091-64999 doi: 10.1073/pnas. 110.57.9....107911..1 URL http: doi org/ 10./////1...0 1073/ 1073/pnas /paas 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. International conference on machine learning, pages 7487-7498.
PMLR, 2020.
Leif E Peterson. K-nearest neighbor. Scholarpedia, 4(2);

---

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 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,
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 Gomez, Lukasz 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, al. Larger language models do in-context learning differently, arXiv preprint arXiv:2303.088 00304400030006,, 2023.
Sang Michael Xie, Aditi Raghunathan. Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. arXiv preprint arXiv:2111.02080 2021.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers leam linear models in-context. arXiv preprint 09Xv:2306.09927, 2023.
12

---

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, [ai et al.) 22024) Abernethy et al. [2024], Akyürek al. [2022] studied the expressiveness of transformers to approximate statistical a algorithms. [i et al. [2023a] studied the generation ability of transformers in ICL tasks. Jeon et al. (2024) studies the information cettiiionnteeeticalll 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 PI. condition. Huuan et al. (2023) studied linear prediction tasks with a nolimax-attention layer trained with gradient descent, with finite support prompt/query distribution.
Chen et al. 2024] studied the gradient flow dynamics in training multi-head softmax attention linear prediction tasks. /hhn et al. (2024), GGannou et al. (2224) studied transformers' ability to learn optimization methods. Prior works also studied transformer optimization beyond in-context learning tasks. Tian al. [2023] studied single-layer transformer with one self-attention layer plus one decoder layer, and proved that the attention layer acts as scanning algorithm. [i et al. [023b)
studied the optimization of transformers in learning semantic structures. Jelassi et al. 2022 studies the spatial localization property vision transformer in optimization.
Notations. We write [N] = (1..... N}. For two vector x' = -x! mb) and x² = (x,
we write their inner product Σn x!2ę as x'Tx² We use On, and 0mxn denote the zero vector and zero matrix of size n and X respectively. We write {x E Rd ||x||2 1} as Sd-1 For two series {ax J220 and [120 we write
[((b)) if there exists 0 < C₁, C2 such that G by S ag S C2 bg- We write dg = O(bx) if there exists C > 0 such that as S C bg. We a poly(b,) if there exits an n-degree polynomial P. (z) such that ag O(,, (bx)). For (ax kEIN] we define its permutation (a(K))EE(N) such that a(1) 2 ((2) .... 2 ((n). We use La to denote the d-dimensional identity matrix and sometimes we also use 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]ig-
B Data-Generating Process in Experiment
In our experiment, we generate the testing dataset with context length N and dimension with separation parameter 8 such that ||x, xx-11l2 2 ||x,- XN+1l13 8 for all j E [N], j A by the following procedure:
(i) We sample x; from the uniform distribution on Sd-1 and y; ~ NO0 1) for all E [N]; (ii) We random sample i* E [N] by uniform distribution and set X = X;- with the I-NN label being y
(iii) If [|x, XN+1 < 8, set -xj.
C Proof for Theorem
In this section, we elaborate on the proof of Theorem Our proof can be broken down into the following steps:
ii) With induction, we prove that the evolution dynamics of w under gradient descent can be captured by two-parameter dynamic system, parametrized by Et and et:
13

---

(ii) By estimating the update dynamics for E and E2, we prove that with proper initialization parameter o, will have Ef. S = (log k), with ei s c E for a constant E (0, 1).
(iii) With the non non-asymptotic behavior of &1 and &2 determined, we further control the loss function L L(W) 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 respectively.
C.1 Dynamic of Gradient Descent
In this section, we prove that can characterize the evolution of W under gradient descent with two-parameter dynamic system. We proceed by mathematical induction:
We prove that when W a diagonal matrix with [
[4] (d] = C₁ for some constant
Ck. have the following break downs of our proof:
W
w = diag ( AC AEt - -.4))
(C.1)
times
where A&T and ACl are two positive constants.
(2) By W being a diagonal matrix with [W*1aaaa = Ckla, combined with Eq. prove that W is diagonal matrix and [W [d]][d Cc+1la.
CC. we
Now since (1) naturally satisfied by the initialization Assumption 2 for k 0, we can conclude the proof by simply proving (1) for & 2 1.
Proof for Step (1) of Induction. Recall that our loss function can be written
L(W) = E (x. hix.ehhen
(Xw((N+)) f(xN+1: (x. YchEEN))
==1 exp (x] WIIX++1 y) W XN+I + +X WI3 + y) w)W,))
j=1 exp (x] WIIXN+1 + y) W +x, TWs + y) Wj) -y. )
E (EINN)NN
where
WI
W
W W1
W
W
21
W W22
22
W 23
W
31
W 32
W 33
W E Llxd W W23, Wa with Wn E Rdxd Wa1 E Ri Wa E 1 W E Rdx1
and Was E IR, and we make the additional definition of y) y,, w Wz,. w W for
j E [N] and and y/++1 = 1, w W W3.. w WS1. Throughout the rest of this paper, we will
also adopt the notation of
exp (x) WIIXN41 + y, W Wfx++1 + X, T Wis +y) wW!
(x, W)
(C.2)
Σ (x] W WIIXN+I
W XN+1 +X TW 13 + +y W]]
when there is no ambiguity, where q W x {xi ( x [yTheeNN- As we have discussed in Section 2.2 W.z does not affect the outcome of the transformer, therefore the second column will remain 0 throughout the training procedure. Now will calculate the gradient Vw,LL(W") respectively. Note that for
Eg (x: = exp((g (I)), we have
Σ, ()) Σ, (x)
Σ3 (x( yi
[ Gol (x)y; Voga (x)(()V E
V
ye
=
Σ, (x)
v:.)
(x( Σx (x
E, (x)) }
14

---

therefore, with some algebra, we have the following closed-form formula for Vw., L(W) respec- tively.
N+1
N+1
N+1
(w,,L(W- = E
Σ
(x
W)(x)X - 1
Σ j=1
Σ
(x WiNs} }
j=1
N+1
Σ
qx(x W) -8-1)
(C.3)
j=1
Vw.:(W) E {{
N+1 (1 W) Σ
q (x, W)(r,&(+1)(5)) - 11 Σ
(x W)
1=1
j=1
N+1
Σ (x W)(r; -x.. -x.)) ||
(C.4)
N+1
N+1
N+1 (x wobn)
Vwi,L(W) = E
Σ (x, W) - { Σ
dx(x, (Xx) j=1
==1
N+1
(x, W) - y.. -y x.
Σ
(C.5)
j=1
N+1 (x, W)(ys)ys - Σ (x, Whys Σ (x worsil
N+1
Vw2e 4(W)= E
Σ
(xxx Wi(ss
(C.6)
Vw,, I(W) B {{- -
N+1
(x, W)(r)) (W)x )
N+1
}1
Σ
Σ (x W)(y, 1 y.-)
/
==1
(C.7)
and
(-(((22) N+1 (x, Win)) (W+)(W) {
N+1 E (()s
Vw,,L(W)= =
(C.8)
By the induction assumption, we have
(x,,, W*) = 1I explet xN+1)) + 1==
explet ||N+1 5
E exp(Et (x;, xN++)) + |xxx (2)
1 Mi*N+1exp(e" (x,, x)) + + - +
(C.9)
EN exp(Ef (x;, x)) + exp(()+ - -
where the first equality comes from ||x ll2 = 1. Therefore, under our induction qj is only a function of [xihe[N+1) and W independent to {yshe|M] Now, with the closed form of Ww,, L(W), we can make the following calculation. First, we calculate Lwi1L(W*):
N+1
N+1
N+1
Tw.,L(W' = E
Σ
(x W+) ( - yx) - { Σ (x, (xxxxX HH Σ (x, (0)s))->)))
N+1 (x w -v.) -v,))
Σ
15

---

E
N+1 Σ
(x,
(D) - yi) {
N+1
Σ ai(x, (y,,
(1)
N+1 (x
N+1
E
Σ j=l
W*xxxxx+1
Σ (x (y)
since q(x, Wh) is only function of by our induction assumption, we have
{ Σ (x,
N+1
N+1
}
()) (1) = B
(x)
Σ (x w
(1)
N+1
N+1
(x + H
N+1 L 9x w
E
Σ
(x,
(x 18,
ys-
Σ (x
(i))
(61)
N+1
+ E
Σ
(x, (
1=1
(iv)
here
N+1
N+1
E
Σ
(x W
(xxxxx()++))
ys Σ (x, W
p'=1
N
Σ (x, w (x
= E
N+1
(ii = E
Σ (x w
(X,xX++1 ) y,yr X1:N+1
N+1
E
Σ 9o (x,xx,
1 (y=argmax,e// Elys,y, [
='=1
E (x,
(x, W Wx)xx
N+1
(iii) = -E
Σ
q (x. w*) W (x
=1
N+1
EE
Σ
q
(x.
Wh)gr- (x, (xXxx+11)))
==1
N+1
(iv) = E
Σ
(x. (x
1=1
16

---

also for (2), we have
N+11
N+1
( - s..) } XLIN+1
(2) E
(x, W xx
E3
Σ a
N+1
N+11
(x)a W$E (xSN - (y3 +yy)y +y | |
= -E
(x, W* x++
Σ
(x
8,'==1
N+1
N+1
ai(x, (r wO (x) yy) [
-E
(x W
|
1+
Σ
(x, w*) - Σ
j=1
1,1'=1
-E {
Σ (x,
W
1
1+ Σ (xxx, (x, w
- -2.. (x, wh) }}
(1)+(2) gives us
N
N+1
Vw.. =E E
Σ
(xx (x, w (x,x(xx))1))
- BB 4x (x*
+ B q= +B W Σ
cx (x)xx)))))
N
N+1
E
Σ
(x, wh
Σ (xx (x, (x)
(C.10)
For W.,, (.)). we have
N
N
N+1 Σ (x WoNrs j=1
Vwn I(WA) = E
Σ
(x, ( -
(x (
W
==1
N+1 Σ (x, W -v.)
note that it is the expectation of multiplication of three linear functions of y:, therefore Lw,,L(Wa. equals to 0 by symmetry. For Ww,, L(Ws), we have
Vw, L(W() k) E {
N+1
(x, Wx)()s)) IN+1(x, We)x
}
{ Σ
N+1
(Ww)(ss
0,
here the second equality comes from (x, W&) being an even function of x and x has symmetric distribution. For V Iw.,L(WW,, k),, we have
{
N+1 Σ
N+1
N+1
Vw..L(WA = E
(W 1
Σ
(x)))
Σ ((WKQ)(r)))
==1
N+1
Σ
( - -x.)
0
here the second equality comes from (x, Wx) function of ||x|| and x has symmetric distribution. For ((W.L(Ws)) k), we have
N
N+1
Vw ((W)) =
Σ q (WKO)(y))))
q
1y;
Σ
q kQ (x5))}
-=1
17

---

N+1
Σ
(WWR)) - -oo.))
0,
due to its symmetry in y. For ss (W&), we have the following calculation:
V W.wLL(Wa) = (-
N+1
Σ 9 WW)x,
a Wt)
I
{
Eg cx Wa)(y; y.-)
N+1
}
}
E 9N+1 (a) (E E
(x)()
Σ 9 Welys)
E 9N+1(x Wo) Wx)
> Cx Wo)
(C.11)
With all previous calculations, we know that Vw. L L(W)) equals 0 except for Ww,, Vw.,L(W') and Ww with our induction assumption.
Now we only need to prove that V,,, L(W) a diagonal matrix. Recall that (E)N)))+1) is identically uniformly distributed on gd-1, thus it naturally satisfies the following Assumption:
Assumption 4 (Rotational Invariance). We say the distribution of x is rotationally invariance, if x Pz, and for every U E Od), where the group of orthogonal matrices, we have Pa = U Pa.
straightforward lemma under Assumption Dis the following.
Lemma 6. Suppose (xi)<E(++] are sampled i..dd from P and Assumption 4 holds, then have 8 -=l
Px.
2 i=1 Pux, for all U E (d).
Proof. Since x,- are sampled i.i.d. from Pz, we have
N+1
N+1
P
Puxi
1=
Lemma 7. When ~ Pa. Pz satisfies Assumption and W is diagonal matrix with L(W*)] = . Id.
[W*] [d][[]] [W*]4x = TT we have E[Vw,,L(W"))
Proof. Note that by the induction assumption, we have
(x, W = 1IAN+1 exp(et #
+ 1 exp(st - +
EN exp(Et (x, xN+)) + exp(st <
therefore, (x W (Ux, Wh) for all U E (d). By Eq. (C.10). there exits set of functions
(of()))eeM) U (of (x)) R RR
such that
V w L(W)
Σ oft xxx]x=+1) xxx+1 + g (x].xN+1) X
here the expectation taken with respect to {x, (xi By Lemma we have Px, @ Px, Pux; 6 Pux, for any orthogonal matrix U E O(d), therefore
Elg& (x] xN+1) x = EEgg ((Ux;) ((UxN+1)) (Ux))(UXNHI]])
UB[gg (x]xx XI- XUT+
18

---

here the first inequality comes from Px, Pxx, = Pux, &ux,, Such identity holds for all orthogonal matrix U, therefore Big (x: xw-1) xxN++1] must be multiplication of the identity matrix la- Now, note that for all orthogonal matrix U,
X argmin || - x
Ux.- = U argmin ||xx; - UXN+1122, X
therefore
V(X/:X++|/XNNNNI X UTIUXNNN
Since P(XI-N+1) = P(UXL:N+1). we have P(x,IXx41) = P(UX, U,,, therefore
Elg (x11xN+1) X x = Elg.. ((Ux)) " (UxN+1)) (Ux, (UxN+1) = -UE (x.1x X, X,+ I
Again, since U is an arbitrary orthogonal matrix, we conclude that Elg (x. xN+1) x is multiplicatiin of the identity matrix. Summing everything together, we conclude the proof that V Ww; L(W*) is multiplication of an identity matrix.
With Lemma
Eq. 2.6 and the previous calculations, conclude induction.
C.2 Evolution of { and &
In this section, characterize the training dynamics of & and & under gradient descent. Our results can be broken down into the following steps:
We prove that when _ and 22 are initialized as in Assumption 3. with N, d satisfy the
condition in Theorem we have ş - 2 O.
(2) We provide an upper bound and lower bound for g\+1 st and ch+11 $
(3) We prove that there exits constant e₁ E (0, 1) such that ei E (0)))
(4) We prove that < 1 & = Qn1 M11 log(poly(N, d)) log(k)).
d) ( we thus conclude that $
Combining (3) and (4), we conclude this section by proving that et, $ and E gt are both of scale 2llog k).
Step (1): Initial incrementation of gl. First, we prove that 5l - -P 2 0. Note that by Lemma 7
and Eq. (C.10), we have
d
( 7
{1 = Ww, I(W9)
Σ (x, (x) xN+1)
-Bax q:- qxx (x3xxx1)1))
E
(x, (xx] (x][(xxxx+)) [mn=1)]
E
qi- q--(x, wo)
Σ
N
N+1
Σ (x (x [x]xxx])+
E
(x, w
==1
Elx
x +
(N 1)3 E|xx J,
( 2 0)
19

---

To prove that EP -1
El < O, we only need
E(xx 1 s 0. v + (N + 1)3
N +
However, by Lemma 18 this can be guaranteed by N 2 O(vdlogd) d).
Step (2): Scale of gf+1 gf and gh+1 E have
Note that in every gradient update with stepsize 7), we
d
-4)) = (gh+1 7
q. (x, ()
E Σ (x, (x, (V(x]xxx.1)
N+1 E fa- W) - Σ (x W)}
Σ (x, W)(x3xxx1))
(C.12)
= E
(a-- (x, W)
Σ (x, w)l ( x -
N
N+1
Σ
(x W)(x3xxHH1)
j=1
1=1
N +
Σ
gي (x, (x - x] -
(C.13)
j=1
we also have the following estimation for the update of Ea:
N
(52
< = Ex(xN+ixx, (h) (x, W 1x w
j=l
N
2 E[qN-I(x, Wh) {q- (x, W 9- W Σ (x, W
3=1
=E|q- (x, W q (x, W
(C.14)
which shows that in each iteration, Ea will at least increment by scale of n Elq- (x, W*). Px41(E (x, W Combining Eq. (C.12). (C.II), we have the following estimation:
Lemma 8. We have
1 [d(sl - ct) + +2 - + 2 (1- 2N Ca Caexp(-6Et
for all k 220 O.
Proof. First, note that we have the following relations:
-. [d(ettt - 5t) + +(E {t0}}
N
qi (x, (x (xx+))
E
Σ
(x, ()(x] xx+1i]
N
+ (x (x. wh)
Σ
of (x, w
j=1
-B (q. -B (x, w) w) E (x, wi)
Σ (x, W) [[(x]xx=1)]
1=1
20

---

- of w)} (xTxNNI++1) (ITTOSXAAA - Σ (x W)(x]xx=1)]
NV
{q (x, w) -
5=1
N
E
Σ
cz(x, W) -XIXXXIIIXXXXIII -xfxN+1I
N
- ()((x w)} qx+i(x, W) + +1]
B
(a- W) -
1=1
N
+ E
Σ
(x W) ()xxxXxN+11 -x]x++11
(C.15)
here the inequality comes from x, XN+1 2 x] XN+1 for all j E [N]. Moreover, note that when { 2 0, we have
N Σ
(x, W) + qx+ W) & Nq, W) + qN41(x, W),
j=1
which is equivalent to q: (x. W) 2 1x Since
(x, W)
QN+1 (x, W) = exp (x - 1) + Ea),
we have
1 (x, W) 2 N +exp(E1(1 - x(xxN+1) Ex) -x)
(C.16)
Therefore we have
E
q Xx,
W)Xx, X0+11XN+11XX3 -x,+11IxxxxI-- x|xx+1I
(x, W
E
(x,
fxxTXN+1- xJ
-x.ww11xxxxx -x]]xx=x x
==1 (x, Wk)
2 N
XN+1 x][xN=1) ) (x]xxxI1-x] -x]xxxx)
>E
Σ exp
2 [x]xx=1
N + exp((1(1 - x(xxx+1)= Es)
1=1
exp(-48t)
Σ (xfxxN+1 x, xxx+1|IxxxxN+ 2 P(x
(N + exp(Ef )²
1
Ca exp( (-4t)
1
2N
(N exp(Et 12
1
1
Caexp((66f)
(C.17)
2N
here Ca is constant that only pertains d. The first inequality comes from Eq. (C.16). and the third Finally, note that
inequality comes from Lemma I7
E (qr (x, W)
Σ (x, W)} 9N+1(x W)xx 1) 2 0,
and we conclude the proof.
Next, we provide the following upper bound for the increment of EL and & whenever &t 2 0.
21

---

Lemma 9. When ef 2 0 and N 2 O(Vd d) as in Theorem
we have the following inequalities:
- (gk+1 ) -
2N
exp(
4
(1
2t
and
exp(2(c) 12
d
(N + 112
d
where and = (2N नेअ and
(52 1
1
cc++1
$ S exp(2 - -E2
Proof.
We first prove the first inequality. Recall that
d ((s <)- = E fai (x, w)
Σ
N w)} [0]3(xxx++++ Σ
(x,
N+1
(x W) [((xx]xx+11
j=1
N
q (x, w) [xx x - x3 x]xx==I]
E
j=1
when Ef 2 0, we have q:- (x, W)
Σnn (x, W) 2 0. therefore
d (5/tt1 Et)= = E
d
N
Σ o (x, W) XN+1 - x
77
i
N Σ (xxx, w (1 -+)(1-x[Xxx111 x(TxX111
E
qi q-- (x, w (x))
1=1
(ii)
N
N
+ E (xx (x, W)
Σ
q (x, w)
L (x, Wx - x] x]xN+1)
j=1
j=1
We have the following bound for (iii) when gt 2
1
E (a-- wh) - Σ ix why Jannils w w*) xx,xx 2 21 an.d
N
(2N
2
1
(x. w ) (x, wh)
(2N
23
1
1
A
B[(1 (x, ( W
VaN
N
1
1
1
exp(2(gt -
N3 (2Nv
where the first inequality comes from Lemma we have
19 Now we aim to bound (i) and (ii) separately. First,
(x w*)
dh (x.
W
[x,xN+11 x3 xN+1
o (x,
W4)
j=1
N
E
exp
-2eh (x,T x x3 (xxx=I))(xxxxxx x]xx-1}
22

---

S N exp
(/x N xx+
-x - XN+1
- x
4 (N + 1) c
S N exp
(C.18)
here the first incquality comes from Jensen's 's inequality and the convexity of a exp((xx between [0, 2]. and the second one comes from E[x;-XN+1l 2 T which is guaranteed by the condition of N in Theorem | and Lemma [8 Thus we conclude the proof for the first inequality. For (ii), denote the second largest order statisties in {x, XN+1 HeeNN by X22;XN+1, we have the following inequalities:
{ q: (x, W) - -Z. (x, W)} {
Σ (x, () [[]]]]xxxIIxxxxM [xx+i)
(ii)
S
E
N
4, (x, W
E (1 qi (x, wb)) Σ
(x.- xx+1N++1 - xxxx+11xxxxxi x]xx+1)
qi- (x, Wh)
< E min
exp ( - E -xXxxxx+xxxxxxxx
- xx]x+++])) (x - x3 x]xx=13]
I N exxp
(Bl <t x XN+1 - x3 xx++ X| xN+1l
s N cxxp(-=*1>
where we define a random variable Çk by
(N 1) exp (5f (x x xxxx+1)) + eexp(E* (1- (0-xXXxN+1) - -22)
the second inequality comes from
exp(st xxIxN+11
1 - d. (x, W I 1 exp(Et - $ + (N 1) exp(et xN++) + exp(Ef x..xx+1))
s - exp (st (x X,- xN+1 + exp(st {1 x
-
and the third inequality comes from Jensen's 's inequality. Combining We now aim to prove the second inequality. Recall that by Eq. (C.1I), we have
N 1 (ck+1 + = Elgn+i(x (xx W - Σ (x Wb)
j=1
therefore have
1
ck+1 52
< < ElqN+1(x, (-- Wh) -x-(x, (x. Wb)
r7
AN+1 W+) min(1,Ct
E
q-- (x, W4)
1
E
1 + exp x 1) + +
< E(exp ((1 xTxxx+11)tt -
exp( E--
Thus conclude the proof of both inequalities.
The next lemma is direct combination of Lemma 8 and Lemma
23

---

Lemma 10. When Et 2 0, we have
d (5t
<
2
1 2
Ca
Caexp(-6E)) 2exp(25f -252)
The next lemma provides lower bound for the update of 33. Notably, when Ef is close to E 4 increases in a larger scale.
Lemma 11. When E 2 O, we have the following inequality holds:
(52 -
ck+t +2)
(N 1) exp (2un,d Et -222)
1
where andd is a constant only pertains to and d that is defined in Lemmal
Proof. With Eq.
C.I1I
we have
1
k++
2
Elq- (x, Wh) I W
11
gi++ (x, W q (x, Wk)
qh- xx W
qi- (x, w*) exp (25 (1 x xx+1) 2(/ E
E
q. (x, w²) exp (2an,d Et -252) |x,_ XN+1 and
(N + 1)3e exp (2an,d ef 2gh
where the first inequality comes from Eq. second inequality comes from Lemma 19 the third inequality comes from Et 0.
Step (3): Upper Bound for c/ /c2. Now, combining Lemma 9 and 11 we immediately get the following result:
Lemma 12. Ifa = <b 2 3 log( and = (moxx(lo(MM)) Nd)), - log (1 (NVZ))))) and g 2
for all k 2 0, then we will have el <
7 ck 1552-
Proof. By Lemma 9 we know that whenever et 2 O, we have
(Et+1 <t) E -tp(-
2N
and dN3e
exp(2(s) +b)).
therefore, if
2N
d exp( (xp(-5t/2)
and
exp(2(5) - =2))
dN3s
which is equivalent
an,d 2.5 log(
+ $
2N dd
we will have ş ck+1 - Et < 0. By Lemma not increase as long as
11
we have th
c 2 3 log and
Therefore et will
2N42
15
ck ck
7
$1
$2
finally, recall that
an
12N ka) d-3
and our result follows.
24

---

Step (4): Scale of č and E- Finally, we conclude the convergence of gradient descent with the following two Lemma:
Lemma 13. With a and N satisfying the condition in Theorem we have
$ = Q2(1 pplypllllNNN d) log k),
with E - 1ek holds all k 2 O.
Proof. We first establish convergence rate for ε2. Note that by Lemma] 12 we have
Lemma
11 and Lemma
1
ck+1 $2 3 S exp( (
15 $2
and
ck c 2 " (N + 13e)
+ 0 poly(N, d) log k)) for all le 2 Now turn to holds for all k > O. Therefore, we have ς By Lemma
when
Ca - T.
g88 < - - log
we have
ck+1
ck
1- Caexp(-66f) 2 0. d
2N
Ca Ca( 1 2v
Since < 0 2 2log
E monotonically increasing, and E 2 0, by induction, we
have ş 2 0 for all k; and et 2 (poly([N N, d) log k) until et 2 & g Meanwhile, Lemma 12 shows that
2N cxp( cxp(-<t122 /2)
ct+l - şt z d
for all & 2 O, which implies ę - Oppoly N, d) log k). Therefore, et = ((poly(N, d) logk). The results of c 1 follows by Lemma and < 2
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 = Ta Ta. W33 = -E2, with E1 2 0, and the rest of items are all zero matrices, we have
N+1 Σ
2
No
qxx Wyy, y2. 1=1
0
+ cxpp(2f1 - Es)
&
Proof. We have
N
)9=1
N - 2E qr (x, W)]
E
q(x,
W)y; -yr
Σ
9 (x, W)
El1 q (x, W)]
Σ
q x, W) g. (x, w)
note that
N
q (x, W) - q- (x,
(x, w)
E
q- (x, W) W) - - (x, W))
j=1
E
4N+1 (W, W)]
25

---

$0.
(E 1 q(x, w)y, y:. )
3
Therefore 18
< - -
q-- (x, W)]. However, we have
E (x, W)] I
E
9,(2, W)]
jFi JE(N+1]
4 W)
= R2 q:- (x, W)
-:- (x. W)
1+1 NE[++]
exp (51 {x XN+1 x x(xx+13) + Eleexp( - -ss] < (N 1)B
(qr 1
$ (N 1)
exp (51 -x+xx++]])++)) +exp(251 - Ess),
By Lemma 20
we have
q(x, w)y, yyx. ) 50 N
E
Σ
exxp(251 -Ea),
and we conclude the proof.
Now, combine Lemma 15 with dynamic bounds for ei and E developed in Section C.2 we have the following convergence result for the loss function.
Lemma 15. When W defined in Eq. (2.4) is updated by gradient descent Eq. 2.6 we have
N+
w)y, y.- )
< <0 poly(NV, d) log K
E
Σ
q(x,
Proof. By Lemma 15
we have
N+1
-v: ) <0
Σ
q(x,
why,, y.
+ exp(2st -
5t
By Lemma 13 have e I
7 ck 1352.
with Ef = I(ppoly[NV d) k). Thus have
N+1
q(x, W w*)y,-y' -y. )
No
1
Σ
<
expl
15
1=1
d)
1
+
log k
1/15
poly(N, d)
log k
D Proof for Theorem 2 and Corollary I
In this section, we discuss the behavior of the pretrained transformer in tasks under distribution shift, and provide proof for Thcorem 2 and Corollary I in Section 3.2
26

---

D.1 Proof for Theorem 2
In this section, we provide proof for Theorem 2 The result comes from the following observation. First, we condition our analysis on the Ag. i.e. assuming that there exists constant such that
||3 - XN+1 d | x
Vj + i",
(D.1)
Then we have
[7(W*)
yr
R
Σ
q W -
1=1
Σ
li(x, W yx.) + +RlaN+1 (x, W
JE(N
-2R
Σ
g,(x, W + R (x+ W
JE(N
(x, W + R 9N+1(x, W
( W
< 2R
Σ
q- W
JE(
22
Σ
exp (SF xx] XN+1 xxxN+1)) + exp (E1 (1 xxIxx-1) -(5)
1
2R eex((=5f -8) R exp
--15
2RN exp poly (N,
log K)
=0 RoNK-oo
(D.2)
uniformly for all ((x,, [E) and XN+1 on Ag. Here the second inequality comes from Theorem I in which we show that 5/ I 7 ck Recall that we defined {,(x, W in Eq. (C.2) as the weight for the j-th label in {ys])e(N+]]- Next, we have
1552
Eprcx (g(W*)) yx-)
Epece ((F(W)) y:-)² [A, + Epec ((((W*) - yz-) il Mag]
s R2 (Poyykpo 4R22 plest (A3).
(D.3)
where the inequality comes from Eq. D.2). Here the expectation is taken over the testing distribution plest By taking inferior on the inequality above for all > O, we conclude our proof for Theorem
D.2 Proof for Corollary
In this section, we provide the proof for Corollary The first statement of Corollary T comes from Markov's s inequality:
plett Round (fw(xN+1)) # y:.) peest
(WW(X)+1))
- y:-
< 4 Blll -y. 12.121
< <OMN inf (K-poly Poly( + ples (A8)))
$20
where the last inequality comes from Theorem
2
Next, we prove the second statement, Suppose
prest (Ag-) for some 5* Then similar the argument in Eq. (.2) we have
((W*) - yv-
1 s PolN
27

---

holds for all {(x) U [xN+1] ~o prest almost surely. Note that
Round (fw(xN+1))
holds whenever
((WW)) yo-l 2"
2'
therefore, suffices to have 1/4 H O POlKK--)))N ), which is equivalent
log(MNN K = (olly(N,d)6*
E Nonconvexity of Loss Function
In this section, we show that the loss function is defined by Eq. (255). We prove by a special subspace of W E (dd+2)x((++)) defined by two scalars E1. E2.
W diag fi.ggl........................ E, 0,
(E.1)
By showing LW is nonconvex under such parametrization, we conclude our proof.
Lemma (Nonconvexity of Transformer Optimization), When W is a two-dimensional subspace of g(d+2)x(d++ defined by Eq. (E.I). the original loss function defined in Eq. (2.3) degenerates to
the following:
Li=1 explEi (x:, XN+1)) + exp(Ei E2) y:" )
Σn (x,,
E2(1, := E
where we use YN+1 = O. Such loss function is still, in general, nonconvex.
Proof. The degeneracy of the original Eq. 2.5 to L(E1, E2) can be shown by basic algebra. We only need show the nonconvexity of T(E1,E> {2) in our proof. Note that by Assumption the gradient of E2) is defined by
E
- yr- )
0g, L(0,E2) = 0c
+ exp( -E
Ej' Vjyr-
(2) y2)²
= 0cg F.
1
N + expp(-o>
(N + exxppp-20 +2
N
2
(N + exp( -E))
N expppp- -2)
- exp(- (2{2)
(E.2)
(N + expl (xp(--2)) 3
where the third equation comes from
|
E
y
IE E
y
{xi
N
EE
i=
ysY:
(x:
1=1
E
Σ j==--
y
=1
28

---

and B{(_== -= v.) E v N.
N. Eq. E2.2 shows that when lim 222 0
and --- 2e, (1, 62) 0. If (2) is convex, then (0, (2) is monotonically increasing. which means 0g ((0,5)) = 0 for all &z. However, this is clearly contradiction.
F
Auxiliary Lemma
Lemma 17 (Distribution of Sphere Inner Product). With the assumption of x; sampled from a uniform distribution on the sphere, Let be the cosine of the angle between an arbitrary -dimensional vector and vector chosen uniformly at random from the unit sphere. Then the probability density
function of random variable T E [-1, is f+(t) ) (1 -22 3
20(
Proof. For convenience, we assume that XN+1 e1. Note that this does not change the distribution of XN+1 x; due rotation invariance. Let X; (0, 1). define
X
Yi X1, Ya-1 Xa-
1 X²
Note that distributed the same way as x,. Calculating the Jacobian:
0
0
0
1
Xa
X; Xa
Ed X²
Ei X?]
ΣN x?|3/2
Ei X? 972
Since is of the form:
I 0
J
a b
the determinant easily evaluated:
Σd
X?
J
3/2
Σ ==
x2
Now solve for the distribution of YN.
10
1
(v)
|-
[ fx(x)xx
where
1
Et.
[x()
E
(2 d/22
29

---

d-1 Writing the distribution in terms of t:
3/22
Yj
Ei=1
i=
Vi
1-12
1z Eist dy .y-------
J (1)
Σ
Li=1
(2Th 2/)a0/2"
d-1
21-12 Eizi ve vidy dy1 dyd-1
fr (t)
yi
8/3322
(2x) (2x)d/2
f+
1
Σ
yff 1/2
(1- (2)312
2m N/2
1
Now can make substitution to remove ya from inside the integral. Let un (1-89) -1/2 we have
yn.
(1 -2)32 (2T)d/2 (1 (1- - 22 22)/2
d/2
600
d-1
f+ (t)
u
E
u dui ... dudd11
-00
i=1
and
1+ = (1 (1-)))(1-0)
(2x)-af2 4/2
100
100
Σ
d-1
u e
u2 dw
dud-
The integral can be seen to be constant so:
d-3
fr (t) = kd (1 - -to) x
or for notational convenience:
f+(t) = kid (1 - p2) 2=3
Since (v) is PDF and is defined for 0 < t 1:
[ ka (1 - (2)
4-a 2
= 1, where kis
f/ (1 y2)
=
dy
2T(2)
Furthermore, for d > 1, we have kid
VaI(
=))
Lemma 18. With N O(logd d- vd). we have E[x][xx+11 (
Proof: Note that we have
B[x 2 P( max x, X a) a {
IP( max x, x|xN+1 20x 201 (-1). iE[N]
KE(N)
To ensure that F [ xN+1] 2 (N+1)2+ we only need
adI( max x XN+1 2 a)) - {1 P
max x, XN-1 2a)
2 (N + 1)2
NE[N]
kE(N!
which is equivalent
( 1+
2
P( max x; XN41 a0) 2 I+a [E[N]
(N + 1)3
30

---

By Lemma 17
we have
kid 0 (1 12) d=3 N
P( max x; XN+1 0x) =1 vE[N]
dt
Therefore, we only need
hi S (1-
-2 (1-
dt
N
= =P( max x; XN+1 2 2a) 2 iE[N]
1+
2
(N - 1)3
I+a
which suffices by
log
1 (1+ [N+1)² T+a
N
log(1
ka Sal (1 12) a-3
" di)
Note that
[
1 (1 (1- +1-t)d-sdt
1
(1 22 dt 2
d-2 (1 -1-11-))-7 a)-2)d
therefore we only need
log(
I+a ([+o][N+I]²
N 2
Ng d-2
(1 a)d-2
Now, let a 2-2- since ka
r(.) r(
O(Vd), we only need N 2 0 ova log(d - 2))
O(vdlog d).
Lemma 19 (Concentration upper bound for x(-xN+1). When x,.XN+1 is defined by
maxie[N] (x] XN+1h. where x; are independently sampled from a uniform sphere distribution,
we have
P XIINN+I
2
(2Nka) +
Proof. By Lemma 17
Ax.XN++ 5a)= = (1 (1- kid [ (1 (1- 2=3 dt
N
Note that
[ (1- (1- d-3 dt S (1 a1 a²) d-3
since (1- ) is monotonically increasing, we only need ku((1 a1 -= s N. Setting
1 (2Nka) S <(1 (N(Kaa)
1/2
suffices.
Lemma 20. Suppose {xitegN+1) are i.i.d. samples from uniform distribution sphere in IRd, with x].xN+1 and xN+x being the largest and second largest order statisfies among {x] XN+1hee(]] respectively. Then we have
N2k g
exp (E(xTe)xN41 - xT. xxxx+1)) S 50(
Moreover, we have
E
exp (E (x))x
-x - xxxxx=x)))) =
where 91-)) hides constant depends on N and d.
31

---

Proof. We denote x; XN+1 by Y,, XN+1 by Y, and (2) XN+1 by ((2)- By Lemma 17 we have
the density function of Y being fa(t) = ka (1 - t²) d-3 where his = 2
) Then by the joint
distribution of order statistics David and Nagaraja 2004 the joint density function of Y,- and ((2)
is
fv (v), v2) = NNN 1) fa(v2)sa(v1) Fd F (12)
for 1 S y2 S y1 S 1. Therefore, we have
E exp (E(x1))xN+1 - x(xx++1))
NNN - 1)
["
(1 v2) == vì) exp (E(82 ys dyydyy
N(N 1)
k2
(2 vi v exp (E(v2 yn un))da dyzdyr
2
N(N 1)
exp(Ev2)dy2dy)
S N2
exp(Ey2)dyz
N2R
E
Thus obtain the upper bound. Next, establish lower bound.
exp ((xx)))XN+ - x(Txx+1))]
N(N 1) k3 / (1- 1/) (1- SS exp (5(12 v1)) dyddy
n
2N
N(N - 1)
(/ (1 v²) exp(-Ev1))
k2
2
N(N 1)
2/v0
k2
(1 -expp-
- E/ vdj)dy
22NN
1/va
N(N 1-311
25 dd
32

---

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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations 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 encouraged to create 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 few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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
33

---

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 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 cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix 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 ex- perimental results of the paper to the extent that 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 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 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 a specific model and empirical evaluation, may be necessary to either make possible for others to replicate the model with the same dataset, 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 large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions 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) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (c.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 deseribe 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
34

---

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 instruc- tions 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 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/
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 new open-source benchmark).
The instructions should contain exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. httpppscccpbblcccuidcccee
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 the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) 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, hyper- parameters, 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 include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary 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.
35

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence 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 given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation 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 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 range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: All experiments are conducted on CPU cluster
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well 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 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration 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]
36

---

Our work discusses the theoretical performance of a well-known architecture, thus the social impacts are insignificant.
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deploymeett 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 other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how 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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make 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 existing assets.
37

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include 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, paperswi thcode. com datastt has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 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 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. Crowdsoureing 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 involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, 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]
38

---

Justification:
Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing 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 adhere 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 the institution conducting the review.
39


Reviews:
[{"review_id": "hrjwO5IbgE", "reviewer": "Reviewer_cdZv", "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.", "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. \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.", "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.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed comments and suggestions. In the following, we will try our best to address your concerns. To accommodate the extensive character count in equations, we will provide our response in multiple parts.\n\n>**Q1**: In-context learning with attention, seems to do a form of non-parametric learning at test time. It would be great for the authors to contextualize the work a bit more in terms of understanding transformers as language models, more generally. \n\n**A1**: Thank you for your advice! We confirm with you that our paper follows the framework of a line of theoretical works in studying In-Context Learning ([1][2][3][4]). Our primary aim is to answer the question of **what type of statistical algorithms could an attention layer approximate** under regular first-order optimization methods. A direct application in the language model would be text categorization, in which the input is a sequence of words/phrases and their corresponding labels, while the query is another word/phrase in which the model aims to predict its label. We promise to add this explanation in the introduction of our new version. \n\n>**Q2**: Can the analysis be extended in any trivial way to k-NN (e.g. using multiple heads)? \n\n**A2**: Thanks for bringing this to our attention! We believe the extension of our results to k-NN is non-trivial, and cannot be easily achieved with one single attention layer even with multiple heads, as the softmax operator only allows an easy approximation for choosing the closest/most distant $x_i$ in the context with $W_{11}$ goes to $\\infty$, but cannot directly approximate tokens selectors of other orders. We would like to leave the interesting question of learning k-NN with multi-head multi-layer transformers to future works. \n\n>**Q3**: Is the independence between $x_i$ and $y_i$ being used throughout, or not. \n\n**A3**: We make the assumption of $x_i$ and $y_i$ being independent for the purpose of ensuring that the attention layer does gain its prediction power by learning the 1-NN algorithm: the prediction power of the model must come from utilizing the proper comparison between query and context, as any estimator in the form of $\\hat{f}(x_{1:N+1})$ or $\\hat{f}(x_{1:N}, y_{1:N})$ can achieve at most 50% accuracy. \n\nFrom a technical perspective, independence condition allows us to achieve a more delicate characterization of gradient descent, which is helpful in achieving the final results.\n\n>**Q4** Minor questions, including typos and writing unclarity\n\n**A4** We highly appreciate your efforts in pointing out those questions and have already revised these issues in our current manuscript.\n\n[1] Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality, Siyu Chen and Heejune Sheen and Tianhao Wang and Zhuoran Yang\n\n[2] In-Context Convergence of Transformers, Yu Huang and Yuan Cheng and Yingbin Liang\n\n[3] Trained Transformers Learn Linear Models In-Context, Ruiqi Zhang and Spencer Frei and Peter L. Bartlett\n\n[4] Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection, Yu Bai and Fan Chen and Huan Wang and Caiming Xiong and Song Mei", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5AN5OlCpBE", "reviewer": "Reviewer_cVwG", "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.\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.", "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).\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$.", "limitations": "This submission would benefit form a discussion on the generality of the technique (see the \"weaknesses\" section).", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "My 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$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed comments and suggestions. In the following, we will try our best to address your concerns. To accommodate the extensive character count in equations, we will provide our response in multiple parts.\n\n>**Q1**: 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\n**A1**: Thanks for bringing this to our attention! When the variance of the random initialization is small enough, with a standard perturbation analysis, by utilizing the Lipschitz condition of the loss function with respect to $W$, we can achieve a similar convergence result with $W_{33}$ being a big negative value. We will add this discussion to our work in the revision. \n\n>**Q2**: 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\n**A2**: We confirm with the reviewer that the extension to $k>1$ is nontrivial, and beyond the ability of one layer of attention even with multiple heads. The reason is that one single attention layer suffices to approximate the one-hot vector of the closest/most distant $x_i$ as $W_{11}$ goes to infinity, but cannot directly approximate other $x_i$ in the context. We would like to leave the interesting question of learning k-NN with multi-head multi-layer transformers to future works. \n\n>**Q3**: Minor questions, including typos and writing unclarity\n\n**A3**: We are grateful to the reviewers for pointing out those issues, and will clean up the notations in our future version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Z4a6Kl8upq", "reviewer": "Reviewer_FWXF", "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.", "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": "Please refer to the weaknesses section.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**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.", "weaknesses": "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. \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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed comments and suggestions. In the following, we will try our best to address your concerns. To accommodate the extensive character count in equations, we will provide our response in multiple parts.\n\n\n\n>**Q1**: Assuming that the data lies on a hypersphere and assuming no class imbalance seems rather strict to me. Further, it is far from practical. What are the insights?\n\n**A1**: Thank you for your advice! We assume no class imbalance and the hypersphere assumption for the training data when studying ICL under 1-NN data setting, as it is the cleanest and the most mathematically tractable setting. Particularly regarding the hypersphere assumption, we consider this setting because it allows us to transform a distance comparison problem to an inner-product comparison problem, making a single self-attention layer capable to compare distances between different tokens. We would also like to point out that this setting could be extended to other data distributions, such as high-dimensional spherical Gaussian distributions, which can be viewed as being almost uniformly distributed on a high-dimensional sphere. \n\n>**Q2**: How well does this generalize to relaxing conditions on the input data lying on the hypersphere under SGD and random initialization?\n\n**A2**: When the input data tokens have significantly different norms, in general, the model could suffer from a large approximation error. This is because the learning of one-nearest neighbor prediction rule by one-layer transformer relies on the fact that when all tokens have the name norm, the a distance comparison problem in one-nearest neighbor is equivalent to an inner-product comparison problem. However, We expect that our result can be extended to settings in which the data points are located around a sphere, such as high-dimensional Gaussian distribution. The same conclusion can also be derived when $W_11$ is randomly initialized with a small variance. Both conclusions can be achieved by utilizing concentration inequalities and a standard perturbation analysis.\n\n>**Q3** Multiple typos throughout the paper.\n\n**A3** We are grateful to the reviewers for pointing this out, and have updated our paper by conducting another grammar check.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Iro7ZgFI8K", "reviewer": "Reviewer_FpKW", "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.", "questions": "N/A", "limitations": "The authors have addressed limitations.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "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.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1**: Can the message of the paper be generalized beyond this specific training distribution or even kNN?\n\n**A1**: We confirm that the data lies on a uniform hypersphere is a key assumption for our analysis, as it allows us to transform a metric comparison problem to an inner-production comparison problem, the latter need not take the $\\ell^2$-norm of the tokens into consideration. This further enables a single softmax attention layer to approximate 1-NN algorithm, since softmax layer can naturally approximate an argmax operator, which is crucial to our final results. Data distribution beyond the hypersphere distribution could result in large function approximation errors and further hinder learning. As the first paper studying ICL in learning a nonparametric estimator, our motivation in making this assumption is to start from the cleanest and most tractable mathematical framework, while extending our results to more general data distribution is a highly important future direction. \n\nWe also agree with the reviewer that learning k-NN is a valuable topic in understanding ICL. However, the extension of our results to k-NN is non-trivial, and cannot be easily achieved with one single attention layer even with multiple heads, as the softmax operator only allows an easy approximation for choosing the closest/most distant $x_i$ in the context with $W_{11}$ goes to $\\infty$, but cannot directly approximate token selectors of other orders. We would like to leave the interesting question of learning k-NN with multi-head multi-layer transformers to future works. \n\n\n>**Q2**: The data assumption that the query is not generated in the same way as the label of the context examples is inconsistent with practice; \"Independence of $\\{x_i\\}$\n and  $\\{y_i\\}$ is essential to study in-context learning of 1NN” is incorrect: $y_i$ can depend on $x_i$ and predictors of both forms mentioned will achieve 50% accuracy.\n\n**A2**: We agree with the reviewer that $y_i$ can depend on $x_i$ and the predictors \nof the form $\\hat{f} = \\hat{f}(x_{1},\\ldots,x_{N+1})$ or  $\\hat{f} = \\hat{f}(x_1,y_1,x_2,y_2,\\ldots,x_N,y_N)$ can still only achieve $50\\%$ prediction accuracy. For example, if the data in the context follows a ground truth linear model, but the linear vector is randomly generated, then these two types of aforementioned predictors can still only achieve $50\\%$ prediction accuracy.\nHowever, we would like to clarify that the goal of this paper is to demonstrate that one-layer transformers can learn one-nearest neighbor in context. If more complicated dependencies exist in the data, it is easy to expect that the transformer model will not learn a clean one-nearest neighbor decision rule. Instead, the transformer may reasonably learn a prediction rule that is some type of “mixture” between in-context linear regression and in-context one-nearest neighbor. Therefore, while studying the case where $x_i$ and $y_i$ are dependent is an interesting future work direction, it may not serve as the best and cleanest example to study the learning of 1-NN decision rule. On the other hand, when assuming $x_i$ and $y_i$ are independent, we can clearly show that the one-layer transformer can cleanly learn to perform 1-NN classification, which motivates us to assume $x_i$ and $y_i$ being independent. Due to the same reason, we are also inclined to assume that the query in the training data also comes from 1-NN label in the context.\n\n>**Q3**: 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, thus it requires no knowledge of the data to achieve strong performance.\n\n**A3**: We confirm with the reviewer that distribution shift results come from the model “remembering” the 1-NN algorithm. However, our results also surprisingly show that when the $\\{x_i\\}_{i\\in[N]}$ in the testing data are strictly bounded away from the decision boundary, the prediction error could be even better than the training loss. This is also shown in our empirical results in Section 5.\n\n>**Q4**: From Lemma 4 and 5, it is not clear how $\\xi_1^k$ grows slower than $\\xi_2^k$, as is claimed and as is needed for the final result. For any $\\xi_1^k$, $\\xi_2^k$, the upper bound on $\\xi_1^{k+1} - \\xi_1^k$ can be dominated by $\\exp⁡(poly(𝑁,𝑑)\\xi_1^k)$, which can be much larger than the lower bound on $\\xi_2^{k+1} - \\xi_2^k$ of $\\exp⁡(−poly(𝑁,𝑑)\\xi_2^k)$ even when $\\xi_1^k=\\Omega(\\xi_2^k)$.\n\n**A4** We apologize for the typo in Lemma 4. The correct form of Lemma 4 should be $$ \\frac{d}{\\eta}(\\xi_1^{k+1} - \\xi_1^k) \\leq c_3\\cdot \\exp\\big(-poly(N,d)\\cdot \\xi_1^k\\big) - c_4\\cdot \\exp\\big(2\\cdot (\\xi_1^k - \\xi_2^k)\\big),$$ where $poly(N,d)$ is a positive polynomial-order term, as we proved in Lemma 9-10 in appendix. Therefore $\\xi_1^k$ will grow slower than $\\xi_2^k$, since $\\xi_1^{k+1} - \\xi_1^k$ would be dominated by $-\\exp\\big(2\\cdot (\\xi_1^k - \\xi_2^k)\\big)$ when $\\xi_1^k$ is big and $\\xi_1^k$ is close to $\\xi_2^k$. We have already revised this error in our current version. \n\n>**Q5**: The ICL tasks are binary classifications but the loss is the squared error.\n\n**A5** Although we choose binary labels in our work for clarity of expression, our results can be easily generated for all distributions with bounded labels. We promise the reviewer we will comment on this in our future version. Meanwhile, even under binary classification, choosing the MSE loss is still beneficial, as it clearly shows that the attention layer can adapt to different ICL tasks under the same objective function.\n\n>**Q6**: Multiple typos.\n\n**A6**: We appreciate the reviewer’s efforts to point these out and have already revised the errors pointed out.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (poster)