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

Q1: Can the message of the paper be generalized beyond this specific training distribution or even kNN?

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.

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

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\}$ 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.

A2: We agree with the reviewer that $y_i$ can depend on $x_i$ and the predictors of 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. However, 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.

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.

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.

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

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.

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

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.

Q6: Multiple typos.

A6: We appreciate the reviewer’s efforts to point these out and have already revised the errors pointed out.

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.

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?

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.

Q2: How well does this generalize to relaxing conditions on the input data lying on the hypersphere under SGD and random initialization?

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.

Q3 Multiple typos throughout the paper.

A3 We are grateful to the reviewers for pointing this out, and have updated our paper by conducting another grammar check.

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.

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.

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.

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.

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.

Q3: Minor questions, including typos and writing unclarity

A3: We are grateful to the reviewers for pointing out those issues, and will clean up the notations in our future version.

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.

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.

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.

Q2: Can the analysis be extended in any trivial way to k-NN (e.g. using multiple heads)?

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.

Q3: Is the independence between $x_i$ and $y_i$ being used throughout, or not.

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.

From a technical perspective, independence condition allows us to achieve a more delicate characterization of gradient descent, which is helpful in achieving the final results.

Q4 Minor questions, including typos and writing unclarity

A4 We highly appreciate your efforts in pointing out those questions and have already revised these issues in our current manuscript.

[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

[2] In-Context Convergence of Transformers, Yu Huang and Yuan Cheng and Yingbin Liang

[3] Trained Transformers Learn Linear Models In-Context, Ruiqi Zhang and Spencer Frei and Peter L. Bartlett

[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