Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR

Technical Summary: Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR

Problem Statement
Reinforcement Learning with Verifiable Rewards (RLVR) has proven effective for enhancing the reasoning capabilities of Large Language Models (LLMs). However, standard reward-maximizing methods like PPO and GRPO often lead to overly sharpened output distributions, reducing generation diversity. While recent approaches utilizing GFlowNets address this by training policies to match a reward-induced target distribution, they typically treat the learned partition function ($Z$) solely as a normalization constant required for training stability. This paper argues that this perspective underutilizes the information encoded within $Z$, missing an opportunity to improve sample efficiency without incurring additional computational costs.

Methodology: PACED-RL
The authors propose Partition Function-Guided RL (PACED-RL), a post-training framework that reinterprets the GFlowNet partition function not just as a normalizer, but as a signal for per-prompt online accuracy. The methodology consists of three core components:

1. Theoretical Reinterpretation of the Partition Function:
   The paper establishes a theoretical link between the optimal partition function $Z^*(x)$ and the online accuracy $p_{old}(x)$ of a question prompt $x$ under the pre-update policy. By modifying the Trajectory Balance (TB) loss to anchor the KL divergence at the current policy ($\pi_{old}$) rather than a fixed reference model, the authors derive that $Z^*(x)$ directly encodes accuracy information. Specifically, $p_{old}(x) \approx \beta \log Z^*(x)$, where the KL divergence term is shown empirically to be negligible during standard training dynamics. This allows the partition function, already being learned for normalization, to serve as a free, online accuracy estimator.

2. Adaptive Prompt Selection:
   Leveraging the accuracy estimates derived from $Z_\phi(x)$, PACED-RL implements a difficulty-aware prompt selection strategy. Instead of sampling prompts uniformly, the method greedily selects the top-$m$ prompts whose estimated accuracies are closest to a target value $\tau$ (typically 0.5). This focuses training on "intermediate difficulty" samples, which prior work suggests offer the most informative learning signals, thereby maximizing sample efficiency.

3. Accuracy Estimation Error-Prioritized Replay:
   To further enhance efficiency and calibrate the partition function, the framework introduces a replay mechanism. A buffer stores prompt-output pairs with correct rewards ($r=1$). Within this buffer, samples are prioritized based on the magnitude of their accuracy estimation error: $| \frac{N_{correct}}{N} - \hat{p}_{old}(x) |$. This ensures that the model revisits samples where the partition function's prediction diverges most from observed reality, refining the accuracy estimates and the target distribution normalization.

Crucially, both the prompt selection and replay prioritization reuse information generated during standard GFlowNet training, amortizing the computational overhead into the existing optimization process.

Key Contributions

• Theoretical Insight: The paper demonstrates that the GFlowNet partition function for LLM post-training is intrinsically linked to online accuracies, enabling its direct use as an accuracy estimator without auxiliary mechanisms.
• PACED-RL Framework: The introduction of a novel training framework that utilizes these accuracy estimates for adaptive prompt selection and error-prioritized replay, improving sample efficiency while preserving the diversity benefits of GFlowNets.
• Empirical Validation: Extensive evaluations across mathematical reasoning and code generation benchmarks show that PACED-RL achieves significant performance gains over strong baselines (GRPO, FlowRL) and adaptive sampling methods (DS, LILO, MoPPS) without the associated computational overhead of over-sampling.

Experimental Results
Experiments were conducted on code generation (DeepCoder dataset, DeepSeek-R1-Distill-Qwen-1.5B) and mathematical reasoning (DeepScaleR dataset, Qwen2.5-Math-1.5B/7B).

• Performance Gains: On the AIME 24 benchmark, PACED-RL improved average pass@1 performance by 29.1% over GRPO and 40.0% over FlowRL using the Qwen2.5-Math-7B model.
• Sample Efficiency: PACED-RL achieved the best performance of GRPO and FlowRL in significantly less time (e.g., 0.42× the training time on LiveCodeBench). Unlike methods requiring over-sampling (DS, LILO), PACED-RL maintained a 1.0× rollout ratio relative to standard training, avoiding the latency penalties associated with generating extra samples for accuracy estimation.
• Diversity: Measured via pass@k metrics, PACED-RL consistently outperformed baselines, showing improvements of up to 14.2% over GRPO and 9.1% over FlowRL, confirming that the method preserves output diversity and exploration capacity.
• Ablation Studies: The authors verified that omitting the KL divergence term in the accuracy estimator has negligible impact on sampling dynamics while saving significant computation. They also showed that the method naturally shifts training toward harder questions as the model improves, preventing collapse to a narrow dataset subset.

Significance and Claims
The paper claims that PACED-RL offers a "principled and practical way" to enhance distribution-matching training for LLMs. By reinterpreting the partition function as an online accuracy signal, the method enables more sample-efficient training without the computational burden of auxiliary accuracy estimators or excessive rollout generation. The authors position this work as a step toward more efficient RLVR post-training, suggesting that the partition function's role extends beyond normalization to act as a dynamic difficulty scheduler. Future work is suggested to explore inference-time strategies (e.g., adaptive self-consistency) and extending the framework to asynchronous GFlowNet settings.