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

TL;DR

This paper introduces PACED-RL, a novel framework that reinterprets the partition function in GFlowNets as an accuracy signal, improving sample efficiency and performance in training large language models for diverse outputs.

Contribution

It redefines the partition function as an online accuracy indicator and develops PACED-RL, a post-training method that enhances sample efficiency by prioritizing prompts based on this signal.

Findings

PACED-RL outperforms GRPO and prior GFlowNet methods across benchmarks.

The approach reuses existing information to reduce computational overhead.

Results show significant improvements in training efficiency and output diversity.

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.