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

TL;DR

This work tackles the diversity-loss problem in reward-maximizing RL for LLM reasoning by reframing the GFlowNet partition function $Z(\\mathbf{x})$ as an online accuracy signal. It introduces PACED-RL, which uses accuracy estimates derived from $Z_{\\phi}(\\mathbf{x})$ to drive difficulty-aware adaptive prompt selection toward a target $\\tau$ and to support accuracy-estimation–error–prioritized replay, all while amortizing computation through reuse of existing TB training signals. Theoretical results connect $Z^{*}(\\mathbf{x})$ to online accuracy via $p_{old}(\\mathbf{x}) = \\beta \\log Z^{*}(\\mathbf{x}) - \\beta \\mathrm{D}_{KL}(\\pi_{old} \\| \\pi_{\\theta})$, enabling practical approximations $p_{old}(\\mathbf{x}) \\approx \\beta \\log Z^{*}(\\mathbf{x})$. Empirically, PACED-RL yields substantial gains on code-generation and mathematical reasoning benchmarks (e.g., up to $29.1\%$ and $40.0\%$ improvements on AIME over GRPO and FlowRL, with consistent pass@k enhancements) and exhibits favorable compute overhead, validating a principled, sample-efficient approach to distribution-matching for LLM reasoning.

Abstract

Reward-maximizing RL methods enhance the reasoning performance of LLMs, but often reduce the diversity among outputs. Recent works address this issue by adopting GFlowNets, training LLMs to match a target distribution while jointly learning its partition function. In contrast to prior works that treat this partition function solely as a normalizer, we reinterpret it as a per-prompt expected-reward (i.e., online accuracy) signal, leveraging this unused information to improve sample efficiency. Specifically, we first establish a theoretical relationship between the partition function and per-prompt accuracy estimates. Building on this key insight, we propose Partition Function-Guided RL (PACED-RL), a post-training framework that leverages accuracy estimates to prioritize informative question prompts during training, and further improves sample efficiency through an accuracy estimate error-prioritized replay. Crucially, both components reuse information already produced during GFlowNet training, effectively amortizing the compute overhead into the existing optimization process. Extensive experiments across diverse benchmarks demonstrate strong performance improvements over GRPO and prior GFlowNet approaches, highlighting PACED-RL as a promising direction for a more sample efficient distribution-matching training for LLMs.

Figures (8)

• Figure 1: Training dynamics of mean of estimated values of $\beta\,D_{\mathrm{KL}}\!\left(\pi_{\mathrm{old}}(\cdot\!\mid \mathbf{x}) \,\!\|\!\, \pi_{\theta}(\cdot\!\mid \mathbf{x})\right)$ of Qwen2.5-Math-1.5B trained on the DeepScaleR dataset, using the TB loss defined in Eq. \ref{['eq:our_tb_loss']}
• Figure 2: Pass@1(%) dynamics of DeepSeek-R1-Distill-Qwen-1.5B plotted with respect to wall-clock time. Pass@1 values obtained via averaging over 8 independent attempts.
• Figure 3: Pass@k(%) curves of Qwen2.5-Math-7B trained on GRPO, FlowRL, and PACED-RL and evaluated on AIME 24 and AIME24 for $k \in \{2,4,8,16,32,64,128,256\}$. Across all evaluated values of $k$, PACED-RL achieves the highest performance.
• Figure 4: Train and test accuracies of Qwen2.5-Math-1.5B trained on the DeepScaleR dataset and evaluated on the MinervaMath benchmark, for target accuracy values $\tau \in \{0.1,0.3,0.5,0.7,0.9\}$.
• Figure 5: Progression of the average difficulty of question prompts sampled in each training batch under the Qwen2.5-Math-7B model.

Theorems & Definitions (1)

• Proposition 4.1