Accurate and Diverse LLM Mathematical Reasoning via Automated PRM-Guided GFlowNets

TL;DR

This work tackles robust mathematical reasoning in LLMs by introducing an automated Process Reward Model (PRM) and a step-level Generative Flow Network (GFlowNet) fine-tuning framework to balance accuracy and diversity. The PRM is trained via Monte Carlo Tree Search with continuous step-quality scores and augmented by rollout reuse with similarity grouping, enabling rich step-level rewards without human annotations. A step-level GFlowNet uses a multiplicative trajectory reward $R(\tau)=\prod_{i=1}^n U(s_i|s_{i-1})$ and the Subtrajectory Balance loss to learn reward-proportional sampling of reasoning steps, thereby fostering diverse, high-quality solutions. Empirically, the approach yields improvements on MATH Level $5$ (+$2.59$pp) and SAT MATH generalization (+$9.4$pp) across model sizes, with notable diversity gains, and demonstrates superior alignment of PRM-guided rewards compared to open-source rewards, highlighting the potential of step-level diversity in educational and broader reasoning tasks.

Abstract

Achieving both accuracy and diverse reasoning remains challenging for Large Language Models (LLMs) in complex domains like mathematics. A key bottleneck is evaluating intermediate reasoning steps to guide generation without costly human annotations. To address this, we first introduce a novel Process Reward Model (PRM) trained automatically using Monte Carlo Tree Search coupled with a similarity-based data augmentation technique, effectively capturing step-level reasoning quality. Leveraging this PRM, we then adapt Generative Flow Networks (GFlowNets) to operate at the reasoning step level. Unlike traditional reinforcement learning focused on maximizing a single reward, GFlowNets naturally sample diverse, high-quality solutions proportional to their rewards, as measured by our PRM. Empirical evaluation shows strong improvements in both accuracy and solution diversity on challenging mathematical benchmarks (e.g., +2.59% absolute accuracy on MATH Level 5 for Llama3.2-3B), with effective generalization to unseen datasets (+9.4\% absolute on SAT MATH). Furthermore, we benchmark our PRM against existing open-source reward models, demonstrating superior alignment with reasoning quality and more consistent guidance for downstream generation. Our work demonstrates the potential of PRM-guided, step-level GFlowNets for developing more robust and versatile mathematical reasoning in LLMs.

Figures (4)

• Figure 1: Data processing workflow for PRM training. Starting from a step $s$ with $MC(s)=1/2$, the diagram shows how subsequent steps are processed based on their Monte Carlo values. Similar steps (indicated by dashed boxes) share MC values. Steps following an incorrect step ($MC=0$) are excluded from the training dataset, as they would be built upon invalid reasoning. Gray boxes indicate steps that become irrelevant to the training process.
• Figure 2: Accuracy of Falcon models on a subset of MATH Hard using PRM-guided search with varying numbers of proposed steps $k$. Horizontal lines indicate the baseline accuracy of unguided Falcon3-3B, 7B, and 10B models using prompt-based decoding. Solid (resp. dotted) curves represent the accuracy of Falcon3-1B, 3B, and 7B guided by our PRM (resp. Skywork-o1-Open-PRM-Qwen-2.5-7B).
• Figure 3: Reasoning steps and corresponding PRM scores. Valid steps from different approaches receive high, comparable scores, while corrupted steps receive lower scores.
• Figure 4: Training dynamics during GFlowNet fine-tuning of Llama3.2-3B-it showing (a) Sub-TB loss convergence, (b) average reward improvement, and (c) proportionality gap reduction. The consistent decrease in proportionality gap demonstrates successful alignment between token selection probabilities and PRM rewards, validating the effectiveness of our Sub-TB learning approach.