One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought
Abstract
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.
Related
- One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
- Diffusion Transformers are Provably Optimal In-context Generators
- Stochastic dominance of first return times for nearest-neighbor random walks on ℤd
- The nearest neighbor interchange proximity measure for the tree-basedness of rooted binary phylogenetic networks
- Vectorized Generalized Nearest Neighbor Decoding for In-block Memory Channel
BibTeX
@misc{wu2026onelayer,
title = {{One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought}},
author = {Lyumin Wu and Yuan Cao},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.16798},
eprint = {2610.16798},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.16798},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought
AU - Lyumin Wu
AU - Yuan Cao
PY - 2026
DA - 2026/10/12
PB - arXiv
DO - 10.48550/arxiv.2610.16798
UR - https://arxiv.org/abs/2610.16798
AB - 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.
ER -