Outcome-Based RL Provably Leads Transformers to Reason, but Only With the Right Data
Yuval Ran-Milo, Yotam Alexander, Shahar Mendel, Nadav Cohen
TL;DR
This work analyzes how sparse terminal rewards in outcome-based RL guide gradient flow to induce multi-step reasoning in transformers, focusing on a synthetic chain-identification task. The authors develop a gradient-flow framework for single-layer transformers, proving that reasoning is necessary and sufficient for solving the task, and showing that the data distribution—specifically, nonzero mass on easy examples—is critical for efficiently learning a chain-traversal strategy and generalizing to longer chains. Experiments on synthetic data and fine-tuning an LLM on mathematical tasks validate the theory: models trained with simple examples learn efficient, interpretable reasoning and generalize to harder instances, while excluding simple examples prevents CoT emergence. The results illuminate how data composition and sparse rewards shape the emergence of chain-of-thought behavior, with practical implications for curriculum design and post-training data strategies that bootstrap complex reasoning in real-world models.
Abstract
Transformers trained via Reinforcement Learning (RL) with outcome-based supervision can spontaneously develop the ability to generate intermediate reasoning steps (Chain-of-Thought). Yet the mechanism by which sparse rewards drive gradient descent to discover such systematic reasoning remains poorly understood. We address this by analyzing the gradient flow dynamics of single-layer Transformers on a synthetic graph traversal task that cannot be solved without Chain-of-Thought (CoT) but admits a simple iterative solution. We prove that despite training solely on final-answer correctness, gradient flow drives the model to converge to a structured, interpretable algorithm that iteratively traverses the graph vertex-by-vertex. We characterize the distributional properties required for this emergence, identifying the critical role of "simple examples": instances requiring fewer reasoning steps. When the training distribution places sufficient mass on these simpler instances, the model learns a generalizable traversal strategy that extrapolates to longer chains; when this mass vanishes, gradient-based learning becomes infeasible. We corroborate our theoretical results through experiments on synthetic data and with real-world language models on mathematical reasoning tasks, validating that our theoretical findings carry over to practical settings.