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

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.

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. - Aside from typos, the results are rigorous, well-formulated and non-trivial. - 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. - The proof sketch is mostly well-written (see below). - 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. 2. The training data distribution is not only specific, but also entails that the label of the query is not generated in the same way as the label of the context examples, which is inconsistent with practice. Specifically each example label must be independent of the corresponding input (as well as all other terms) while the query label does depend on the query and the context examples. The reasoning behind the statement that independence of $\{\mathbf{x}\_{i}\}\_{i\in[N+1]}$ and $\{\mathbf{y}\_{i}\}\_{i\in[N]}$ “is also essential to properly study in-context learning of one-nearest neighbor” is incorrect; $\mathbf{y}\_i$ can depend on $\mathbf{x}\_i$ and predictors of both forms mentioned will achieve 50% accuracy. 3. The distribution shift result is solid but not surprising since 1-NN should behave similarly on all distributions for which the true label of the query is in fact the label of the closest nearest neighbor in the context — it requires no knowledge of the data to achieve strong performance. 4. From Lemma 4 and 5, it is not clear how $\xi^k_1$ grows slower than $\xi^k_2$, as is claimed and as is needed for the final result. For any $\xi^k\_1,\xi\_2^k$, the upper bound on $\xi\_1^{k+1}-\xi^k\_1$ can be dominated by $\exp( \text{poly}( N, d ) \xi\_1^k)$, which can be much larger than the lower bound on $\xi\_2^{k+1}-\xi^k\_2$ of $\exp( -\text{poly}( N, d ) \xi\_2^k)$ even when $\xi_1^k = \Omega(\xi_2^k)$. Minor - The ICL tasks are binary classifications but the loss is the squared error. - In Theorem 1, the term $\log(1 - (N\sqrt{d}^{1/d}))$ should be $\log(1 - (N\sqrt{d}^{-1/d}))$? - Equation 2.6 describes gradient ascent, not gradient descent, and the step size is inverted - Lemma 4 clearly has multiple typos, one of which makes it not clear how to deduce the true statement. - The term “epoch” is improperly used in the experiments section. - The caption of Figure 1 says that error bars are plotted but this is false.

Questions

N/A

Rating

4

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors have addressed limitations.

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.

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. **Quality.** The paper is well written and the claims are well-substantiated with detailed proofs and theorems **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. 2. The results show that the model converges to 1-NN predictor on the training data even under SGD and random initialization. How well does this generalize to relaxing conditions on the input data lying on the hypersphere?

Questions

1. Authors mention in the caption of Figure 5 about “error bars”. However, I don’t see any shaded error curve in their figure.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Please refer to the weaknesses section.

Summary

This submission considers learning to implement 1-NN in-context with a single self-attention layer. In particular, they consider training on in-context learning (ICL) sequences of form $(x1, y1), (x2, y2), … (x_N, y_N), x_{N+1}$, where $x_i$ are sampled from the $d$-dimensional unit sphere independently and with uniform probability, $y_i$ are i.i.d. $\pm 1$ labels independent of $x_i$, and $x_{N+1}$ is a test point with a prediction target equal to the prediction of 1-NN. The authors show that with a specific weight initialization gradient descent on the population loss converges to a global minimum that corresponds to implementing 1-NN in-context. The proof relies on the observation that under the proposed weight initialization most parameters stay zero during optimization and it becomes possible to describe the dynamics with only 2 variables. While the loss function written as a function of these 2 variables is nonconvex, they show that these two variables converge to infinity, with their difference converging to infinity too. This limit corresponds to the 1-NN algorithm. The authors also prove that the learned algorithm is robust to distribution shifts, with increasing robustness with the number of gradient descent iterations. Finally, they conduct experiments with random weight initialization and show that a single self-attention layer can be trained to implement 1-NN in-context and be robust to distribution shifts.

Strengths

1. Overall the paper is well-written. The related work is properly referenced. 2. 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. * While the technique is interesting, it depends critically on the initialization scheme. As the experimental results hint, a single self-attention might be able to learn to implement 1-NN in-context even with standard initialization. It would be great to see a discussion on how the technique employed in this work can be useful for proving convergence under standard initialization. * As I understand the employed technique is also tied to the *k=1* case of k-NN. It is unclear whether the technique is general enough to be useful for $k>1$.

Questions

* Lines 127-133: I recommend expanding this part a bit. Also, $W_{3,3}$ should be $-\xi_2$ so that the softmax attention peaks on the example with highest dot product (i.e., the closest point as all points are on the unit sphere). * In the first equation of Lemma 4, the rightmost term should be $\exp(2\xi_1^k - \xi_2^k)$. * Denoting sequences with index $k$ as $\xi^k_1$ and $\xi^k_2$ is confusing. I recommend using either $\xi^{(k)}_1$ and $\xi^{(k)}_2$ notation or even better, $\xi_k$ and $\zeta_k$.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

This submission would benefit form a discussion on the generality of the technique (see the "weaknesses" section).

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.

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. - 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. - 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. - Can the analysis be extended in any trivial way to k-NN (e.g. using multiple heads)? There seems to be a lack of commentary on this, and if this is doable in some simple way, this result should be included. If it turns out to be non-trivial, the paper could benefit from commentary on this as well. - Throughout the paper, the authors refer to the input examples (xs and ys) as being either independent or not independent. It was somewhat unclear to me whether independence was being used throughout, or not.

Questions

- In equation 2.2, the softmax output is directly multiplied with the embedding matrix, meaning that the W_v matrix is set to the identity. While this is spelled out in the text, including this in the equation would improve clarity. - In Assumption 2, I think $\sigma_1$ should just be $\sigma$. - Some of the equations run off of the right-hand side of the page in the appendix.

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

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.