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

Reward-maximizing RL methods have shown to be capable of enhancing the reasoning performance of LLMs, but often lead to reduced generation diversity. 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 \textbf{Pa}rtition Fun\textbf{c}tion-Guid\textbf{ed} \textbf{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.

To identify the best training samples to learn from, previous works utilize expensive auxiliary mechanisms that incur additional computational overheads. We show that the partition function in GFlowNets training, previously under-utilized during the training process, naturally encodes enough information to select the best training samples. By doing away with the need to employ auxiliary mechanisms, for a fixed compute budget, our method trains LLMs to converge faster and higher compared to baseline approaches.