Flow of Reasoning: Training LLMs for Divergent Reasoning with Minimal Examples

Thanks for your insightful comments and suggestions.

• Q1. How representative are the benchmarks? Are there reasoning types where it fails?

A1: Some tasks like GSM8K and Game24, are popular LLM benchmarks, while others, like BlocksWorld and 1D-ARC, are known to be challenging for LLMs. FoR remains a general framework for multi-step reasoning: whenever a problem can be decomposed into intermediate states and has a well-defined reward, we expect FoR to be applicable.

• Q2: Tradeoff between amortized inference vs. slower inference with faster training. When training cost of FoR is justified by its better inference performance?

A2: Unlike methods like ToT/RAP without diversity, FoR enables efficient amortized inference with diversity. Though training time is long, it eliminates abundant inference costs in ToT/RAP, keeping the overall cost low when the number of tasks needed to be solved is numerous. Once trained, it avoids repeated, expensive searches. In such cases, the initial training cost is largely offset by the benefits of efficient inference. Compared to SFT/PPO, FoR has higher training costs but is highly data-efficient, needing only ~15 examples, reducing data collection costs.

• Q3: Diversity metric (human annotation & automatic)

A3: Following your advice, we use a paraphrase-based method [1] with GPT-4o to automatically evaluate solution diversity on GSM8K. Please see detailed results in A2 to reviewer C9xd due to page limit. The results show that FoR’s diversity is still higher than baselines by using GPT4o, although it consistently overestimates.

• Q4: Evaluation of diversity for search-based methods and O1-series.

A4: We ran ToT-DFS multiple times on the Blocksworld task with the same setup as in the paper. The results are shown below:

Search methods (ToT/RAP) show limited diversity despite multiple runs. Table 1 in the paper shows ToT/RAP needs 4-40× FoR's inference time per sample, making equal-sample comparison time-unfair. We'd like to clarify that Table 1 shows that O1-mini achieves decent accuracy but limited diversity across multiple runs.

• Q5: Replay buffer is more important without local search?

A5: We add an additional ablation study by removing the local search to assess the impact of the other components in FoR. See the results below:

Table 5 in the paper and the above results show that removing the replay buffer causes a smaller performance drop in FoR with local search (4-6%) than without it (11-19%), indicating that the replay buffer is more critical when the local search is removed.

• Q6: While applying GFlowNets to LLM reasoning is novel, core techniques are adapted from existing GFlowNet literature.

A6: Thanks for recognizing our novel formulation, which allows us to adapt existing approaches to the new multi-step LLM reasoning domain. This ability to reuse and generalize, rather than reinvent everything from scratch, is a key advantage of our work.

• Q7: Related works.

A7: We'll add more discussion on Takase et al. (2024) in the revision. Their work generates diverse math solutions via token-level diversity, similar to Hu et al. (2024) (mentioned in our Introduction, Lines 95–96). Oh et al. (2024) [2] follow these lines and apply GFlowNets to LLM preference optimization. In contrast, FoR targets broader reasoning tasks beyond math or preference learning, modeling thought process structures with GFlowNets.

• Q8: Details on training data sampling and replay buffer update mechanism.

A8: We used random sampling from the full training set. For the replay buffer, we set the capacity to store 50 trajectories, new trajectories will replace the lowest-reward ones when the buffer is full. These details will be added to Appendix D.

• Q9: More methods (e.g. efficient exploration) to reduce the training cost.

A9: We use (1) local search and (2) ϵ-sampling (Sec. 4.7) to encourage efficient exploration. In Game24, we use offline data to reduce model's exploration costs. Other methods to accelerate training are a promising future direction.

• Q10: Different base LLMs.

A10: We test LLama-3-8B, Qwen-2.5-7B, and InternLM2.5-7B-Chat on Blocksworld (one task due to rebuttal constraints), and FoR consistently outperforms baselines (see A5 to reviewer c9xd due to page limit).

References:

[1] Michail et al. "PARAPHRASUS: A Comprehensive Benchmark for Evaluating Paraphrase Detection Models." COLING 2025.

[2] Oh Joon Kwon, et al. GDPO: Learning to Directly Align Language Models with Diversity Using GFlowNets. EMNLP 2024

Thanks for recognizing our strong performance across 6 benchmarks, ablation studies, data efficiency, and case studies.

• Q1: Hyperparameters selection and analysis? Guidelines for designing rewards? Handcrafted rewards?

A1: Please refer to Section 4.7 and Figure 3 for the hyperparameter analysis on a small training set (i.e., 15). We search for the best-performing λ values on each task's training set.

We further assess how the intermediate reward weight λ affects test set performance. Results for varying λ values are shown below:

The results show that performance consistently improves as λ increases up to 1.5, but drops a bit at λ= 2 across 4- to 6-step settings. The influence of λ becomes more pronounced with more steps.

For the guidelines, we use the generalizable common principles: 1. For tasks with likely correct outputs (e.g., GSM8K), a binary success reward suffices. 2. For challenging reasoning tasks (e.g., BlocksWorld, Game24) with sparse rewards, intermediate rewards from prior works (Hao et al., 2023) (e.g., LLM log-likelihood, rule-based) help.

Regarding handcrafted rewards, in GSM8K, we simply use a standard outcome reward (0/1). Additionally, we use the reward model Qwen2.5-Math-PRM-7B, achieving 62.62% accuracy and 1.31 diversity (please see A2 to reviewer c9xd due to page limit). These two reward designs show that FoR does not rely on sophisticated reward design, supporting general applicability.

We will include all these in the revision.

• Q2: Varying amounts of training data.

A2: We ran additional experiments on the 6-step BlocksWorld task using Llama-3-8B with varying training sizes {1, 15, 30, 45, 60}, following the same setup as Section 4.7. The results are shown below:

FoR consistently outperforms SFT across all data sizes—for example, with just one training example, FoR achieves a relative 200% higher accuracy and a relative 30% higher diversity. This highlights FoR’s data efficiency even in low-resource settings.

• Q3: Theoretical derivations.

A3: Please refer to Appendix B for background information on GFlowNets.

For the trajectory balance (TB) constraint, please see Section 3.2. Below is a short derivation of the TB constraint:

Let $P_F(τ) = \prod_{t=1}^{n} P_F(s_t \mid s_{t-1})$ be the forward trajectory distribution, and $P_B(τ) = \prod_{t=1}^{n} P_B(s_{t-1} \mid s_t)$ be the backward trajectory distribution.

The TB constraint requires: $P_F(τ) \cdot Z(s_0) = P_B(τ) \cdot R(s_n)$.

This constraint aligns the flow allocated to $τ$ under the forward policy with reward-scaled backward probability, enabling proper credit assignment.

We will add these in the revision.

• Q4: Scalability and open-ended tasks.

A4: For larger models, we run additional experiments with LLaMA-3-70B and Qwen2.5-72B, and FoR consistently achieves better accuracy and diversity (see A5 to reviewer c9xd due to page limit). On the open-ended GSM8K benchmark, FoR also achieves stronger performance, highlighting its scalability.

• Q5: Limitations and potential failure modes.

A5: Please see the Limitations in Appendix H and will move it to the main text. For potential failure modes, FoR, like other on-policy RL methods, may struggle in sparse-reward settings. However, FoR mitigates this by the augmented intermediate rewards (Sec. 4.7), which provide dense guidance, and by leveraging off-policy data (Sec. 4.3 Game24).

• Q6: More reasoning tasks, such as Math and Code.

A6: We’d like to clarify that FoR has already been evaluated on GSM8K, a math benchmark (Sec. 4.6), as well as on ARC-1D (Sec. 4.5), which involves Python program synthesis.

• Q7: Optimizations to be computationally efficient?

A7: We'd like to clarify that although FoR’s training is slower, it is data-efficient (Section 4.7), requiring only ~15 training examples—keeping the overall training cost low. Two potential further optimization directions are: 1. Off-Policy Data: Parallel off-policy training can reduce the computational cost (Appendix H). 2. Intermediate rewards: accurate reward functions (e.g., robust reward models) enable faster convergence.

• Q8: A large number of reasoning paths? Scaling limitations?

A8: FoR performs well with large reasoning spaces. For example, the Game24 task involves ~8,000 distinct reasoning trajectories per sample, showing that FoR could scale to complex reasoning tasks.

• Q9: Other techniques like verifiers?

A9: We use rule-based rewards as verifiers in Rubik’s Cube and 1D-ARC (Sec. 4.4&4.5). Your suggestion is also a promising direction.

Thank you for recognizing our innovative Markovian flow approach to multi-step LLM reasoning with GFlowNets, our strong performance across six benchmarks, and our method’s data efficiency.

• Q1: More complex real-world tasks?

A1: We evaluate FoR on six benchmarks that pose significant challenges for current LLMs (e.g., embodied reasoning in BlocksWorld and math reasoning in GSM8K). These tasks capture core difficulties encountered in real-world scenarios. For instance, the BlocksWorld task requires not only spatial reasoning but also practical world knowledge, closely mirroring real-world planning problems.

• Q2 The diversity metric to capture semantic differences?

A2: Thank you for your suggestion. We intentionally use human annotation for GSM8K to ensure the evaluation is accurate and reliable, despite the high cost of it. Here, we run additional experiments using a paraphrase-based method [1] to automatically evaluate the diversity with GPT-4o to measure the number of different solutions for each problem. The results on GSM8K are shown below:

Compared to human annotation, GPT-4o overestimates the reasoning diversity, indicating that these automatic metrics currently do NOT provide more reliable evaluation. So, we leave this evaluation method as future work.

• Q3: Reward formulations that enhance diversity?

A3: Thanks for your inspiring comment. We agree that explicitly incorporating diversity into the reward formulation may further enhance solution diversity. We will discuss this in the next version.

• Q4: Analysis of diversity truly benefits downstream applications

A4: Thanks, we'll further highlight this in the revision. We'd like to clarify that diversity indeed benefits downstream tasks by improving accuracy, as mentioned in the Introduction (Lines 101-104). For a detailed analysis, please refer to Section 4.7 (Experiment Discussion) and Appendix F. We find that diversity promotes exploratory behaviors (e.g. in Game 24, it explores a large trajectory space), thereby enhancing the robustness of the model (Fig. 6).

• Q5: Scale with larger models, and diversity gains compare to baselines?

A5: Given the time and resource constraints in our research lab during the rebuttal period, we ran additional experiments to evaluate FoR's scalability and diversity gains using LLaMA-3 (8B & 70B) and Qwen2.5 (7B & 72B) on the BlocksWorld task. Results, shown below, indicate that FoR consistently improves accuracy and diversity with larger models. Compared to CoT baselines, FoR exhibits greater diversity gains as model size increases. This suggests that the gains of FoR become more pronounced as model capacity increases. Even with minimal data (15 examples), FoR yields clear improvements, highlighting its potential for robust performance across different base models. The results are shown below and will be included in the revision.

• Q6: Iterative reasoning refinement based on feedback?

A6: Yes, FoR supports iterative refinement. During both training and inference, by incorporating feedback—from the model itself or an additional model—into the state, subsequent state predictions refine solutions iteratively. This aligns with R1/O1's approach, where each step either refines the current state or advances the reasoning. For example, in BlocksWorld, given a current block position, the model predicts an action like 'move blue onto yellow,' and then generates feedback on whether to proceed or refine. Based on the state including this feedback and block positions, the model predicts the next state, iteratively refining its actions or proceeding. While we haven't tested this, iterative refinement based on feedback is a promising future direction.

• Q7: A meta-learning approach to reduce the need for task-specific reward engineering?

A7: While we haven't directly explored meta-learning approaches, we ran an additional experiment on an existing process reward model, Qwen2.5-Math-PRM-7B, for the GSM8K task. Table in A2 indicates that using this reward model with FoR improves performance over the previous 0/1 outcome reward function, achieving a 9% relative increase in accuracy and a 4% relative gain in diversity. This shows a promising step towards reducing the need for task-specific reward engineering.

References:

[1] Michail et al. (2025) "PARAPHRASUS: A Comprehensive Benchmark for Evaluating Paraphrase Detection Models." COLING 2025.