PaperlayerView exact revision on arXiv

Source-linked AI summary

From $P(y|x)$ to $P(y)$: Investigating Reinforcement Learning in Pre-train Space

Yuqiao Tan, Minzheng Wang, Bo Liu, Zichen Liu, Tian Liang, Shizhu He, Jun Zhao, Kang Liu

arXiv:2604.14142v1cs.LGcs.AIcs.CL

TL;DR

The paper addresses the limitation that RLVR optimizes P(y|x) within the base model’s existing output distribution and introduces reward-driven updates to P(y) through PreRL. It validates gradient alignment between the two objectives, identifies NSR as an effective reasoning driver, and combines NSR-PreRL with standard RL in DSRL, which consistently outperforms strong baselines.

  • Problem

    RLVR’s reasoning gains are bounded by the base model’s output distribution, while static-corpus pre-training causes distribution shift that limits targeted reasoning enhancement.

  • Method

    PreRL optimizes the marginal policy P(y), and DSRL uses Policy Reincarnation to initialize with NSR-PreRL before applying standard RL to P(y|x).

  • Results

    DSRL consistently outperforms strong baselines across extensive experiments, achieving superior accuracy, better sample efficiency, and robust Pass@K improvements.

  • Takeaways & Limitations

    Pre-train-space pruning steers the policy toward a refined correct reasoning subspace while NSR stimulates endogenous reasoning capabilities.

  • Takeaways & Limitations

    PSR-PreRL can degrade performance through on-policy learning collapse, and maximizing P(y) may require high-quality out-of-distribution expert demonstrations.

Abstract

from arXiv · show

While reinforcement learning with verifiable rewards (RLVR) significantly enhances LLM reasoning by optimizing the conditional distribution P(y|x), its potential is fundamentally bounded by the base model's existing output distribution. Optimizing the marginal distribution P(y) in the Pre-train Space addresses this bottleneck by encoding reasoning ability and preserving broad exploration capacity. Yet, conventional pre-training relies on static corpora for passive learning, leading to a distribution shift that hinders targeted reasoning enhancement. In this paper, we introduce PreRL (Pre-train Space RL), which applies reward-driven online updates directly to P(y). We theoretically and empirically validate the strong gradient alignment between log P(y) and log P(y|x), establishing PreRL as a viable surrogate for standard RL. Furthermore, we uncover a critical mechanism: Negative Sample Reinforcement (NSR) within PreRL serves as an exceptionally effective driver for reasoning. NSR-PreRL rapidly prunes incorrect reasoning spaces while stimulating endogenous reflective behaviors, increasing transition and reflection thoughts by 14.89x and 6.54x, respectively. Leveraging these insights, we propose Dual Space RL (DSRL), a Policy Reincarnation strategy that initializes models with NSR-PreRL to expand the reasoning horizon before transitioning to standard RL for fine-grained optimization. Extensive experiments demonstrate that DSRL consistently outperforms strong baselines, proving that pre-train space pruning effectively steers the policy toward a refined correct reasoning subspace.

1 Introduction

The paper introduces reward-driven reinforcement learning in the pre-train space to optimize P(y), address static-corpus distribution shift, and improve subsequent conditional-policy optimization. It identifies NSR as a mechanism for pruning incorrect reasoning and combines NSR-PreRL with standard RL in DSRL.

  • Motivation: Static-corpus pre-training creates distribution shift from pre-trained knowledge to post-training task distributions, hindering targeted reasoning enhancement.The paper motivates reward-driven updates in the pre-train space as a way to combine knowledge internalization with task-directed optimization.
  • Pre-train Space RL: Pre-train Space RL optimizes the marginal distribution P(y) within reinforcement learning rather than only the conditional policy P(y|x).The approach targets the model’s broader output distribution while preserving its exploration capacity.
  • Pre-train Space RL: Gradient alignment between log P(y) and log P(y|x) supports PreRL as a surrogate for standard RL while retaining broad exploration capacity.The paper provides both theoretical justification and empirical verification for this alignment.
  • Negative Sample Reinforcement: NSR-PreRL prunes incorrect trajectories and increases transition and reflection thoughts by 14.89× and 6.54×, respectively.The mechanism contrasts with PSR and stimulates endogenous reasoning behaviors through richer probabilistic exploration.
  • Dual Space RL: DSRL initializes models with NSR-PreRL to expand reasoning capabilities before switching to standard RL for direct policy optimization.This Policy Reincarnation strategy combines pre-train-space exploration with post-train-space refinement.
  • Results: DSRL consistently outperforms strong baselines with superior accuracy, better sample efficiency, and robust Pass@K improvements.The paper presents this as evidence that pre-train-space learning steers policies toward a refined correct reasoning subspace.

2 From Post-train Space P(y|x) to Pre-train Space P(y)

The section contrasts conditional post-train optimization of P(y|x) with PreRL’s marginal optimization of P(y), then analyzes their alignment and training dynamics. It finds strong gradient alignment, but shows that positive and negative sample reinforcement behave differently, motivating negative-sample initialization for later RL.

  • Gradient Alignment: Conditional and marginal log-probabilities often share token rankings, but diverge for early-sequence or highly uncertain tokens.The distributions align most closely for high-probability, deterministic tokens.
  • Post-train Space Optimization: Standard RL optimizes the conditional policy πθ(y|x), updating token generation while conditioning on the input question.The policy operates over states containing the question and previously generated tokens, with sparse task-success rewards.
  • Pre-train Space Optimization: PreRL directly optimizes the marginal policy πθ(y), removing input dependency from the gradient update while preserving sequence-level pre-training behavior.This formulation is proposed as a pre-train-space alternative to conditional policy optimization.
  • Gradient Alignment: Aligned gradients imply that a PreRL update can improve conditional log πθ(y|x) when the marginal and conditional gradients share direction.The first-order analysis identifies their inner product as the cross-gradient term governing the conditional change.
  • Gradient Alignment: 100% of AMC23 samples had non-negative marginal–conditional gradient inner products, with an average inner product of +9.2.The accompanying cosine-similarity analysis was strongly positive, supporting alignment between the two objectives on Qwen3-4B.
  • Training Dynamics: PSR-PreRL initially raises conditional probability but eventually collapses on self-generated on-policy trajectories, unlike teacher-trained QFFT.The comparison suggests that maximizing πθ(y) may require high-quality out-of-distribution expert demonstrations or off-policy data.
  • Training Dynamics: After 20 NSR-PreRL steps, transition thoughts increased 14.89× and reflection thoughts 6.54× versus the vanilla model, while accuracy reached 86% with 3× fewer steps than standard RL.NSR suppresses incorrect trajectories and elicits longer, deeper reasoning, though excessively long outputs can eventually hinder continued training.

3 Dual Space RL: Combining PreRL and RL via Policy Reincarnation

DSRL combines an initial NSR-PreRL phase with standard RL through a policy-reincarnation transition, using pre-train-space pruning before conditional optimization. This framework is evaluated through performance, reasoning-behavior, and problem-solving-status comparisons.

  • Policy Reincarnation: DSRL replaces the base model midway with an NSR-PreRL checkpoint, then restarts on-policy RL to unify pre-train and post-train optimization.The transition is controlled by a threshold S: before S, inputs are masked and negative samples drive pruning; after S, standard RL uses all samples.
  • Policy Reincarnation: During the initial phase, the indicator constraint applies updates only to negative samples, pruning incorrect reasoning paths before standard RL begins.After the transition threshold, conditioning is restored and standard RL utilizes all samples.
  • Performance Evaluation: Figure 4 compares DSRL and GRPO using Pass@K performance across LLMs.The supplied caption identifies the comparison, while the section’s cited results describe DSRL as outperforming GRPO across most K budgets.
  • Reasoning Behaviors: Figure 5 tracks reasoning behaviors over RL training with different methods, including Subgoal Setting, Enumeration, Verification, and Backtracking.The cited results report faster and more sustained behavior growth for DSRL than RL across these patterns.
  • Problem-Solving Status: Figure 6 tracks Fully Solved and Fully Unsolved questions during training to assess foundational problem-solving ability.The cited results report a sharp increase in Fully Solved counts and a decline in Fully Unsolved counts during NSR-PreRL.

4 Experiments

Experiments show that DSRL improves accuracy, training efficiency, reasoning-behavior emergence, and Pass@K scalability over strong RL baselines. Its gains depend on NSR-PreRL initialization and a carefully chosen 10–25-step warmup.

  • DSRL consistently improves Avg@K over strong RL baselines across all evaluated benchmarks and model scales.
  • DSRL outperforms GRPO across the entire Pass@K range, indicating stronger sampling scalability and a more exploration-friendly policy landscape.
  • DSRL generalizes beyond mathematics, improving GPQA-Diamond by +3.79 and MMLU-Pro by +5.37 for Qwen3-4B, while HumanEval gains +2.44 for Qwen3-8B.
  • Training Efficiency in Dual-Space RL: 61.6 average accuracy versus 57.7 for GRPO, with 45.0% and 58.0% accuracy reached using 2.5× and 1.6× fewer steps, respectively.
  • Evolution of Reasoning Behaviors: DSRL produces faster growth and higher frequency ceilings for subgoal setting, enumeration, verification, and backtracking than RL.
  • Foundation Reasoning Enhanced via NSR-PreRL: NSR-PreRL rapidly increases Fully Solved questions, after which standard RL refines problem-specific nuances and yields more Fully Solved questions with lower error.
  • Impact of Warmup Steps: An inverted-U warmup trend makes 10–25 steps optimal; fewer steps provide insufficient stimulation, while excessive warmup hinders later fine-grained optimization.
  • Pre-train Space vs. Post-train Space Warmup: With 20-step warmup, DSRL scores 57.54 versus 55.79 for GRPO and 54.38 for NSR-RL, while NSR-PreRL leads NSR-RL by 6.6 average points by step 20.

5 Related Works

Related work contrasts post-training RL and pre-train-space optimization. The paper positions PreRL against static, passive pre-training approaches by adding reward-guided online learning over the marginal distribution.

  • PreRL is distinguished from these approaches by integrating reward-driven reinforcement learning into the pre-train space.
  • RLVR enhances LLM reasoning through post-training, while related work also examines how its distribution sharpening may constrain exploration.
  • Continual pre-training encodes broad knowledge by optimizing P(y) on static corpora, providing a stronger downstream foundation but retaining passive data learning.

6 Conclusion

The paper concludes that PreRL and NSR provide a stronger reasoning foundation for DSRL, which combines pre-train-space initialization with standard RL and outperforms strong baselines.

  • PreRL aligns marginal and conditional objectives, while NSR prunes incorrect paths and elicits endogenous reasoning capabilities.
  • DSRL uses Policy Reincarnation to combine NSR-PreRL initialization with standard RL for stronger final reasoning performance.
  • DSRL consistently outperforms strong baselines across extensive benchmarks, supporting pre-train-space optimization as a robust reasoning foundation.

A.1 Comparison of PreRL with Pre-training and Continual Pre-training

Pre-training, continual pre-training, and PreRL all optimize P(y), but differ in data sources, learning signals, and training paradigms. PreRL replaces passive corpus learning with online reward-guided self-rollouts while retaining the pre-train-space objective.

  • All three paradigms optimize the marginal distribution P(y), but differ in data sources, learning signals, and training paradigms.
  • Pre-training: Pre-training uses large-scale web corpora and next-token prediction to establish broad language capabilities, but its distribution is not aligned with downstream reasoning tasks.
  • Continual Pre-training: Continual pre-training narrows the distribution gap with task-specific corpora, yet remains limited by static, pre-collected data and passive next-token prediction.
  • PreRL: PreRL generates reasoning trajectories online through self-rollouts, uses verifiable rewards to reinforce or suppress them, and updates only on response y with x removed.

A.2 Comparison with Reinforcement Pre-Training Paradigm

The section distinguishes PreRL from RLPT by applying reinforcement learning to reasoning tasks with verifiable rewards while optimizing the pre-train-space marginal distribution P(y).

  • Paradigm comparison: RLPT methods apply RL objectives during pre-training, generally using static pre-training corpora.They share the goal of extending RL beyond standard post-training to strengthen foundational reasoning.
  • Paradigm comparison: RLPT commonly converts Next-Token Prediction into a pseudo-reward because raw corpora lack objective verification signals.Self-generated rollouts are rewarded according to whether they predict subsequent ground-truth tokens.
  • PreRL: PreRL operates directly on reasoning tasks within the standard RLVR framework and uses objective task success as the verifiable reward.This replaces static dataset continuations and NTP pseudo-rewards with task-level online exploration.
  • PreRL: PreRL restricts policy updates to the pre-train space by masking the question during gradient updates, thereby optimizing the marginal distribution P(y).This is the central optimization-space distinction from conventional post-train-space RL.

B.1 Training and Evaluation Setup

Experiments use Qwen3 models trained on MATH with fixed rollout and sampling settings, then evaluate performance using high-sample reasoning metrics and specified templates.

  • Training: Training uses Qwen3-4B and Qwen3-8B base models on the 7,500-problem MATH dataset.Models are trained with the verl framework, prompt batch size 128, and eight rollouts per prompt.
  • Training: Training samples use temperature 1.0 and a maximum response length of 8192 tokens.The passage also specifies a mini-batch size of 128 for model updates.
  • Evaluation: Evaluation uses vLLM with temperature=1.0 and top p=1.0, reporting Avg@32 because reasoning outputs have high variance.Pass@K is estimated from 300 generated responses to reduce evaluation variance.
  • Evaluation: All Qwen experiments use a Qwen3 non-thinking prompt template.The template is identified as the standard format used for the experiments.

B.3 Implementation of Reflection, Transition, and Execution Thoughts

The implementation classifies generated reasoning steps as reflection, transition, or execution using case-insensitive prefix and phrase heuristics.

  • Step taxonomy: Steps beginning with keywords such as “wait” or “alternatively” are classified as transition or reflection thoughts.This is the prefix-based classification rule.
  • Step taxonomy: Steps containing predefined mid-sentence phrases indicating self-correction or redirection are likewise classified as transition or reflection thoughts.This is the phrase-based classification rule.
  • Step taxonomy: Steps matching neither heuristic are classified as execution thoughts.Classification is case-insensitive, and category proportions are computed across all generated-response steps.

B.4 Implementation of Reasoning Behaviors

Reasoning behaviors are evaluated from generated chains of thought using GPT-4o annotations and predefined behavioral categories, with figures comparing distributions and training dynamics.

  • Behavior annotation: GPT-4o identifies four reasoning behaviors in model-generated responses: subgoal setting, enumeration, verification, and backtracking.The analysis follows a prior taxonomy and reports the proportion of responses exhibiting each behavior throughout training.
  • Behavior definitions: Backtracking marks explicit error identification and revision of previously used methods.This category captures active correction during response generation.
  • Behavior definitions: Verification marks systematic checking of intermediate results for correctness.The category focuses on checking intermediate reasoning rather than only the final answer.
  • Behavior definitions: Subgoal setting marks decomposition of complex problems into smaller, manageable steps.This category captures structured problem decomposition.
  • Behavior definitions: Enumeration marks exhaustive consideration of multiple cases or possibilities.The category captures branching exploration during problem solving.
  • Behavior annotation: The evaluation prompt targets departures from typical linear, monotonic reasoning patterns in generated chains of thought.Figure 9 provides the system prompt used to evaluate these beneficial behaviors.
  • Diagnostic figures: Figures 10–13 compare conditioned and unconditioned token distributions, token-level log-probabilities, and NSR-PreRL versus NSR-RL warmup dynamics.The aligned and misaligned cases illustrate comparisons between P(y|x) and P(y) under the same context.
Loading 2604.14142v1…