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

How reinforcement learning after next-token prediction facilitates learning

ICLR 2026 · Large language models › Reasoning and chain of thought · approach: RL with verifiable rewards for reasoning

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.

Theoretically and empirically shows that RL after next-token pretraining on a mixture of short and rare long chain-of-thought sequences can learn parity efficiently, where pretraining alone fails, via length-increasing reweighting.

Problem

Why next-token prediction alone fails to learn hard tasks (e.g., parity of d bits) when long chain-of-thought demonstrations are rare in pre-training data, and how RL after pre-training succeeds and increases response length.

Approach

Models a pre-training mixture of short and long chain-of-thought sequences; shows next-token prediction stays length-calibrated and fails under greedy decoding when long data fraction is below 1/3, while STaR-style RL reweights toward long correct responses with odds doubling each round, giving logarithmic rounds to generalization.

Key results

• Pre-training with next-token prediction on parity (d=50, p_cot=0.25) stays near 50% greedy accuracy, while GRPO after pre-training reaches ~100% accuracy after a handful of RL samples.
• Theory: next-token pretraining fails under greedy decoding when p_cot < 1/3; STaR reaches perfect generalization in O(log((1-p)/p)) rounds.
• Llama 3.2 3B on GSM8K with p_cot=0.05: GRPO after SFT rapidly reaches full-CoT SFT baseline accuracy while response length increases.

Surprising findings

• Pre-training accuracy at temperature 1 equals (p_cot+1)/2 while greedy decoding stays at chance, because the model is length-calibrated but fails on short responses.
• Pre-training failure on parity is due to sample complexity rather than model capacity; smaller d eventually generalizes with enough iterations.

Limitations

• Theory uses a linear autoregressive model and an idealized parity task.
• Long, correct chain-of-thought assumed present in pre-training data with small but nonzero probability.
• Some GSM8K/MATH gains may stem from out-of-distribution pretraining knowledge, confounding attribution.

Experimental setup

Datasets: synthetic parity sequences multiplication traces (4x4, 5x5, 7x7) GSM8K MATH

Benchmarks: GSM8K MATH parity (synthetic) multiplication (synthetic)

Metrics: test accuracy under greedy decoding response length accuracy at temperature 1

Baselines: next-token prediction only (pre-training) full chain-of-thought SFT

Builds on: STaR REINFORCE GRPO chain-of-thought

Closest papers by meaning

• The Pitfalls of Next-Token Prediction (ICML 2024 · similarity 0.89)
• Reinforcement Mid-Training (ICLR 2026 · similarity 0.89)
• Well Begun, Half Done: Reinforcement Learning with Prefix Optimization for LLM Reasoning (AAAI 2026 · similarity 0.89)
• Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs (ICLR 2026 · similarity 0.89)
• TTRL: Test-Time Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
• Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning (NeurIPS 2025 · similarity 0.89)