arXiv CS open access

One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context

2026 · arxiv_cs

arXiv CS · Papers · 2026

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

arXiv:2609.01311v1 [cs.LG] 1 Sep 2026

One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context Skanda Athreya∗

Yutong Wang†

James B. Conant High School Hoffman Estates, IL

Illinois Institute of Technology Chicago, IL

Abstract We extend recent work establishing an equivalence between one-layer transformers and nearest-neighbor classifiers in the binary setting to the multiclass case. By leveraging the simplex encoding, we show that one-layer transformers with an argmax classification head behave identically to a one-nearest-neighbor classifier in the multiclass setting. This closes a gap left by prior work, whose multiclass result relied on a non-standard rounding-based approach rather than the typical argmax head used in practice.

Code: https://github.com/skandaka/onenn-multiclass

Introduction

Transformers have become the dominant architecture in large language models and other applications such as computer vision. Understanding the theoretical properties of transformers is therefore of considerable importance. A growing body of work studies theoretical properties of transformers, e.g., in the framework of in-context learning (Bai et al., 2023), and their connection to classical machine learning models such as nearest neighbors (Li et al., 2024). Li et al. (2024) showed that one-layer transformers, under appropriate conditions, behave identically to a one-nearest-neighbor classifier. Their work is, however, limited to the binary classification setting. While they do provide a result with respect to a task resembling multiclass classification (Li et al., 2024, Corollary 1), the result is for a rounding-based approach that departs from the more standard argmax classification head used in practice. In this work, we close this gap by leveraging the simplex encoding (Mroueh et al., 2012). We show that one-layer transformers with an argmax head are equivalent to one-nearest-neighbor classifiers in the multiclass setting. This is a natural and important extension, since the primary application of transformers (i.e., next-token prediction) is inherently a multiclass classification problem, and the argmax head is the standard mechanism for producing predictions.

Related works

2.1

Theory for transformers’ generalization

In-context learning (ICL) is the ability of a trained transformer to solve a new task from a handful of input/output examples placed in its prompt, without any updates to its weights (Bai et al., 2023). ∗ †

athreya1428@students.d211.org ywang562@illinoistech.edu

Understanding how and why this behavior emerges from ordinary training is a central open problem in transformer theory, and a growing amount of work has aimed to investigate it in controlled settings. One influential line of work casts ICL as implicitly running a learning algorithm on the in-context examples: e.g., Zhang et al. (2024) proves that the behavior of a single-layer attention model trained on linear-regression prompts provably converges to a predictor with provable in-distribution and distribution shift guarantees. In a similar vein, the one-nearest neighbor (Li et al., 2024) and our work show that transformers can implicitly act like a one-nearest neighbor algorithm.

2.2

Simplex encoding

Representing K classes as vectors is an important algorithmic design choice which decides how to convert a classifier’s real-valued output to a discrete decision. The simplex encoding assigns each class a vertex of a regular simplex, so that distinct classes are equidistant and symmetric. Mroueh et al. (2012) introduced this coding for multiclass learning and analyzed its geometric properties. Later, Pouliot (2018) established equivalences between simplex-based methods and other multi-category support vector machine formulations (Lee et al., 2004). The simplex encoding has also been applied fruitfully in representation learning (Papyan et al., 2020).

2.3

Nearest neighbor algorithms

The nearest-neighbor algorithm is one of the oldest and most well-studied methods in classification. Cover and Hart (1967) established foundational results on its convergence properties, showing that the risk of the one-nearest-neighbor classifier is bounded above by twice the Bayes risk as the number of samples tends to infinity. Subsequent work established consistency and finite-sample rates for the k-nearest-neighbor rule under margin and smoothness conditions: Chaudhuri and Dasgupta (2014) gives distribution-dependent, finite-sample rates in metric spaces and shows that under the Tsybakov margin condition the nearest-neighbor rate matches known lower bounds for nonparametric classification, and Efremenko et al. (2020) construct a Bayes-consistent classifier that keeps fast query time while attaining rates comparable to minimax. The margin event that controls our distribution-shift result (Section 4.1) plays a role akin to classical margin conditions, where performance is controlled precisely when queries are unlikely to fall close to a boundary between classes. A separate line of work connects nearest-neighbor and kernel methods to attention. Tsai et al. (2019) formulate attention as a kernel smoother over the inputs, with the kernel scores given by input similarities. This perspective supports the equivalence between attention and the nearestneighbor rule that we study. Beyond its role in classical statistics, the nearest-neighbor rule has re-emerged as an important component of modern language models. Khandelwal et al. (2020) introduced kNN-LM which is an “interpolation” of transformers and k-nearest neighbor algorithms. Xu et al. (2023) analyzed why the retrieval-like component in kNN-LM helps generalization. These results indicate that nearestneighbor is a fruitful avenue to understanding transformer-based language models.

Setup and assumptions

Throughout this work, let N be the number of labeled examples in a prompt (the context length) and K the number of classes, both integers. Let [N ] := {1, . . . , N }. Moreover, let d be the dimension of the input token. 2

Testing Theorem 1

Recovery of the 1-NN label 1.0 N = 16 N = 32 N = 64

0.30 0.25

0.8

accuracy

training MSE loss

0.35

0.20 0.15

0.6 0.4

N = 16 N = 32 N = 64 1=K = 0:33

0.2

0.10 0

0.0

100

gradient steps (thousands)

100

gradient steps (thousands)

Figure 1: Convergence (Theorem 1). Training loss (left) decreases and separated-test argmax accuracy (right) approaches 1, faster for smaller context length N , matching the predicted dependence on N . Bands are ±2 standard deviations over 8 seeds. See Section 5 for experimental details.

3.1

Problem Setup

An in-context learning instance consists of a prompt of N labeled examples {(xi , ci )}i∈[N ] and a query xN +1 , where each xi ∈ S d−1 (the unit sphere in Rd ) and each class label ci ∈ [K]. The prompt {(xi , ci )}i∈[N ] and the query xN +1 are known to the learner, while the label cN +1 of the query is hidden from the learner. The objective is to predict cN +1 . Throughout we reserve y for the binary {+1, −1} label encoding used by Li et al. (2024) and write c for a multiclass label in [K]. For the given prompt, we define the index of the nearest neighbor to the query by i∗ := arg min∥xN +1 − xj ∥2 . j∈[N ]

Thus, the one-nearest-neighbor (1-NN) predictor for the prompt is ci∗ . Simplex encoding. To represent K classes as vectors, we use the simplex encoding (Mroueh et al., 2012), implemented as centered one-hot vectors: uc := ec − K1 1K ∈ RK ,

c ∈ [K],

(1)

where ec is the c-th one-hot vector and 1K is the all-ones vector. The vectors {uc }c∈[K] are the vertices of a regular simplex centered at the origin and satisfy: • (zero mean)

c∈[K] uc = 0K ,

• (constant norm) ∥uc ∥22 = K−1 K for each c ∈ [K], and • (constant inner product) ⟨uc , uc′ ⟩ = − K1 for c ̸= c′ . Consequently, the uc ’s are equidistant with distance ∥uc − uc′ ∥2 =

√

These three properties are the only structural facts about the encoding that the convergence proof uses, and they act as the multiclass counterparts of the {+1, −1} label encoding used in the binary case in (Li et al., 2024). 3

Loss Ratio vs various (»1 ; »2 )

0.9

0.8

LMC =Lbin

Loss ratio vs. K

0.7 0.6

0.6

theory (K ¡ 1)=K

measured (MC, (2.0, 5.0))

0.5 2

0.7

K=2 K=3

K=5 K = 10

(K ¡ 1)=K

0.5

(1,2)

number of classes K

(2,5)

(3,6)

(5,8)

(»1 ; »2 )

Figure 2: Scaling identity (Lemma 2). The measured ratio LMC /Lbin , estimated by sampling directly from the expected-version of the loss definitions, matches (K − 1)/K as a function of K (left) and is constant in (ξ1 , ξ2 ) (right). Prompt Embedding. lowing matrix

Following (Li et al., 2024), we embed the prompt and query as the fol

x1  H = uc1 0

··· ··· ···

xN ucN 0



xN +1  0K  ∈ R(d+K+1)×(N +1) 1

(2)

Note that the second row of H, representing the labels, consists of columns that are K-dimensional vectors, in contrast to that of (Li et al., 2024, Eq. (2.1)). Denote the columns of H by h1 , . . . , hN +1 , so that hj = [xj ; ucj ; 0] for j ∈ [N ] and hN +1 = [xN +1 ; 0K ; 1]. We note that the query’s label is 0K since its label is unknown to the learner. Model. The one-layer softmax attention model has a single weight matrix W ∈ R(d+K+1)×(d+K+1) as parameter. The model takes H as input and outputs HW := H · softmax(H⊤ WH),

(3)

where the softmax is applied column-wise. Compared with (Li et al., 2024, Eq. (2.2)), which writes ⊤ W H with separate key and query matrices before merging them the attention score as H⊤ WK Q into a single matrix, we parametrize the model directly by the merged matrix W, playing the role ⊤ W : the model depends on the key and query matrices only through this product, and of WK Q analyzing the merged matrix is the standard reduction in this line of work (Li et al., 2024; Zhang et al., 2024). As in (Li et al., 2024), the value matrix is frozen to the identity, so each output token is a convex combination of the input tokens. The prediction is the K-dimensional vector obtained by taking the (d + 1)-st to (d + K)-th rows of the last column of HW : ℓW (xN +1 ) = [HW ]d+1:d+K, N +1 =

N X

qj (x, W) ucj ∈ RK ,

j=1

(4)

where x := (x1 , . . . , xN +1 ) collects the prompt inputs and the query, and the attention weights are exp(hj⊤ WhN +1 ) qj (x, W) = PN +1 . ⊤ l=1 exp(hl WhN +1 )

(5)

The predicted class is produced by the standard argmax head ĉW (xN +1 ) := arg max [ℓW (xN +1 )]k .

(6)

k∈[K]

3.2

Assumptions on training distribution

Assumption 1 (Training distribution). The inputs {xi }i∈[N +1] are sampled i.i.d. from the uniform distribution on S d−1 , and the class labels ci are sampled uniformly from [K], independently across tokens and of the inputs. We write P train for the resulting joint law of the instance ({(xi , ci )}i∈[N ] , xN +1 ). Note that the following hold: 1. E[uci | x1:N +1 ] = 0K , 2. E[uci uc⊤j | x1:N +1 ] = 0K×K for i ̸= j, and 3. the law of x1:N +1 is invariant under x1:N +1 7→ −x1:N +1 . Expectations. An expectation E[ · ] or probability P( · ) written without a subscript is always taken with respect to the training distribution P train of Assumption 1. Every other law carries an explicit subscript, as in (7) and Theorem 2. Remark 1. Assumption 1 is the multiclass counterpart of the training distribution of (Li et al., 2024, Assumption 1), adopted for the same reasons. Because the labels are independent of the inputs, the prompt contains no parametric relationship between x and c that the model could exploit instead. The only way to reduce the loss is to retrieve the label of the nearest context example. The task is nonetheless nontrivial: each prompt induces a geometrically complicated partition of the sphere. Moreover, the resulting objective is nonconvex (we show this in Lemma 5). Moreover, the analysis uses Assumption 1 only through properties 1,2, and 3, which decouple the label codes from the attention dynamics and force the off-diagonal gradient blocks to vanish. Assumption 2 (Initialization). For a parameter σ > 0, gradient descent is initialized at W = 0

0(d+K)×(d+K) 0d+K 0⊤ −σ d+K

∈ R(d+K+1)×(d+K+1) ,

that is, every entry vanishes except the final diagonal entry, which equals −σ and suppresses the query’s self-attention. The initialization will be exploited as follows. First, the entry −σ plays the role of the masking used in practical self-attention training, by suppressing the query’s attention to its own token, whose label slot is 0K , so that attention concentrates on the labeled context examples. (Li et al., 2024) motivate the same initialization this way, with σ playing the role of a mask value that is effectively set to infinity in practice. Second, it is similar to Trockman and Kolter (2023), who observe that in pretrained transformers the merged key–query product has a dominant positive diagonal, approximately a scaled identity, 5

and show that imposing this structure at initialization improves trainability. Gradient descent started at W0 remains on the “diagonal family” diag{ξ1t Id , 0K×K , −ξ2t } with ξ1t increasing, as shown in the appendix, so the trained key–query block remains a multiple of the identity throughout training. Training. We train W with gradient descent on the expected (population) mean-squared error between the logits and the simplex code of the 1-NN label, L(W) = 12 EP train ∥ℓW (xN +1 ) − uci∗ ∥22 ,

(7)

initialized according to Assumption 2.

Main results

This section presents our main theoretical results. Omitted proofs are all included in the appendix. The analysis reduces the dynamics to a two-parameter family of weight matrices, which Assumption 2 initializes and which gradient descent provably preserves. Lemma 1 (Two-parameter reduction). Under Assumptions 1 and 2 there exist real sequences {ξ1t }t≥0 and {ξ2t }t≥0 with Wt = diag{ξ1t Id , 0K×K , −ξ2t } for all t ≥ 0. Moreover the blocks W12 , W22 , W32 , which multiply the query’s zero label block, have identically vanishing gradient, so the loss does not depend on them at any point of the parameter space. Write LMC (ξ1 , ξ2 ) := L(W) for the loss (7) restricted to this 2-parameter space. Next, we briefly recall the binary case from (Li et al., 2024). Let Lbin (ξ1 , ξ2 ) be the binary objective from (Li et al., 2024, Eq. 2.5), which we recall below. For the binary case, the prompt and query are the same as in Assumption 1. The simplex codes ucj are however replaced by scalar labels yj ∈ {±1} drawn uniformly and independently of the P inputs, the logit ℓW (xN +1 ) = N j=1 qj (x, W) yj is a scalar, and the target is yi∗ , so that Lbin (ξ1 , ξ2 ) = 21 E ℓW (xN +1 ) − yi∗

(8)

The attention weights qj are identical in the two models, because by Lemma 1 they depend on the inputs only through their inner products and not on any label. Now, the two losses are related by the identity: Lemma 2 (Scaling identity). On the two-parameter family, LMC (ξ1 , ξ2 ) = K−1 K Lbin (ξ1 , ξ2 ). Consequently ∂ξm LMC = K−1 ∂ L for m ∈ {1, 2}. ξm bin K Every increment bound in the convergence analysis is therefore the binary bound scaled by

K−1 K . The factor rescales the growth of (ξ1 , ξ2 ) but cancels in the ratio that controls the trajectory.

√ Theorem 1 (Convergence). Suppose Assumptions 1 and 2 hold with N ≥ O( d log d) and n √ σ > 2 max log(N d), − log 1 − (N d)−1/d , 1 Cd 1 − 2N

where Cd = poly(d). These are the conditions of (Li et al., 2024, Theorem 1), which we inherit unchanged.1 Then gradient descent on (7) satisfies L(Wt ) ≤ O

K − 1 poly(N, d) , · K log t

so L(Wt ) converges to 0 as t → ∞.

4.1

Behavior Under Distribution Shift and the Argmax Head

For δ > 0 define the event Aδ := ∥xj − xN +1 ∥22 ≥ ∥xi∗ − xN +1 ∥22 + δ,

for all j ̸= i∗ with cj ̸= ci∗ . (9) which captures the set of instances with large margin. Theorem 2 (Distribution shift). Let P test be any distribution over instances ({(xi , ci )}i∈[N ] , xN +1 ) supported on S d−1 with labels in [K]. After T steps of gradient descent, EP test ∥ℓWT (xN +1 ) − uci∗ ∥22

≤O

K−1 inf K δ>0

2 − poly(N,d) δ N T + P test (Acδ )

Remark 2 (Training vs. test distributions). Here P test is a distribution over the entire instance ({(xi , ci )}i∈[N ] , xN +1 ). Beyond the stated “supported on S d−1 ” conditions, no independence is assumed between inputs and labels, or across tokens, and no relation to P train is assumed. This mirrors (Li et al., 2024, Assumption 3). The label boundedness required there holds automatically here since ∥uc ∥2 ≤ 1. A guarantee under an arbitrary P test is possible because the underlying estimate is pointwise rather than distributional: training on P train produces the fixed parameters WT = diag{ξ1T Id , 0K×K , −ξ2T }, and on the event Aδ the deviation ∥ℓWT (xN +1 ) − uci∗ ∥2 admits a deterministic bound, decaying in ξ1T δ and ξ2T , which holds for every prompt–query configuration on the sphere (the pointwise bound is derived in the appendix). The test distribution enters only through P test (Acδ ), the probability that the query falls within δ of a 1-NN decision boundary. Corollary 1 (Argmax classification). Under the setting of Theorem 2, P test ĉWT (xN +1 ) ̸= ci∗

≤ O inf N 2 T − poly(N,d) δ + P test (Acδ )

δ>0

Moreover, if P test (Aδ∗ ) = 1 for some δ ∗ > 0, then ĉWT (xN +1 ) = ci∗ almost surely once log T ≥ O log(KN )/(poly(N, d) δ ∗ ) . √ (Li et al., 2024, Theorem 1) print the middle term with exponent 1/d. Since (N d)1/d > 1, the logarithm is then applied to a negative number. We state the well-posed form with exponent −1/d, which is the quantity their Lemma 12 identifies this condition with, namely − log aN,d up to constants. 1

Testing Corollary 1 argmax accuracy w.r.t 1-NN

Testing Theorem 2

0.8

MSE k` ¡ uc ¤ k 2

K=3 K=5

0.6 0.4 0.2 0.0 0

gradient steps (thousands)

1.0 0.8 0.6 0.4 K = 3 separated K = 3 uniform K = 5 separated K = 5 uniform

0.2 0.0

gradient steps (thousands)

Figure 3: Distribution shift (Theorem 2) and argmax classification (Corollary 1). On the separated test distribution the logit MSE ∥ℓ − uci∗ ∥22 → 0 (left) and the argmax matches the 1-NN label (right). The uniform-data accuracy is shown for contrast.

Experiments

We validate the theory with six experiments. Complete configurations and code to reproduce all experimental results are available at https://github.com/skandaka/onenn-multiclass. Compute environment is listed in Section D. Protocol. All experiments use the one-layer attention model of (3) with the score H⊤ WH, the centered one-hot vector encoding (1), and the loss (7). Training data follows Assumption 1 (inputs i.i.d. uniform on S d−1 , labels uniform on [K]), and W is initialized as in Assumption 2. We run stochastic gradient descent on minibatch estimates of (7) with analytical gradients, which we verify against finite differences in double precision to relative error ≈ 10−10 . Training runs in PyTorch (Paszke et al., 2019). The theory analyzes gradient descent on the population objective (7). The minibatch updates are a stochastic approximation of the population gradient. (Li et al., 2024) likewise validate their theory with stochastic gradient descent, in a less restrictive setting still with random initialization whereas we retain the initialization of Assumption 2. Unless stated otherwise, d = 8, N = 16, σ = 3.0, learning rate η = 0.05, and batch size 128. Additionally, we evaluate on a well-separated test distribution, extending the construction of (Li et al., 2024, Appendix B) to K classes, which places the query strictly closer to examples of a single class than to all others, so that P test (Aδ∗ ) = 1 for a fixed margin δ ∗ > 0, matching the setting of Theorem 2 and Corollary 1. Scaling identity (Figure 2). We test Lemma 2 directly. The identity is an exact identity between two population expectations that admit no closed form, so we verify it by estimating both sides directly from their definitions. For fixed (ξ1 , ξ2 ), we sample prompts and queries from Assumption 1, estimate LMC and Lbin on this sample, and plot their ratio. At (ξ1 , ξ2 ) = (2.0, 5.0) the measured ratio traces (K − 1)/K over K ∈ {2, 3, 4, 5, 7, 10} (80000 samples per point, left), and it is constant over the grid (ξ1 , ξ2 ) ∈ {(1, 2), (2, 5), (3, 6), (5, 8)} for K ∈ {2, 3, 5, 10} (40000 samples per point, right), with deviations within error tolerance confirming Lemma 2. 8

10 6

training MSE loss

0.1

0.6

0.2

0.8

0.3

1.0

K=2 K=3 K=5 K = 10

0.4

Final argmax accuracy vs K

final separated-test accuracy

Convergence for each K

0.4 0.2 0.0

number of classes K

gradient steps (thousands)

Figure 4: Convergence holds for every K ∈ {2, 3, 5, 10} (left), with final separated-test accuracy 1 across K (right).

Dist to diagonal initialization 10

0.8 0.6

" = 3, inert blocks only " = 3, active blocks only

0.4

random 1=K perturb all blocks

kW t ¡ ¦(W t )k F

final separated-test acc

Robustness to noise at initialization 1.0

0.2

10 0.0

−2

−1

" = 0:01 " = 0:1 " = 1:0 " = 3:0

−1

gradient steps (thousands)

noise scale "

Figure 5: Perturbed initialization. Left: final separated-test argmax accuracy against the noise scale ε, with the two block-restricted ablations at ε = 3. Right: distance from the two-parameter family of Lemma 1 over training, for the blocks that enter the attention scores. Convergence (Figure 1). We train with K = 3 and context lengths N ∈ {16, 32, 64} for 105 steps and 8 random seeds. The training loss decreases towards zero at the slow logarithmic rate of Theorem 1, more slowly for larger N , consistent with the poly(N, d) dependence of the bound (left). Argmax accuracy on the separated test distribution rises from the random-guess baseline 1/K ≈ 0.33 to 1 (right). Bands show ±2 standard deviations over the seeds. Distribution shift and the argmax head (Figure 3). For K ∈ {3, 5} we train for 6×104 steps and 6 random seeds, and evaluate on the separated distribution, which differs from the training distribution but satisfies the margin condition. The logit error ∥ℓ − uci∗ ∥22 decays toward zero (left) and argmax accuracy reaches 1 (right), as Theorem 2 and Corollary 1 predict. On uniform test data, on the other hand, accuracy plateaus below 1: a constant fraction of queries falls near a 1-NN decision boundary, so P test (Acδ ) does not vanish for any fixed δ and the second term of the bound in Theorem 2 remains. We note that the plateau is as predicted by the 9

1.0

Final accuracy across the loss family 1.0

0.8

final accuracy

separated-test accuracy

(1 ¡ c) L sq + c L ce

0.6 c = 0:0 c = 0:25 c = 0:5 c = 0:75

0.4 0.2 0.0

c = 0:9 c = 1:0 squentropy 20

0.8 0.6 0.4 separated test uniform test squentropy (separated)

0.2 0.0

gradient steps (thousands)

0.0

0.2

0.4

0.6

0.8

1.0

c (where 0 ! L sq and 1 ! L ce )

Figure 6: Loss function. Left: separated-test argmax accuracy over training for (1 − c)Lsq + cLce and for squentropy. Right: final accuracy on the separated and uniform test distributions across the family. theory. Dependence on the number of classes (Figure 4). We vary K ∈ {2, 3, 5, 10}, training for 5 × 104 steps and 4 random seeds. Convergence holds for every K: the prefactor K−1 K shifts the loss curves but leaves the logarithmic rate unchanged (left), and the final separated-test accuracy is 1 for every K (right). The guarantees of Theorem 1 and Corollary 1 are therefore uniform in the number of classes. Dependence on the diagonal initialization (Figure 5). Assumption 2 prescribes an exactly diagonal W0 , which has precedent in the mimetic initialization of (Trockman and Kolter, 2023). We test how far the 1-NN phenomenon depends on it by adding i.i.d. N (0, ε2 ) noise to every entry of W0 and sweeping ε, training for 4 × 104 steps and 3 random seeds at each value. Alongside accuracy we track the distance from the two-parameter family of Lemma 1, that is ∥Wt − Π(Wt )∥F where Π projects onto diag{ξ1 Id , 0K×K , −ξ2 } (Figure 5 right). The separated-test argmax accuracy is 1.000 for every ε ≤ 1.0, degrades at ε = 2.0 (0.78 ± 0.28) and fails at ε = 3.0 (0.44 ± 0.07), so the guarantee survives perturbations up to roughly the scale σ of the mask entry itself. We observe first that, wherever the result survives, the distance to the family contracts over training, to about a third of its initial value (1.056 → 0.343 at ε = 0.1 and 10.558 → 3.449 at ε = 1.0); where it fails, the contraction weakens to 0.79 (31.673 → 25.145 at ε = 3.0). Lemma 1 states that the family is invariant under gradient descent, and this indicates that it is also locally attracting: the trajectory returns to the diagonal structure rather than merely remaining on it. Moreover, failing to return to the diagonal structure seems to imply failing to learn the 1-NN rule. Second, we observe that the component of the perturbation lying in W12 , W22 , W32 is unchanged to three decimal places throughout training (5.400 → 5.400 at ε = 1.0), as the identically zero gradients established in Lemma 1 require. Consistently, noise of magnitude ε = 3.0 confined to those blocks leaves the accuracy at 1.000, while the same magnitude confined to the remaining blocks reproduces the failure (0.44).

Dependence on the loss function (Figure 6). The analysis thus far (and also in (Li et al., 2024)) is specific to the square loss: Lemma 2 is an identity between coordinatewise expansions of (7) and has no cross-entropy analogue. We test whether the one-nearest-neighbor phenomenon extends to “convex-combination” losses of the form (1 − c) Lsq + c Lce , interpolating from the square loss at c = 0 to cross entropy at c = 1. We also test against the “squentropy” loss of (Hui et al., 2023), Lsqen = Lce +

X 1 [ℓ]2 , K − 1 k̸=c ∗ k i

Unlike the convex-combination losses above, the squentropy loss has no tunable parameter. Both families are trained for 4 × 104 steps and 3 random seeds, using the same hyperparameters as in the square-loss experiments. The separated-test argmax accuracy is 1.000 for every c > 0 and for squentropy, and 0.999±0.002 at c = 0, so the recovery of the 1-NN label does not depend on the choice of loss. On uniform test data the accuracy increases mildly and monotonically with c, from 0.908 at c = 0 to 0.924 at c = 1, with squentropy at 0.916. Recall that the square-loss gradient ≈ 0 when ℓ ≈ uci∗ , whereas the cross-entropy gradient never vanishes. Training under cross entropy therefore continues to boost attention on the nearest neighbor, reducing the near-boundary errors seen in the uniform-data plateau. Taken together, our result shows that the one-nearest-neighbor phenomenon is robust to changing the loss function. While our theory covers the square loss only, an interesting line of future work will analyze the cross entropy.

Discussion, Limitations, and Future Work

We have rigorously shown the equivalence between a one-layer softmax attention model and the one-nearest-neighbor rule, extending the binary setting result by Li et al. (2024). We show that the convergence rate, the conditions on σ and N , and the distribution-shift guarantee all carry over without K-dependence. We discuss several limitations in and future directions of this work. The analysis exploits the onelayer architecture and the special diagonal initialization of Assumption 2. While these restrictions have precedent in the literature (e.g., in (Trockman and Kolter, 2023)) an interesting avenue of future research is to remove these restrictions. Our experiments show the one-nearest-neighbor phenomenon persists after switching the loss function from square loss to squentropy. It will be interesting to develop a rigorous theory to explain this phenomenon. The scaling identity is specific to the square loss and has no cross-entropy analogue, though the empirical behavior does not depend on that choice. Finally, the guarantee concerns recovery of the 1-NN rule for a single head. Extending the argument to multiple heads is a natural next step.

References 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:57125–57211, 2023. 11

Kamalika Chaudhuri and Sanjoy Dasgupta. Rates of convergence for nearest neighbor classification. In Advances in Neural Information Processing Systems (NeurIPS), volume 27, 2014. Thomas Cover and Peter Hart. Nearest neighbor pattern classification. IEEE transactions on information theory, 13(1):21–27, 1967. Klim Efremenko, Aryeh Kontorovich, and Moshe Noivirt. Fast and Bayes-consistent nearest neighbors. In International Conference on Artificial Intelligence and Statistics (AISTATS), pages 1276–1286. PMLR, 2020. Like Hui, Mikhail Belkin, and Stephen Wright. Cut your losses with squentropy. In Proceedings of the 40th International Conference on Machine Learning, pages 14114–14131. PMLR, 2023. Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference on Learning Representations (ICLR), 2020. Yoonkyung Lee, Yi Lin, and Grace Wahba. Multicategory support vector machines: Theory and application to the classification of microarray data and satellite radiance data. Journal of the American Statistical Association, 99(465):67–81, 2004. Zihao Li, Yuan Cao, Cheng Gao, Yihan He, Han Liu, Jason M. Klusowski, Jianqing Fan, and Mengdi Wang. One-layer transformer provably learns one-nearest neighbor in context. In Advances in Neural Information Processing Systems, volume 37, pages 82166–82204, 2024. Youssef Mroueh, Tomaso Poggio, Lorenzo Rosasco, and Jean-Jacques Slotine. Multiclass learning with simplex coding. Advances in Neural Information Processing Systems, 25, 2012. Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40):24652– 24663, 2020. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32, 2019. Guillaume Pouliot. Equivalence of multicategory SVM and simplex cone SVM: Fast computations and statistical theory. In International conference on machine learning, pages 4133–4140. PMLR, 2018. Asher Trockman and J. Zico Kolter. Mimetic initialization of self-attention layers. In International Conference on Machine Learning (ICML), pages 34456–34468. PMLR, 2023. Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer’s attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 4344–4353, 2019. Frank F. Xu, Uri Alon, and Graham Neubig. Why do nearest neighbor language models work? In International Conference on Machine Learning (ICML), pages 38325–38341. PMLR, 2023. 12

Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models incontext. Journal of Machine Learning Research, 25(49):1–55, 2024.

Proof of the Convergence Theorem

We prove Theorem 1 as follows. Lemma 3 (Properties of the simplex code). Under Assumption 1, the simplex codes uci defined in (1) satisfy (i) E[uci | x1:N +1 ] = 0K , (ii) E[uci uc⊤j | x1:N +1 ] = 0K×K for i ̸= j, (iii) ∥uci ∥22 = K−1 K . Proof. For (i), the labels are independent of the inputs and the K classes are equally likely, so P 1 1 1 E[uci | x1:N +1 ] = E[uci ] = K1 K c=1 (ec − K 1K ) = K 1K − K 1K = 0K . For (ii), ci and cj are independent for i ̸= j and both are independent of x1:N +1 , so E[uci uc⊤j | x1:N +1 ] = E[uci ] E[ucj ]⊤ = 0K×K by (i). For (iii), for any class c, ∥uc ∥22 = ∥ec − K1 1K ∥22 = ∥ec ∥22 − K2 ⟨ec , 1K ⟩ + K12 ∥1K ∥22 = 1 − K2 + K1 = K−1 K , which is independent of c.

A.1

Closed-Form Gradient

Lemma 4 (Closed-form gradient). Under Assumptions 1 and 2, for all t ≥ 0, X N

∇W11 L(W ) = E t

⊤ git ⟨xi , xN +1 ⟩ xi xN +1

(10)

i=1

⊤ + git∗ ⟨xi∗ , xN +1 ⟩ xi∗ xN +1

for scalar-valued functions {git }i∈[N ] ∪ {git∗ } of one real variable, and ∇Wij L(Wt ) = 0 for every block (i, j) except (1, 1) and (3, 3). Proof. We induct on t. The form holds at t = 0 by Assumption 2. Assume it holds at step t ≥ 0, so that Wt = diag{ξ1t Id , 0K×K , −ξ2t }. Using ∥xN +1 ∥2 = 1, the score hj⊤ Wt hN +1 of (3) reduces to ξ1t ⟨xj , xN +1 ⟩ for j ∈ [N ] and to ξ1t − ξ2t for the query, so the attention weights take the scalar form qj = PN

l=1 exp

exp ξ1t ⟨xj , xN +1 ⟩

ξ1t ⟨xl , xN +1 ⟩ + exp(ξ1t − ξ2t )

(11)

for j ∈ [N ], with exp(ξ1t − ξ2t ) in the numerator for j = N + 1. Below we abbreviate qj = qj (x, Wt ). This is the multiclass counterpart of (Li et al., 2024, Eq. (C.9)).

The consequence used repeatedly below is that no label code appears in (11): each qj is a measurable function of x1:N +1 alone. Hence for any function Φ of the labels and any j, j ′ ∈ [N ], E qj qj ′ Φ(c1:N ) = E qj qj ′ E[ Φ(c1:N ) | x1:N +1 ] ,

(12)

by the tower property: the weights may be taken outside the inner expectation, which is then evaluated by Lemma 3. This step is the only place where the independence of the labels from the inputs in Assumption 1 is used. Partition W conformally with hj = [xj ; ucj ; zj ], 



W11 W12 W13   W = W21 W22 W23  , W31 W32 W33 with W11 ∈ Rd×d , W22 ∈ RK×K , W33 ∈ R, and zj = 0 for j ∈ [N ], zN +1 = 1. Since hN +1 = [xN +1 ; 0K ; 1], the score expands as hj⊤ WhN +1 = xj⊤ W11 xN +1 + W13 + uc⊤j

W21 xN +1 + W23

(13)

+ zj W31 xN +1 + W33 .

We treat the blocks in three groups. (i) The blocks W12 , W22 , W32 . These do not occur in (13): they multiply the second block of hN +1 , which is 0K . Every score, hence every attention weight, hence L, is therefore constant in these blocks, and ∇W12 L = ∇W22 L = ∇W32 L = 0 identically — not merely at the initialization. This is the multiclass form of the observation that the label-to-output column does not influence the prediction (Li et al., 2024, Section 2.2). (ii) The blocks W21 , W13 , W23 . By (13) these enter every score linearly, so by the chain rule P (·)] with weights γl built from the qj and each gradient is an expectation of the form E[ l γl uc⊗m l the residual ℓ − uci∗ , and with m ≤ 3; this is the counterpart of (Li et al., 2024, Eqs. (C.4)– (C.6)). Applying (12) and Lemma 3 term by term: a monomial containing a code ucl whose label appears exactly once has conditional expectation E[ucl | x1:N +1 ] = 0K by Lemma 3(i); a monomial containing ucl uc⊤l reduces to the constant K−1 K by Lemma 3(iii), leaving a residual factor that is again linear in a single code and so has zero mean; and a monomial containing two distinct codes has conditional expectation 0K×K by Lemma 3(ii). Hence ∇W21 L = ∇W13 L = ∇W23 L = 0. (iii) The block W31 . By (13) it appears only in the term zj W31 xN +1 , that is, only in the query’s ⊤ own score, contributing xN +1 to the gradient. Under the induction hypothesis the weights qj are ⊤ given by (11) and are therefore even in x1:N +1 , while xN +1 is odd; since the input law is invariant under x1:N +1 7→ −x1:N +1 (Assumption 1, property 3), the expectation vanishes. This argument uses no label structure and is identical to the binary case. The indicator block W33 has a partial derivative of 1 with respect to the query’s self-attention score, guaranteeing a non-zero update during gradient descent: 1 η

ξ2t+1 − ξ2t = −∂ξ2 L(Wt ) = ∇W33 L(Wt )

∗ = K−1 K E qN +1 qi −

2 j=1 qj

(14)

Expanding ∇W11 L and using Lemma 3 to discard the label terms leaves (10), in which the inputs enter only through inner products. By (Li et al., 2024, Lemma 7) (rotational invariance of the uniform distribution on S d−1 ), such a gradient is a scalar multiple of the identity, ∇W11 L(Wt ) = 0 = 0 this gives Wt+1 = ξ t+1 I , completing the induction. at Id . With W11 d 11 1 14

A.2

Two-Dimensional Reduction

Proof of Lemma 1. The claim is the diagonal form established in the induction of Lemma 4; it remains to be seen why the K × K label block does not contribute free parameters. At initialization 0 =0 W22 K×K by Assumption 2. By the structural-zero argument above, W22 multiplies the query’s (zero) label rows in every score, so ∇W22 L = 0 at every step; a block that starts at zero and receives a zero gradient stays at zero. Thus, although the ambient dimension grows from (d + 2)2 to (d + K + 1)2 , the trajectory explored by gradient descent is governed by the two scalars ξ1t and ξ2t , exactly as in the binary case.

A.3

The Scaling Identity

We now prove Lemma 2. Proof of Lemma 2. By (7) and (4), expanding the squared norm over coordinates, LMC (ξ1 , ξ2 ) = 12

K X

  N X 2 E qj [uc ]m − [uc ∗ ]m  , j

m=1

j=1

with the input-only weights qj = qj (x, W) of (11). Expanding the square gives a quadratic term, a cross term, and a constant; we sum each over m using Lemma 3. The quadratic P P PK K−1 P 2 K−1 2 2 term is j,j ′ qj qj ′ m=1 E[[ucj ]m [ucj ′ ]m ] = j qj , since m E[[ucj ]m ] = ∥ucj ∥2 = K K by Lemma 3(iii) while distinct tokens are uncorrelated by Lemma 3(ii). The cross term P P K−1 ∗ j qj m E[[ucj ]m [uci∗ ]m ] contributes only through j = i , giving K qi∗ in expectation. The K−1 times the corresponding term of the constant term is E[∥uci∗ ∥22 ] = K−1 K . Each is exactly K 2 binary-analog loss, in which y ∈ {±1} gives E[y ] = 1 and E[yj yj ′ ] = 0 for j ̸= j ′ . Therefore LMC = K−1 K Lbin , and differentiating gives the gradient identity. Lemma 5 (Nonconvexity). The reduced loss LMC (ξ1 , ξ2 ) is nonconvex. K−1 > 0. The original binary reduced loss Lbin is Proof. By Lemma 2, LMC = K−1 K Lbin with K nonconvex (Li et al., 2024, Appendix E), and a positive scalar multiple of a nonconvex function is nonconvex; this is why LMC is nonconvex, and why the trajectory analysis below is needed.

A.4

Evolution of the Two Parameters

Each bound below is the binary estimate of (Li et al., 2024) multiplied by K−1 via Lemma 2. K √ −2/(d−3) ′ ′ are inherited unchanged from the The constants c1 , . . . , c4 , c1 , c2 , Cd , and aN,d = (2N d) sphere geometry estimates (Li et al., 2024, Lemmas 17–20), which have no labels. Lemma 6 (Increment bounds for ξ1 ). For ξ1t ≥ 0, there exist constants c1 , c2 , c3 , c4 > 0 such that t

d η

−6ξ1 ξ1t+1 − ξ1t ≥ K−1 − c2 e2ξ1 −ξ2 , K c1 e

d η

poly(N,d) ξ1 ξ1t+1 − ξ1t ≤ K−1 − c4 e2(ξ1 −ξ2 ) . K c3 e

Proof. By Lemma 1 the ξ1 update is ξ1t+1 − ξ1t = −η ∂ξ1 LMC (ξ1t , ξ2t ), and from Lemma 2 we have ∂ξ1 LMC = K−1 K ∂ξ1 Lbin . Applying (Li et al., 2024, Lemma 4) to the binary gradient gives both inequalities.

Lemma 7 (Increment bounds for ξ2 ). For ξ1t ≥ 0, there exist constants c′1 , c′2 > 0 such that t

1 η

′ − poly(N,d) ξ2 ξ2t+1 − ξ2t ≥ K−1 , K c1 e

1 η

′ − poly(N,d) ξ2 ξ2t+1 − ξ2t ≤ K−1 . K c2 e

Proof. The ξ2 update is ξ2t+1 − ξ2t = −η ∂ξ2 LMC , with ∂ξ2 LMC = K−1 K ∂ξ2 Lbin by Lemma 2, and (Li et al., 2024, Lemma 5) bounds the binary gradient. In particular the parameter updates follow the recurrence relation bt+1 − bt ≥ Ce−αbt with C, α > 0, so ξ2t = Ω poly(N, d) log t and ξ2t → ∞. Note that the factor K−1 K , like the step size η, is absorbed into C and therefore contributes only an 1 additive α1 log(αC) to ξ2t ; since K−1 K ∈ [ 2 , 1), this offset is bounded uniformly in K and does not affect the rate. Lemma 8 (Combined lower bound). For all t ≥ 0, 1 η

−6ξ1 1 d(ξ1t+1 − ξ1t ) + 2(ξ2t+1 − ξ2t ) ≥ K−1 . K 1 − 2N Cd e

Proof. The left-hand side equals K−1 K d ∂ξ1 Lbin + 2 ∂ξ2 Lbin by Lemma 2, to which (Li et al., 2024, Lemma 8) applies. The binary estimate uses only ⟨xi∗ , xN +1 ⟩ ≥ ⟨xj , xN +1 ⟩ and the constant Cd from (Li et al., 2024, Lemma 17), both unchanged by K. √ Lemma 9 (Sharp upper bound for the ξ1 increment). For ξ1t ≥ 0 and N ≥ O( d log d),

1 η

ξ1t+1 − ξ1t ≤ K−1 K

4 t a 2(ξ1t −ξ2t ) 2N − (N +1)2 ξ1 − dNN,d 3e e d e

Proof. Apply Lemma 2 to (Li et al., 2024, Lemma 9); the constant aN,d comes from the sphere tail bound (Li et al., 2024, Lemma 19) and is unchanged. Lemma 10 (Refined lower bound for the ξ1 increment). For ξ1t ≥ 0, d η

1 ξ1t+1 − ξ1t ≥ K−1 1 − 2N Cd e−6ξ1 − 2 e2ξ1 −2ξ2 . K

Proof. Combine d times Lemma 8 with 2 times the upper bound of Lemma 9; the common factor K−1 K factors out. Lemma 11 (Sharp lower bound for the ξ2 increment). For ξ1t ≥ 0, 1 η

2aN,d ξ1 −2ξ2 1 ξ2t+1 − ξ2t ≥ K−1 . K · (N +1)3 e e

2 t Proof. The binary proof bounds η1 (ξ2t+1 − ξ2t ) ≥ E[qi∗ (x, Wt ) qN +1 (x, W )]; by Lemma 2 the multiK−1 class quantity acquires K . The remaining estimates — the concentration bound 1−⟨xi∗ , xN +1 ⟩ ≥ aN,d , which holds with probability at least 1e by (Li et al., 2024, Lemma 19), and qi3∗ ≥ (N + 1)−3 — involve only geometry and are unchanged.

A.5

Ratio Bound and Growth Rates

Lemma 12 (Ratio bound). If σ = ξ20 ≥ 3 log t ≥ 0.

2N 4 d aN,d

7 t and ξ1t ≥ 0 for all t, then ξ1t ≤ 15 ξ2 for all

Proof. By Lemma 9, the increment ξ1t+1 − ξ1t is nonpositive once 4 t aN,d 2(ξ t −ξ t ) K−1 2N − (N +1)2 ξ1 1 2 . < K−1 K · d e K · dN 3 e e

t t The factor K−1 K > 0 cancels from both sides, so the resulting condition on (ξ1 , ξ2 ) is exactly the threshold comparison of the binary analysis: no label code enters it, and it involves only the two scalars and the sphere constant aN,d . Under the stated bound on σ, (Li et al., 2024, Lemma 12) t+1 t t − ξ1t ≤ 0 there. Hence ξ1t cannot grow shows that the condition is met whenever 15 7 ξ1 ≥ ξ2 , so ξ1 7 t t past 15 ξ2 ; since ξ2 does not decrease by Lemma 11, the bound holds for all t. Only ξ1t ≤ c ξ2t for some constant c < 21 is used downstream, in Theorem 1 and Theorem 2.

Lemma 13 (Growth rates). With σ and N as in Theorem 1, ξ1t , ξ2t = Ω poly(N, d) log t , with 7 t ξ2 for all t ≥ 0. ξ1t ≤ 15

Proof. For ξ2 , Lemmas 7 and 11 bound the increment between two exponentially decaying rates, η K−1 −2ξ2t ≤ ξ t+1 − ξ t ≤ K−1 η e−ξ2t /15 ; comparison with the corresponding continuous flow 2 2 K · (N +1)3 e e K gives ξ2t = Ω poly(N, d) log t , the factors K−1 K and η contributing only a bounded additive offset.

C (1− 1 )

For ξ1 , Lemma 10 gives ξ1t+1 −ξ1t ≥ 0 while 8ξ1t ≤ 2ξ2t +log d 2 2N . This condition holds through an initial phase because ξ20 = σ is large and ξ2t increases. During that phase ξ1t grows logarithmically, and the upper bound of Lemma 9 caps it at O(poly(N, d) log t). Hence ξ1t = Ω poly(N, d) log t . 7 t The bound ξ1t ≤ 15 ξ2 is Lemma 12. Proof of Theorem 1. By Lemma 2 applied to the binary loss bound (Li et al., 2024, Lemma 14), the multiclass loss obeys h

E ∥ℓWt − uci∗ ∥22

≤ K−1 K O

N 3 kd2 ξ1t

2ξ1 −ξ2 . + K−1 K e

Substituting the growth rates of Lemma 13, the first term is O K−1 K · poly(N, d)/ log t , and since 7 t 1 t t t t ξ1 ≤ 15 ξ2 by Lemma 12 gives 2ξ1 − ξ2 ≤ − 15 ξ2 = O(− log t), the second term decays as O(t−1/15 ). Both vanish as t → ∞, which is the bound of Theorem 1.

Distribution Shift and Argmax Classification

We prove Theorem 2 and Corollary 1. Throughout, ℓ := ℓWT (xN +1 ) and qj := qj (x, WT ). By the convergence analysis, WT is diagonal with x-block ξ1T Id and indicator entry −ξ2T , so the attention weights depend on the inputs only through inner products: exp(ξ1T ⟨xj , xN +1 ⟩) T T T l=1 exp(ξ1 ⟨xl , xN +1 ⟩) + exp(ξ1 − ξ2 )

(15)

qj = PN for j ∈ [N ].

Proof of Theorem 2. Since the query contributes 0K to (4), we have ℓ = N j=1 qj ucj and 1 − qN +1 . Writing the target as uci∗ = (1 − qN +1 ) uci∗ + qN +1 uci∗ and grouping, P

ℓ − uci∗ =

N X

j=1 qj =

qj ucj − uci∗ − qN +1 uci∗

j=1

qj ucj − uci∗ − qN +1 uci∗ ,

j∈[N ] cj ̸=ci∗

(16)

where the terms with cj = ci∗ cancel because ucj = uci∗ . Taking norms in (16) and using ∥ucj − √ p uci∗ ∥2 = 2 and ∥uci∗ ∥2 = (K − 1)/K ≤ 1, √ X qj + qN +1 . (17) ∥ℓ − uci∗ ∥2 ≤ 2 j∈[N ] cj ̸=ci∗

On the event Aδ of (9), every j with cj ̸= ci∗ satisfies ∥xj − xN +1 ∥22 ≥ ∥xi∗ − xN +1 ∥22 + δ. Since ∥xa − xN +1 ∥22 = 2 − 2 ⟨xa , xN +1 ⟩ on the sphere, this is equivalent to ⟨xj , xN +1 ⟩ ≤ ⟨xi∗ , xN +1 ⟩ − 2δ . Hence, from (15) and qi∗ ≤ 1, X

qj ≤

j∈[N ] cj ̸=ci∗

qj ∗ q j∈[N ] i X

cj ̸=ci∗

exp ξ1T (⟨xj , xN +1 ⟩ − ⟨xi∗ , xN +1 ⟩)

j∈[N ] cj ̸=ci∗

≤ N exp − 12 ξ1T δ .

(18)

For the self-attention weight, qN +1 /qi∗ = exp(ξ1T (1 − ⟨xi∗ , xN +1 ⟩) − ξ2T ) ≤ exp(2ξ1T − ξ2T ), and by 1 T 7 T the convergence analysis ξ1T ≤ 15 ξ2 , so 2ξ1T − ξ2T ≤ − 15 ξ2 and 1 T qN +1 ≤ exp − 15 ξ2 .

(19)

Combining (17), (18), and (19), ∥ℓ − uci∗ ∥2 ≤

√

1 T ξ2 2 N exp − 12 ξ1T δ + exp − 15

(20)

uniformly on Aδ . By Lemma 13, ξ1T , ξ2T = Ω(poly(N, d) log T ), so both terms are O N T − poly(N,d)δ . to take expectations. Splitting on Aδ and its complement, and using ∥ℓ − uci∗ ∥22 ≤ √ It remains K−1 2 ( 2 + 1) · K as a hard uniform bound on Acδ (every qj ≤ 1 and ∥uc ∥22 = K−1 K ),

EP test ∥ℓ − uci∗ ∥22

= EP test ∥ℓ − uci∗ ∥22 1{Aδ } + EP test ∥ℓ − uci∗ ∥22 1{Acδ }

P test (Acδ ). ≤ O N 2 T − poly(N,d)δ + O K−1 K

1 K−1 Since K−1 K ≥ 2 for every K ≥ 2, the first term also carries the factor K after absorbing the constants and a further factor of two. Taking the infimum over δ > 0 gives the claim.

√ 2

Lemma 14 (Argmax margin). For any ℓ ∈ RK and any class c ∈ [K], if ∥ℓ − uc ∥2 < 12 then arg maxk∈[K] [ℓ]k = c. Proof. Because ∥·∥∞ ≤ ∥·∥2 , the hypothesis gives [ℓ]k − [uc ]k < 12 for every coordinate k. The 1 code uc has entry [uc ]c = K−1 K and [uc ]k = − K for k ̸= c, so 1 [ℓ]c > K−1 K − 2,

[ℓ]k < − K1 + 12

(k ̸= c).

Subtracting, for every k ̸= c, [ℓ]c − [ℓ]k >

K−1 1 1 1 K − 2 − − K + 2 1 = K−1 K + K − 1 = 0.

Thus [ℓ]c strictly exceeds every other coordinate, so the argmax is c. 18

Proof of Corollary 1. Taking the contrapositive of Lemma 14 with c = ci∗ , ĉWT (xN +1 ) ̸= ci∗ ⊆ ∥ℓ − uci∗ ∥2 ≥ 12 .

By Markov’s inequality, P test ĉWT (xN +1 ) ̸= ci∗

≤ P test ∥ℓ − uci∗ ∥22 ≥ 14

≤ 4 EP test ∥ℓ − uci∗ ∥22 .

2 − poly(N,d)δ + Applying Theorem 2 and noting K−1 K ≤ 1 bounds the right-hand side by O inf δ>0 {N T test c P (Aδ )} , which is the first claim. claim, suppose P test (Aδ∗ ) = 1. Then (20) holds almost surely, so ∥ℓ − uci∗ ∥2 ≤ √ For the second ∗ 1 T ∗ 1 T 2 N exp(− 2 ξ1 δ ) + exp(− 15 ξ2 ) = O N T − poly(N,d)δ almost surely. By Lemma 14, the argmax is correct as soon as this bound falls below 12 , i.e. once log T ≥ O log(KN )/(poly(N, d) δ ∗ ) .

Auxiliary Sphere-Geometry Results

The convergence proof has several facts about the uniform distribution on S d−1 . None of them involve the labels, so they are identical to their counterparts in the binary analysis. We state them here for completeness and refer to (Li et al., 2024) for the proofs, which are unchanged. Lemma 15 (Rotational invariance: (Li et al., 2024, Lemmas 6–7)). If {xi }i∈[N +1] are i.i.d. uniform on S d−1 , then their joint law is invariant under any orthogonal matrix U. Consequently, if W is diagonal with x-block ξ1 Id , the gradient ∇W11 L(W) is a scalar multiple of Id . Lemma 16 (Inner-product density: (Li et al., 2024, Lemma 17)). Let τ = ⟨x, x′ ⟩ for a fixed unit d−3 vector x and a uniform x′ ∈ S d−1 . Then τ has density fτ (t) = kd (1 − t2 ) 2 on [−1, 1], where 2 Γ( d2 ) kd := √ . π Γ( d−1 2 ) Lemma 17 (Expected nearest-neighbor alignment: (Li et al., 2024, Lemma 18)). If N ≥ √ 2 , where xi∗ = arg maxi∈[N ] ⟨xi , xN +1 ⟩. O( d log d), then E[⟨xi∗ , xN +1 ⟩] ≥ (N + 1)2 Lemma 18 (Concentration of the nearest-neighbor alignment: (Li et al., 2024, Lemma 19)). With xi∗ as above and kd as in Lemma 16,

P ⟨xi∗ , xN +1 ⟩ ≤ 1 − (2N kd )−2/(d−3) ≥ 1e . √ Equivalently, with the constant aN,d = (2N d)−2/(d−3) used in Lemmas 9 and 11, the gap satisfies 1 − ⟨xi∗ , xN +1 ⟩ ≥ aN,d with probability at least 1e ; that is, aN,d lower bounds the separation of the √ nearest neighbor from the query with constant probability. Note kd = Θ( d), so the two forms agree up to constants. D

Lemma 19 (Order-statistic gap: (Li et al., 2024, Lemma 20)). Let ⟨xi∗ , xN +1 ⟩ and x(2) , xN +1 be the largest and second-largest of {⟨xi , xN +1 ⟩}i∈[N ] . Then, for ξ > 0, h

E exp ξ( x(2) , xN +1 − ⟨xi∗ , xN +1 ⟩) N 2 kd2 =O ξ

and 19

1 . ξ

=Ω

Compute Environment

All experiments were conducted on a single workstation running Linux (kernel 6.8.0-58-generic) with glibc 2.39. The system was equipped with a 16-core (32-thread) CPU, 128 GB of RAM, and one NVIDIA RTX 5000 Ada Generation GPU.

One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context

Analyze & explore related research in Conceptio

Related documents