Table of Contents
Fetching ...

The Two-Stage Decision-Sampling Hypothesis: Understanding the Emergence of Self-Reflection in RL-Trained LLMs

Zibo Zhao, Yuanting Zha, Haipeng Zhang, Xingcheng Xu

TL;DR

This work tackles how self-reflection emerges in RL-trained LLMs by introducing the Gradient Attribution Property and the Two-Stage Decision-Sampling (DS) Hypothesis, which decomposes a unified policy into $π_{ ext{sample}}$ for generation and $π_d$ for verification. It shows that surrogate rewards produce Balanced Gradient Attribution, while KL penalties induce Unbalanced Attribution, explaining why RL can strengthen the decision-making component $π_d$ more effectively than sampling. Empirical validation on arithmetic reasoning tasks demonstrates that RL's generalization largely stems from improved $π_d$, with $p_d|W$ remaining substantial under RL but collapsing under SFT, supporting the claim that self-correction arises from better verification rather than generation. The framework further explains phenomena such as SFT's echo chamber, memorization vs generalization, the benefits of Dynamic Fine-Tuning (DFT), and the power of negative feedback for learning when to revise outputs. Overall, the results offer a principled, mechanistic lens for designing and evaluating training objectives to cultivate genuine self-reflection in thinking models.

Abstract

Self-reflection capabilities emerge in Large Language Models after RL post-training, with multi-turn RL achieving substantial gains over SFT counterparts. Yet the mechanism of how a unified optimization objective gives rise to functionally distinct capabilities of generating solutions and evaluating when to revise them remains opaque. To address this question, we introduce the Gradient Attribution Property to characterize how reward gradients distribute across policy components, formalized through the Two-Stage Decision-Sampling (DS) Hypothesis, which decomposes the policy into sampling ($π_{sample}$) for generation and decision ($π_{d}$) for verification. We prove that surrogate rewards exhibit Balanced Gradient Attribution, while SFT and KL penalties exhibit Unbalanced Gradient Attribution, with length-weighting creating asymmetric regularization that constrains $π_{sample}$ while leaving $π_{d}$ under-optimized, providing an theoretical explanation of why RL succeeds where SFT fails. We also empirically validate our theoretical predictions on arithmetic reasoning demonstrates that RL's superior generalization stems primarily from improved decision-making ($π_{d}$) rather than sampling capabilities, providing a first-principles mechanistic explanation for self-correction in thinking models.

The Two-Stage Decision-Sampling Hypothesis: Understanding the Emergence of Self-Reflection in RL-Trained LLMs

TL;DR

This work tackles how self-reflection emerges in RL-trained LLMs by introducing the Gradient Attribution Property and the Two-Stage Decision-Sampling (DS) Hypothesis, which decomposes a unified policy into for generation and for verification. It shows that surrogate rewards produce Balanced Gradient Attribution, while KL penalties induce Unbalanced Attribution, explaining why RL can strengthen the decision-making component more effectively than sampling. Empirical validation on arithmetic reasoning tasks demonstrates that RL's generalization largely stems from improved , with remaining substantial under RL but collapsing under SFT, supporting the claim that self-correction arises from better verification rather than generation. The framework further explains phenomena such as SFT's echo chamber, memorization vs generalization, the benefits of Dynamic Fine-Tuning (DFT), and the power of negative feedback for learning when to revise outputs. Overall, the results offer a principled, mechanistic lens for designing and evaluating training objectives to cultivate genuine self-reflection in thinking models.

Abstract

Self-reflection capabilities emerge in Large Language Models after RL post-training, with multi-turn RL achieving substantial gains over SFT counterparts. Yet the mechanism of how a unified optimization objective gives rise to functionally distinct capabilities of generating solutions and evaluating when to revise them remains opaque. To address this question, we introduce the Gradient Attribution Property to characterize how reward gradients distribute across policy components, formalized through the Two-Stage Decision-Sampling (DS) Hypothesis, which decomposes the policy into sampling () for generation and decision () for verification. We prove that surrogate rewards exhibit Balanced Gradient Attribution, while SFT and KL penalties exhibit Unbalanced Gradient Attribution, with length-weighting creating asymmetric regularization that constrains while leaving under-optimized, providing an theoretical explanation of why RL succeeds where SFT fails. We also empirically validate our theoretical predictions on arithmetic reasoning demonstrates that RL's superior generalization stems primarily from improved decision-making () rather than sampling capabilities, providing a first-principles mechanistic explanation for self-correction in thinking models.
Paper Structure (71 sections, 10 theorems, 104 equations, 4 figures, 7 tables)

This paper contains 71 sections, 10 theorems, 104 equations, 4 figures, 7 tables.

Key Result

Lemma 1

The probability of trajectory $\tau$ under policy $\pi_\theta$ factorizes as:

Figures (4)

  • Figure 1: Single Parameter Illustration
  • Figure 2: (RL Thinking Model): Calibrated policy parameters across arithmetic task sizes. $p_s$: sampling accuracy approximated by accuracy of $1st$ try, $p_{d|C}$: probability of stopping conditional on getting correct answer, $p_{d|W}$: probability of resample conditional on getting incorrect answer.
  • Figure 3: (SFT Thinking Model): Calibrated policy parameters across arithmetic task sizes. $p_s$: sampling accuracy approximated by accuracy of $1st$ try, $p_{d|C}$: probability of stopping conditional on getting correct answer, $p_{d|W}$: probability of resample conditional on getting incorrect answer.
  • Figure 4: chu2025sftmemorizesrlgeneralizes: ID/OOD performance comparison between SFT/RL

Theorems & Definitions (18)

  • Lemma 1
  • Corollary 1
  • Lemma 2
  • Definition 3.1: Gradient Attribution Property
  • Remark : perfect attribution
  • Remark : When imperfect attribution arise?
  • Remark
  • Theorem 3.1
  • Theorem 3.2
  • Remark
  • ...and 8 more