Chain-of-Thought (CoT) enables transformers to solve complex reasoning tasks by generating intermediate reasoning steps. Despite strong empirical success, the mechanisms by which such reasoning abilities emerge during training remain poorly understood, particularly from a theoretical perspective. In this work, we study this question in a stylized yet fully analyzable setting: in-context K-nearest-neighbor (K-NN) prediction with a one-layer transformer. Earlier work has demonstrated that one-layer transformers can be trained to perform 1-NN without CoT (Li et al., 2024a) by having the softmax attention attend to the nearest neighbor in context. However, this result on 1-NN cannot be extended to K-NN for K>1, as softmax attention is much less naturally suited to attending to the 2nd through K-th nearest neighbors. We first give a rigorous negative result showing that this limitation already appears on a well-separated class of binary classification tasks for odd K>1. We then give an explicit construction of a one-layer transformer, and show that it can solve K-NN via CoT reasoning. We further prove that, under a stylized training setup, gradient descent can recover this construction, thereby showing that transformers can acquire the capability to solve K-NN through training. In addition, we show that the resulting trained model applies to a broader test class than that assumed in the training analysis. Our results shed light on how CoT expands the ability of transformers to learn and solve tasks that are otherwise hard to solve.
Figures & tables
Figure 1: Training loss, learned parameter heatmap, and test accuracy. The heatmap and test accuracy are shown for the trained transformer with d=10 and N=20 .
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 2: Joint training of W and V . We train both matrices jointly from random initialization. The learned V and W recover the construction in Theorem 4.1 , and the resulting transformer achieves high test accuracy on the K -NN task.
Model
Training / Inference
K=1
K=3
K=5
1-head, 1-layer
CoT / CoT rollout
100.0
99.8
98.3
Direct ( K=1 ) / Single-pass
100.0
–
–
Direct ( K=3 ) / Single-pass
100.0
74.7
–
Direct ( K=5 ) / Single-pass
99.6
–
68.8
1-head, 2-layer
CoT / CoT rollout
100.0
99.1
92.8
Direct ( K=1 ) / Single-pass
100.0
–
–
Appendix
Table 1: Test accuracy of CoT and direct models across different values of K .
Chain-of-Thought (CoT) prompting substantially improves the sample efficiency of transformers, reducing the complexity of tasks like parity learning from exponential to polynomial in the input length. However, generating explicit reasoning steps at inference is computationally expensive. Implicit Chain-of-Thought (ICoT) has emerged as a promising empirical remedy that trains models to internalize intermediate steps within their hidden states, but its theoretical foundations remain poorly understood. We give the first theoretical analysis of ICoT, proving that an L-layer transformer trained under our proposed Log-ICoT curriculum learns k-parity with poly(n) samples and L=log2k training stages. This matches the sample efficiency of explicit CoT while eliminating its inference overhead, and extends prior one-layer parity guarantees to multi-layer architectures. Compared to standard ICoT, which removes thinking tokens one at a time, Log-ICoT removes them in geometric chunks, reducing the number of stages from linear in k to logarithmic. Experiments on multi-layer transformers confirm the theory and visualize how reasoning is progressively absorbed into deeper layers.
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.
Skanda Athreya, Yutong Wang
James B. Conant High School · Hoffman Estates, IL · Illinois Institute of Technology +1
Modern transformers often solve complex tasks by generating and refining intermediate reasoning steps. To develop a theoretical understanding of chain-of-thought (CoT) reasoning, a notable recent study by Huang et al. (2025) viewed transformers as algorithm executors, and showed that when transformers' CoT steps are supervised by a gradient-descent teacher, they can learn to implement gradient descent for in-context weight prediction in linear regression. In this paper, we study the same in-context weight prediction task but ask whether CoT reasoning can go beyond merely imitating the teacher algorithm. We show that, after initial supervision from a gradient-descent teacher, further trajectory-level supervision using the true regression weights leads the model to implement the Heavy-Ball method, thereby achieving acceleration over the original teacher. Notably, this accelerated algorithm emerges without supervision from any Heavy-Ball teacher. Our results provide a concrete theoretical example in which transformer reasoning can discover new algorithms through training, rather than merely imitate the algorithm used for supervision.
Yuxuan Cai, Yuan Cao
School of Computing & Data Science, The University of Hong Kong