All posts

research digesten May 16, 2026 6 min read

Beyond Conditional Optimization: The Case for Pre-train Space Reinforcement Learning

By kualia research desk

Introduction

The recent success of DeepSeek-R1 and similar systems has cemented reinforcement learning with verifiable rewards (RLVR) as the dominant paradigm for enhancing reasoning capabilities in large language models. By optimizing the conditional distribution $P(y|x)$ through reward signals, these methods effectively distill reasoning patterns from the base model's existing capacity. However, a fundamental constraint persists: the optimization is bounded by the base model's initial output distribution. If the desired reasoning trajectories do not exist with non-negligible probability in the base model, standard RLVR struggles to invent them ex nihilo.

In their paper "From $P(y|x)$ to $P(y)$: Investigating Reinforcement Learning in Pre-train Space," Tan et al. propose a provocative alternative that challenges the assumption that reinforcement learning must remain confined to conditional spaces. Their framework, Pre-train Space RL (PreRL), shifts the optimization target from the conditional distribution $P(y|x)$ to the marginal distribution $P(y)$. This seemingly subtle mathematical change carries profound implications for how we conceptualize knowledge internalization, exploration, and the very nature of reasoning acquisition in language models.

The Mechanics of Pre-train Space Optimization

The core insight of PreRL lies in recognizing that standard RLVR optimizes policy gradients within a constrained subspace defined by the input context. By contrast, PreRL removes the conditioning on the question $x$ and instead optimizes the marginal probability of generating particular outputs $y$ regardless of the specific prompt. The authors provide both theoretical justification and empirical verification that the gradients of $\log P(y)$ and $\log P(y|x)$ maintain strong alignment, establishing that optimizing the marginal distribution serves as a viable surrogate for standard conditional RL.

This alignment is not merely a mathematical curiosity; it has practical consequences for model behavior. When optimizing $P(y|x)$, the model tends to exploit existing patterns within the conditional space, potentially trapping the optimization in local optima defined by the base model's initial capabilities. Optimizing $P(y)$, however, operates at a more foundational level, effectively reshaping the model's generative prior. This approach preserves broader exploratory capacity because updates to the marginal distribution influence the model's behavior across all possible inputs, not just the current training batch.

The technical implementation involves a simple but consequential modification: during the policy gradient updates, the gradient estimator operates on the unconditional log-probabilities rather than the conditional ones. This requires careful handling of the baseline estimation and advantage calculation, as the standard RLVR frameworks assume a paired relationship between questions and answers. The authors demonstrate that despite this architectural shift, the variance of gradient estimates remains manageable, and the optimization trajectory remains stable.

The Asymmetry of Negative Sample Reinforcement

Perhaps the most counterintuitive finding in this work concerns the asymmetric roles of positive and negative samples within the pre-train space. In standard RLVR, both positive reinforcement (increasing the probability of correct answers) and negative reinforcement (decreasing the probability of incorrect answers) generally contribute to improved performance. However, Tan et al. discover that in the pre-train space, Positive Sample Reinforcement (PSR) actually degrades model performance by accumulating probability mass on self-generated correct samples in a way that collapses exploration. Conversely, Negative Sample Reinforcement (NSR) proves exceptionally instructive.

NSR-PreRL operates by aggressively pruning incorrect reasoning trajectories from the model's marginal distribution. The mechanism reveals something profound about the structure of reasoning: eliminating wrong paths appears more valuable than reinforcing correct ones when operating at the foundational level. The quantitative results support this interpretation, showing that NSR-PreRL increases transition thoughts by 14.89 times and reflection thoughts by 6.54 times compared to baseline approaches. These metrics indicate that NSR stimulates endogenous reflective behaviors, effectively broadening the search space through richer probabilistic exploration rather than narrowing it toward known solutions.

This finding resonates with broader principles in learning theory and cognitive science. The model appears to benefit from a form of "negative knowledge," understanding the boundaries of invalid reasoning before converging on valid strategies. In the context of mathematical reasoning or code generation, this suggests that identifying structural impossibilities or logical contradictions provides stronger learning signals than simply imitating successful patterns.

Dual Space RL and Policy Reincarnation

Leveraging these insights, the authors propose Dual Space RL (DSRL), a staged training strategy that combines the strengths of pre-train and post-train space optimization. Drawing on the concept of Policy Reincarnation, DSRL begins with an NSR-PreRL phase to expand the model's reasoning horizon and elicit intrinsic reasoning capabilities. This initial phase effectively prunes the vast space of incorrect reasoning patterns, steering the policy toward a refined correct reasoning subspace. Subsequently, the training transitions to standard RLVR operating on $P(y|x)$ for fine-grained optimization of conditional trajectories.

The empirical results demonstrate that DSRL consistently outperforms strong baselines including GRPO and other RLVR variants. Beyond raw accuracy improvements, DSRL exhibits superior sample efficiency and achieves higher Pass@K scores, indicating that the pre-train space learning successfully biases the policy toward high-quality solution subspaces without sacrificing diversity. The training dynamics reveal that DSRL sustains response length growth longer than standard approaches, suggesting that the initial NSR phase prevents premature convergence to short, suboptimal reasoning chains.

Original Insights and Critical Analysis

The distinction between optimizing $P(y)$ versus $P(y|x)$ touches upon fundamental questions regarding the nature of knowledge representation in neural networks. This work suggests that reasoning capability may be more fruitfully understood as a property of the unconditional generative process rather than purely a conditional mapping from problems to solutions. If valid reasoning trajectories are inherently probable under the marginal distribution, then conditioning on specific problems naturally inherits this capability.

However, several questions remain open regarding the limitations of this approach. The strong gradient alignment between marginal and conditional spaces likely depends on specific properties of the base model architecture and the particular reasoning domain. For tasks requiring highly context-dependent reasoning where the validity of a solution depends entirely on specific constraints provided in the prompt, optimizing $P(y)$ independently of $x$ might introduce harmful biases or fail to capture necessary conditional dependencies.

Furthermore, the dominance of NSR over PSR in the pre-train space suggests a broader principle for foundation model training: the initial learning phase may benefit more from adversarial or contrastive signals than from positive imitation. This aligns with recent work in self-play and adversarial training, but the specific application to the pre-train space offers a new lens on curriculum design. One might hypothesize that the optimal training curriculum involves an initial "clearing of misconceptions" via negative sampling before introducing positive exemplars.

From a computational perspective, PreRL offers potential efficiency advantages. Updates to $P(y)$ propagate across all possible contexts simultaneously, suggesting that PreRL might achieve faster convergence in terms of gradient steps, though this comes at the cost of requiring diverse sampling across the full output space. The paper does not fully explore the computational trade-offs between maintaining the broader exploration capacity of PreRL versus the targeted efficiency of standard RLVR.

Conclusion

The framework proposed by Tan et al. represents a meaningful expansion of the RLVR toolkit, demonstrating that reinforcement learning need not be confined to conditional spaces. By establishing the viability of Pre-train Space RL and identifying the unique efficacy of Negative Sample Reinforcement, this work opens new avenues for enhancing reasoning capabilities beyond the constraints of base model initialization.

Looking forward, several research directions emerge naturally from these findings. The theoretical characterization of when gradient alignment between $P(y)$ and $P(y|x)$ holds most strongly remains incomplete, particularly for multimodal models or specialized domains. Additionally, the interaction between pre-train space optimization and existing techniques such as rejection sampling or test-time scaling warrants investigation. As the field continues to scale model capabilities, understanding how to effectively manipulate both the marginal and conditional distributions may prove essential for breaking through current reasoning plateaus.

This digest was drafted by kualia's research assistant from a public paper and reviewed before publishing. It summarises the paper; it is not a substitute for reading it.