Research map / Large language models / Reasoning and chain of thought
Transformers Provably Solve Parity Efficiently with Chain of Thought
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 theoretically that a one-layer softmax transformer trained by gradient descent can learn k-parity efficiently with chain-of-thought: impossible end-to-end without intermediate supervision, learnable in one step with teacher forcing, and in logarithmic steps with self-consistency filtering and data augmentation.
Problem
Understanding how transformers can be optimized with chain-of-thought (intermediate-state) supervision to solve complex compositional tasks, using k-parity as a test case where end-to-end gradient learning is provably hard.
Approach
Applies a one-layer transformer recursively to its own output over a binary tree decomposition of parity, with positional attention weights only; intermediate-state losses (teacher forcing) make parity learnable in one gradient step, and block-autoregressive masking, augmented random data as a self-consistency filter, and integer weight quantization allow learning without teacher forcing in log2 k
Key results
- Under finite-sample, polynomially bounded gradients, any iterative algorithm with an approximate gradient oracle needs exponentially many queries to learn parity when n=e^{Omega(d)} (Theorem on hardness).
- With teacher forcing, one gradient step with learning rate Theta(d^{2+eps/16}) from zero init yields ||y_test - y_test_hat||_inf <= O(d^{-eps/8}) for n=Omega(d^{2+eps}).
- Without teacher forcing but with self-consistency augmentation and weight quantization, the transformer reaches exponentially small loss exp(-Omega(d^{eps/16})) in log2 k iterations; experiments show direct and basic CoT models plateau at prediction loss 0.5 while CoT with teacher forcing and self-c
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train only on the final prediction instead of supervising intermediate parity states | one-layer transformer · 64-bit k-parity (k=32) | prediction loss: prediction error is stuck at 0.5 | Hurt |
| add intermediate-state loss without teacher forcing or self-consistency | one-layer transformer · 64-bit k-parity (k=32) | prediction loss: never improves beyond 0.5 | No clear effect |
| add teacher forcing to intermediate-state supervision | one-layer transformer · 64-bit k-parity (k=32) | prediction loss | Helped |
| add data augmentation and self-consistency filtering to end-to-end CoT | one-layer transformer · 64-bit k-parity (k=32) | prediction loss | Helped |
| add filtering to control error accumulation in end-to-end CoT | basic CoT model · 64-bit k-parity (k=32) | prediction loss | Helped |
Surprising findings
- Finite-sample parity becomes harder as the number of samples increases in the hardness result, because empirical correlations between parities do not vanish.
- The direct end-to-end transformer and basic CoT model fail on k=32 parity, with weights becoming nearly uniform and prediction loss stuck at 0.5.
Limitations
- Analysis restricted to a one-layer transformer with reparametrized attention weights and fixed positional attention (keys/queries zero on content).
- Results require k=Theta(d) or specific sample-size and learning-rate regimes; the non-teacher-forcing result relies on quantization and carefully designed filters and masks.
- Theorem bounds depend on an approximate gradient oracle model, and the hardness result has a caveat that the original adversarial perturbation statement holds only for certain schemes.
Experimental setup
Datasets: synthetic random k-parity samples (d=64 bits, k=8,16,32)
Benchmarks: synthetic k-parity (d=64, k=8,16,32)
Metrics: CoT loss prediction loss
Baselines: Direct end-to-end training Basic CoT without teacher forcing
Builds on: Wies23 (RNN parity with intermediate supervision) statistical query hardness of parity via gradient variance Transformer CoT expressivity (Feng et al.)
Closest papers by meaning
- Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization (NeurIPS 2025 · similarity 0.92)
- Grokking of Implicit Reasoning in Transformers: A Mechanistic Journey to the Edge of Generalization (NeurIPS 2024 · similarity 0.90)
- Chain of Thought Empowers Transformers to Solve Inherently Serial Problems (ICLR 2024 · similarity 0.90)
- One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought (NeurIPS 2026 · similarity 0.90)
- The Expressive Power of Transformers with Chain of Thought (ICLR 2024 · similarity 0.89)
- Multi-head Transformers Provably Learn Symbolic Multi-step Reasoning via Gradient Descent (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.