From P(y|x) to P(y): Investigating Reinforcement Learning in Pre-train Space
| Authors | Yuqiao Tan et al. |
| Year | 2026 |
| HF Upvotes | 26 |
| arXiv | 2604.14142 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
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.
Engineering Breakdown
Plain English
This paper addresses a fundamental limitation in reinforcement learning for large language models: while RLVR (reinforcement learning with verifiable rewards) improves reasoning by optimizing P(y|x), it's constrained by what the base model already knows. The authors propose PreRL, which instead optimizes the marginal distribution P(y) directly during pre-training using reward signals, letting the model explore and learn new reasoning patterns rather than just recombine existing ones. They prove mathematically that gradients from log P(y) align strongly with log P(y|x), meaning you can improve general reasoning ability while still maintaining conditional correctness. The key insight is that applying RL in pre-training space—before conditioning on inputs—unlocks fundamentally better reasoning capabilities than post-hoc fine-tuning approaches.
Core Technical Contribution
The core novelty is decoupling RL optimization from the conditional probability framework and applying it to the marginal distribution P(y) in pre-training space rather than P(y|x) in fine-tuning space. The authors establish theoretical gradient alignment between ∂log P(y) and ∂log P(y|x), proving that optimizing the unconditional distribution doesn't degrade conditional performance and can actually enhance it. This moves beyond static corpus-based pre-training toward active, reward-driven pre-training where the model's base output distribution itself evolves. The technical contribution is showing that this distribution shift—normally considered harmful—can be weaponized to inject reasoning capability into the model's priors rather than layering it on top through fine-tuning.
How It Works
PreRL operates in three stages: first, take a pre-trained base model and its marginal output distribution P(y). Second, collect or generate reasoning trajectories with verifiable rewards (e.g., math problems with correct/incorrect solutions), then compute reward signals for entire outputs. Third, apply gradient updates directly to model parameters using ∂log P(y) weighted by rewards, rather than conditioning on input x. The mechanism treats the unconditional generation problem as a Markov decision process where the 'action' is the full output token sequence, and the 'reward' comes from task verification. Unlike standard RLHF which fine-tunes P(y|x) and requires input-output pairs, PreRL modifies base distributions through online updates, allowing the model to internalize which types of outputs are generally valuable. The gradient alignment theorem ensures these updates to P(y) propagate beneficial signals to P(y|x) without collision.
Production Impact
Adopting PreRL would fundamentally change how you structure LLM training pipelines. Instead of: base model → RLHF fine-tuning → deployment, you'd now do: base model → PreRL pre-training with reward signals → optional conditional fine-tuning → deployment. This eliminates the ceiling imposed by base model capacity—you're directly growing reasoning ability into the model rather than constraining reasoning to recombinations of what it already knows. In production, this means smaller models trained with PreRL may match larger models trained conventionally, reducing inference costs and latency. The trade-off is that pre-training RL requires on-policy data collection and verification infrastructure at scale (you need reward signals for many diverse outputs), making it compute-intensive upfront; you'll need reward models or external verifiers, adding orchestration complexity. For reasoning-heavy workloads (coding, math, science), the improved base model quality likely justifies the extra compute; for simple classification, marginal gains may not warrant the effort.
Limitations and When Not to Use This
The paper assumes access to high-quality, scalable reward signals during pre-training—a major practical constraint rarely satisfied outside narrow domains like math or code verification. The theoretical gradient alignment holds asymptotically but doesn't account for optimization landscape challenges, local minima, or the actual convergence speed in large-scale settings; empirical validation on models >1B parameters with diverse task sets is needed. Distribution shift caused by RL updates could destabilize multi-task performance if not carefully regularized, and the paper doesn't deeply explore catastrophic forgetting or task interference when pre-training on multiple reward signals simultaneously. The approach also assumes static reward definitions; it's unclear how PreRL handles distribution shifts in real-world reward functions or adversarial scenarios where reward signals become misaligned with true downstream objectives.
Research Context
This work extends the RLHF paradigm (Ouyang et al., Christiano et al.) by moving optimization from the conditional distribution to the marginal, addressing a recognized bottleneck in reasoning improvement. It builds on recent work showing that pre-training objectives significantly constrain fine-tuning capacity, and proposes active rather than passive pre-training. The paper sits at the intersection of offline RL (using pre-trained models as behavior policies) and online RL (collecting reward signals dynamically), opening a new frontier in 'reward-driven pre-training.' This likely influences future work on curriculum learning during pre-training, multi-objective RL in base model training, and efficient reasoning capability injection without trillion-parameter scaling.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::