Research map / Architectures and efficiency / Attention and transformers

One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows a one-layer softmax transformer cannot solve odd-K nearest-neighbor by direct inference but can via chain-of-thought with iterative token selection and masking, and that gradient descent from zero init recovers this construction.

Problem

Whether one-layer transformers can learn in-context K-nearest-neighbor prediction (K>1), which direct inference cannot do for odd K, and how chain-of-thought enables it and training recovers it.

Approach

Reformulates K-NN as a sequential procedure: each CoT step softmax-attends to the next unretrieved nearest neighbor, using positional encodings to penalize previously selected tokens (via a sharp-attention limit), and proves gradient descent from zero initialization converges to this sharp-attention parameter pattern.

Key results

• Proposition 3.2: a one-layer transformer without CoT cannot solve odd K-NN (3<=K<=N-2) on a well-separated class for any fixed V, W.
• Theorem 4.1 / Prop 4.2: a one-layer transformer with sharp attention (alpha to infinity, beta/alpha to -2) solves K-NN via K-1 CoT steps with uniform correctness under a fixed margin.
• Theorem 5.2: GD from zero initialization on population loss gives W(t)=diag(alpha Id,0,beta I) with beta to -inf, alpha to +inf, beta/alpha to -2, and loss to zero; experiments show CoT accuracy approaches 1 for K=1,3,5 while direct inference baselines are substantially lower for K=3,5.

Surprising findings

• Direct inference by a one-layer softmax transformer provably fails for odd K>1 even on well-separated instances, while the same architecture succeeds with CoT.

Limitations

• Frozen value matrix V and simplified merged-key-query one-layer architecture in training analysis
• Population-loss training with one-step state supervision and synthetic data; not finite-sample
• Experiments use random initialization and SGD, deviating from the theorem setup

Experimental setup

Datasets: synthetic K-NN data on unit sphere with Rademacher labels

Benchmarks: synthetic well-separated K-NN test set (size 1000, margin 0.1)

Metrics: test accuracy training loss

Baselines: direct-inference one-layer transformer (K=3, K=5) Li et al. 2024a 1-NN construction

Builds on: Li et al. 2024a in-context 1-NN with one-layer transformer Wang et al. 2024 sparse token selection with stochastic positional encodings Huang et al. 2025a CoT linear regression Yang et al. 2025 CoT path finding

Closest papers by meaning

• One-Layer Transformer Provably Learns One-Nearest Neighbor In Context (NeurIPS 2024 · similarity 0.90)
• Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis (ICLR 2025 · similarity 0.90)
• Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought (ICLR 2025 · similarity 0.90)
• Transformers Provably Solve Parity Efficiently with Chain of Thought (ICLR 2025 · similarity 0.90)
• Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization (NeurIPS 2025 · similarity 0.89)
• Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent (NeurIPS 2025 · similarity 0.87)