Outcome-Based RL Provably Leads Transformers to Reason, but Only With the Right Data

TL;DR

This paper demonstrates that outcome-based reinforcement learning enables Transformers to develop reasoning abilities, but only when trained on data emphasizing simpler reasoning tasks, with theoretical and empirical validation.

Contribution

It provides a theoretical analysis of how sparse rewards guide Transformers to learn systematic reasoning, highlighting the importance of training data distribution.

Findings

Policy gradient drives Transformers to learn iterative reasoning algorithms.

Simple examples in training data are critical for generalizable reasoning.

Training on insufficient simple examples hampers reasoning ability.

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 policy gradient to discover such systematic reasoning remains poorly understood. We address this by analyzing the policy gradient dynamics of single-layer Transformers on a synthetic graph traversal task that cannot be solved without Chain-of-Thought but admits a simple iterative solution. We prove that despite training solely on final-answer correctness, policy gradient drives the Transformer 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 examples, the Transformer learns a generalizable traversal strategy that extrapolates to longer chains; when this mass vanishes, policy gradient 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.