Research map / Large language models / Reasoning and chain of thought

Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis

ICLR 2025 · Large language models › Reasoning and chain of thought · approach: Chain-of-thought training and analysis · 0 citations (OpenAlex)

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.

Theoretical analysis of training a one-layer nonlinear (softmax) attention-only Transformer on synthetic K-step reasoning tasks, giving sample/iteration bounds and generalization guarantees for chain-of-thought (CoT) and ICL at inference on distribution-shifted tasks.

Problem

Why a Transformer trained by gradient methods can acquire generalizable multi-step chain-of-thought reasoning ability, and when CoT outperforms in-context learning (ICL), with nonconvex training of nonlinear attention.

Approach

Analyzes SGD training dynamics of a one-layer single-head softmax attention-only Transformer on prompts with step-wise examples, proving the attention concentrates on same-step, same-pattern context examples; derives CoT generalization bounds under noisy, erroneous testing examples and shows ICL additionally requires a dominant-correct-label condition.

Key results

• On a 3-step synthetic task with d=30, M=20, M'=10, one-layer Transformer CoT error decreases as the number of testing examples grows, with more examples needed when alpha', tau, or rho is small.
• On real arithmetic with GPT-4, CoT accuracy stays 100% with up to 2 implementation-error examples while ICL drops to 60%.
• Attention on same-step, same-pattern context columns rises to about 1 during training, with a two-stage dynamic where different-step attention decays to zero first.

Measured effects reported in the paper

Surprising findings

• With erroneous examples in the prompt, GPT-4 CoT accuracy drops to 0% once all three context examples carry an error in the o3 operator, showing sensitivity to example errors.

Limitations

• Analysis restricted to one-layer single-head attention-only Transformers with simplified orthonormal patterns and fixed value matrix.
• Positional encoding is simplified to orthonormal vectors.
• Real-world validation is limited to a small custom arithmetic task with 50 prompts per setting.

Experimental setup

Datasets: synthetic orthonormal-pattern reasoning tasks

Benchmarks: synthetic pattern reasoning tasks arithmetic chain of thought (4-operand expression, custom)

Metrics: CoT generalization error ICL generalization error accuracy

Baselines: ICL

Builds on: in-context learning theory with one-layer attention CoT-I/O inference scheme Transformer training dynamics analyses

Closest papers by meaning

• From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency (ICLR 2025 · similarity 0.91)
• How Does Chain of Thought Think? Mechanistic Interpretability of Chain-of-Thought Reasoning with Sparse Autoencoding (AAAI 2026 · similarity 0.91)
• Eliciting Chain-of-Thought in Base LLMs via Gradient-Based Representation Optimization (AAAI 2026 · similarity 0.91)
• Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thought (ICLR 2025 · similarity 0.91)
• CoT Vectors: Transferring and Probing the Reasoning Mechanisms of LLMs (ICLR 2026 · similarity 0.91)
• Enhancing Chain of Thought Prompting in Large Language Models via Reasoning Patterns (AAAI 2025 · similarity 0.91)