Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR
This work addresses sample efficiency in distribution-matching training for LLMs, representing an incremental improvement over existing GFlowNet methods.
The paper tackles the problem of reduced output diversity in reward-maximizing RL methods for LLMs by reinterpreting the partition function in GFlowNets as a per-prompt expected-reward signal, proposing PACED-RL which improves sample efficiency by prioritizing informative prompts and using error-prioritized replay. Experiments show strong performance improvements over GRPO and prior GFlowNet approaches.
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.