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

The evolution of large language models (LLMs) has transformed NLP, with methods like supervised fine-tuning and reinforcement learning (RL) significantly boosting reasoning and generation capabilities. Early RL approaches, such as PPO, improved reasoning but often caused output mode collapse, reducing diversity. GFlowNets, introduced by Bengio et al., offered a new paradigm by focusing on reward distribution matching, which naturally preserves diversity. Prior works like FlowRL adapted GFlowNets for LLM post-training, but their core reliance on the partition function as a normalization constant limited the exploitation of its informational potential. Despite progress, challenges remain in balancing sample efficiency, diversity, and training stability.

Existing RL methods for LLM fine-tuning face a fundamental trade-off: maximizing reward often leads to reduced output diversity, while distribution matching approaches like GFlowNets, though preserving diversity, treat the partition function as a static normalizer. This neglects the rich information embedded within the partition function, particularly its potential to reflect prompt difficulty. Consequently, training remains sample-inefficient, especially in complex reasoning tasks where data and compute costs are high. The core issue is how to leverage the partition function to dynamically guide training, improving efficiency without sacrificing diversity.

The key innovation is reinterpreting the GFlowNet partition function as an online accuracy estimator, establishing a theoretical link that allows it to serve as a difficulty scheduler. This enables adaptive prompt selection, focusing training on questions of intermediate difficulty, which maximizes learning progress. Additionally, the authors introduce an accuracy estimation error–prioritized replay mechanism, emphasizing samples where the model's accuracy estimate is uncertain. These mechanisms are integrated into the existing GFlowNet training pipeline, reusing existing signals to avoid extra computation. The approach balances exploration and exploitation, significantly boosting sample efficiency while maintaining diversity.

• �� Derive the theoretical relationship between Zϕ and per-prompt accuracy p_old(x) using a modified Trajectory Balance loss. • Replace the reference policy with the old policy to interpret Zϕ as an estimate of prompt accuracy. • Use Zϕ to select prompts with accuracy estimates near 0.5, ensuring informative training samples. • Generate rollouts for selected prompts, updating both the policy and Zϕ by minimizing the modified loss. • Implement an accuracy estimation error–prioritized replay buffer, storing prompt-output pairs with high estimation errors. • During training, sample from the buffer to focus on uncertain prompts, refining accuracy estimates and normalization. • Validate the approach on datasets like DeepCoder and DeepScaleR, with hyperparameters β=0.05, buffer size 128, and 8 rollouts per prompt.

Experiments involve training models such as Qwen2.5-Math-7B on datasets like DeepScaleR and DeepCoder, comparing PACED-RL to baselines like GRPO and FlowRL. Metrics include pass@1 and pass@k, with evaluations on benchmarks like AIME, HumanEval+, and Math500. Hyperparameters include a β of 0.05, buffer capacity of 128, and 8 rollouts per prompt. Ablation studies examine the impact of prompt difficulty targeting and replay prioritization. Results show faster convergence and higher accuracy, with improvements up to 29.1% in pass@1 and significant diversity gains.

PACED-RL outperforms baselines across multiple tasks, achieving up to 29.1% higher pass@1 on AIME and faster training times—about half the epochs needed for comparable performance. On code generation benchmarks, it reaches top performance with only 42-49% of the training time of previous methods. The pass@k scores indicate enhanced exploration and diversity, especially at larger k values, demonstrating the method’s ability to balance accuracy and variety. The ablation confirms that adaptive prompt selection based on accuracy estimates is crucial for these gains.

This approach is immediately applicable to tasks requiring high reasoning accuracy and output diversity, such as automated theorem proving, code synthesis, and complex question answering. It enables models to focus on optimally challenging prompts, reducing training costs and improving generalization. Long-term, integrating this dynamic difficulty scheduling into multi-task and multi-modal systems could foster autonomous AI that learns more efficiently across diverse domains, ultimately reducing reliance on massive labeled datasets and computational resources.

The method assumes stable training dynamics where KL divergence remains small; in unstable or highly noisy environments, accuracy estimates may degrade. Its effectiveness diminishes in tasks with extremely high complexity or where prompt difficulty cannot be reliably inferred. Scalability to larger models and broader tasks remains to be validated. Further, the approach depends on the quality of the initial GFlowNet training, which may limit its applicability in scenarios with limited data or unstable training conditions.