Adapting Monte Carlo Tree Search for Generative Flow Network Training

Comparison with other baselines:

Since the authors position the purpose of their approach to be improving trajectory sampling for GFlowNet training, I would like to see some baselines other than the standard on-policy training. E.g. replay buffer approaches [3, 4], and using a mixture of a forward policy and a uniform distribution [5], which were demonstrated to also improve the efficiency of GFlowNet training. As stated by the authors, combining MCDS with such methods could lead to further improvements, but putting some direct comparison to them is also crucial for understanding the efficiency of MCDS.

We have added some experiments about epsilon-exploration in the Hypergrid environment to the Appendix, see Section 7.4. Overall, MCDS still offers an improvement over on-policy.

L1 metric computation:

In addition, I have a concern over which policy is used for evaluation. A standard approach to calculate L1 metric on hypergrids is to use an empirical distribution of terminal states of the last $N$ trajectories sampled for training, and then compare it to the true reward distribution. If this approach is also used here, then the policy being evaluated would be not the GFlowNet forward policy $P_{F}$, but the combination policy $P_{M}$. Then, such evaluation would show the improvements coming not only from introducing a different distribution over trajectories for GFlowNet training (motivation stated by the authors), but also from the adjustment to the forward policy introduced by MCDS. Can the authors please provide more details on how the metric is calculated? For factor graph environment I have the same question.

In both the Hypergrid and Factor Graph experiments, we do not employ search during inference for our MCDS/MCTS evaluations. While it is true that we draw samples from $P_{M}(s'|s)$ to train the models, we re-sample trajectories (independently) using $P_{F}(s'|s;\theta)$ for the purposes of approximating $P(x;\theta)$ to estimate L1 error and JSD. As the reviewer inferred, we maintain a buffer of the previous $N$ sampled trajectories to obtain a running estimate of $P(x;\theta)$. For both Hypergrid and Factor Graph experiments, we use $N=2e5$. In the original submission, we mention this in the Appendix, see Section 7.3.

Runtime measurements and computational overhead:

It will also be good to include some runtime measurements to see how much* of a computational overhead MCDS adds in comparison to standard on-policy training.

The overhead of MCDS depends primarily on the construction budget, the number of parallel workers, and the frequency with which the search DAG is cleared.

Table 4 summarizes the relative slowdown of different MCDS variants used in the Hypergrid and Factor Graph experiments (DB parameterization). Note that the reported metrics include time associated with the calculation of rewards, losses, gradients, and evaluation metrics. With the configurations we tested, the total MCDS runtime penalty ranges from a factor of 1.30 to 3.47.

Correction to Eq. 3:

A minor remark, if $Q(s,s')$ in Equation 3 is an entropy-regularized action value function defined consistently with the standard RL notation, the correct way to write it would be $Q(s,s') = R(s,s') + \log \sum_{s''} \exp Q(s',s'')$. On the other hand, if it is defined as an exponent of action value function, I believe there should be no logarithm in the equation on line 187.

This is correct, we apologize for the error and will update the equation on line 187 to reflect this change.

Equation 3 will read: $Q_{T}(s,s') = R(s,s') + Q_{T}(s)$

The equation on line 187 will read: $Q_{T}(s) = \log \sum_{s' \in Ch(s)} \exp Q_{T}(s,s')$

Flow estimates for hypergrid in Section 4.1:

How are flow estimates $\log F_{D}$ initialized for the newly added MCDS edges in the hypergrid experiment in Section 4.1 (when there are no neural networks)?

We assume the author is referring to the experiments in Figure 2. The flow estimates are all initialized to be zero (meaning the initial policy $P_{D}$ is random uniform). Note that this does not affect the properties of the MCDS or MCTS flows after exhaustive sampling. Some other initialization strategy might be more appropriate (i.e. using a function that depends on the length of the shortest path between the root and the newly added node), but we think this simple strategy is sufficient for a simple demo that is supposed to highlight the differences between MCDS and MCTS at convergence, rather than behaviour of the two algorithms with small sample size.

Note on theoretical contribution:

The theoretical statements seem to be a trivial corollary of the GFlowNet-RL connection by (Tiapkin et al. 2024) and (Deleu et al. 2024), applied to a GFlowNet problem over the subgraph. Thus, I cannot consider the theoretical contribution as a crucial part of this paper.

We agree that the theoretical results are a corollary of previous connections, although for clarity we emphasize that Theorem 2 applies to all GFlowNet problems with deterministic transitions. The assumption of a DAG-structured environment is standard in most GFlowNet formulations.

Scale of environments:

The hypergrid and factor graph environments seem limited in scale, as the final distribution over terminal states is tractable in both, which was used to compute the success metric.

This is correct, although the sizes of our hypergrid and factor graph environments are consistent with previously published works (Malkin et al 2022, Madan et al 2022, Deleu et al 2024), and are generally accepted as useful GFlowNet benchmarks. As noted below, some of the Factor Graphs experiments can be quite large: for example, the Permuted Chain environment has $|\mathcal{S}| \approx 2\mathrm{e}9$.

The only large-scale problem presented lacks sufficient details for reproducibility. Specifically, there is no information on how MCDS is applied in the language reasoning task. If MCDS is used on underlying structures rather than at the language level, it could be possible to construct an exhaustive or near-exhaustive DAG, simplifying the language reasoning task to an almost trivial problem. To assess whether this is the case, it is essential to know: (a) the size of the MCDS search and (b) the size of the underlying graphs for the reasoning task. The underlying GFlowNet graph appears relatively small, which could make MCTS-like techniques applied to these structures overly effective.

The reviewer has correctly identified that we perform MCDS search on the underlying reasoning graphs. For the reasoning problems we consider, it is possible to construct an exhaustive search DAG in a reasonable time. However, with the exception of some of the 2-step environments, the budgets we use in training are smaller than what is required for exhaustive search (see Table 3 for details).

Note that we do not apply MCDS for inference (in this experiment or any others), we only use on-policy sampling. Also, in the reasoning task we only evaluate the model on problems (environments) which are not in the training set; exhaustively exploring the training environments would not necessarily lead to perfect performance at inference time.

Reasoning baselines:

Additionally, there are no comparisons with SFT and RLHF baselines for the reasoning task. Comparing only to an inference-level CoT baseline is not enough, as training data is used. It would be more informative to compare the CoT baseline with few-shot examples included in the training set.

If our goal was simply to improve reasoning (by any means), we agree that it would be useful to consider other formulations of the problem that are not using GFlowNets (for example RLHF and DPO, which essentially formulate reasoning as an RL task with a Bradley-Terry reward model). However, we were more interested in demonstrating our method's performance on a challenging GFlowNet benchmark task.

Standard Deviation in the plots:

The plot states that it shows mean and standard deviation however there no markers, fills or any visual indication of the standard deviation.

The standard deviations are plotted in Figure 3, but they are small so it is difficult to see.

The standard deviation is not a statistical significance test --- it is a measure of dispersion. Confidence intervals would be an appropriate choice for visualizing statistical significance. Without any statistical testing it is difficult to verify whether the difference in performances are significant. Hence, I believe the results do the support the authors' claims that MCDS improves the performance of GFlowNets over on-policy sampling in Hypergrid.

Statistical testing is only one tool to establish significance in an experiment. We chose to plot the standard deviation because it can give an idea of the variability of the compared algorithms, not just the difference in average performance. The reviewer notes that the standard deviations are too narrow to see; it would be even narrower if we plotted standard error. If we increased the sample size, the standard error would also get narrower.

Comparision to MCTS sampling:

Secondly, inspecting the plots do not show that their is indeed any performance gain over conventional MCTS sampling.

Please see the curves in Figure 3c and 3f in which the MCTS curves overlap with the on-policy curves, while the MCDS curves in Figure 3a and 3d descend faster than the on-policy curves, hence showing improvement.

Number of seeds:

Lastly, even if CI plots were produced it is hard to imagine that averaging over a mere 3 seeds would produce results of significance. With 3 seeds the authors may just plot the results over each seed.

We did not see much variation for the runs with different seeds for our method, and hence we do not expect the plots for Hypergrid to change by adding more than 3 seeds. Please note that for our factor graph experiments (Figure 4), where the results are much more variable, we included more seeds (at least 10) and plotted individual seeds as dashed lines.

Details of the Hypergrid domain:

The explanation of the Hypergrid domain in Sec 5.1 did not make any sense and there was no explanation of the performance metric. I had to read the original GFlowNet paper to understand it.

We apologize for the confusion. Hypergrid is well-known in the GFlowNet community, although it might be confusing for readers who are less familiar with this area. We will elaborate on the definition in the Appendix (7.3).

Clarification on Fig 2 and MCTS-tree and terminating distribution:

The explanation for MCTS in Sec 2.2 describes it such that all states in a MCTS-tree, and consequently MCDS-tree, are unique which is not necessarily true. In fact, if I understand the diagrams in Fig 2, the equivalent trees all will have duplicate states.

The version of MCTS we are describing is similar to the one described in Buesing et al 2019. In this setting (and the GFlowNet setting), we assume that all states in the tree are unique (even though it might be possible to reach them with multiple paths). This assumption of uniqueness also applies to MCDS. In Figure 2, each state in the environment is a cell in the 2-dimensional hypergrid, so there are no duplicate states.

Terminating distribution:

Additionally, in Fig 2, I cannot understand why the forward policy for MCTS is the way it is. If the reward function for MCTS is using the one shown in Eq 2, then if sampled enough, MCTS should terminate at the bottom left corner and observe a reward. Am I missing something?

Because we are using maximum entropy MCTS, the terminating distribution $P_{D}(x)$ is biased by the number of trajectories leading to each state $x$ (see Bengio et al 2021 for more details).

The number of paths from the initial state $s_0 = (0,0)$ (bottom-left) and an arbitrary state $s_{m,n} = (m,n)$, where $m,n \leq 7$, is equal to the binomial coefficient ${m+n \choose n} = \frac{(m+n)!}{m!n!}$. Consider two modes in the 8x8 Hypergrid: the mode at $(1,1)$ and the mode at $(6,6)$. The former only has $\frac{2!}{1!} = 2$ trajectories that lead to it, while the mode at $(6,6)$ has $\frac{12!}{6!6!} = 924$ trajectories that lead to it. The MCTS solution (i.e. the maximum entropy solution) greatly favours termination at states that are further from the initial state $s_0$, because these states can be reached with a greater number of trajectories.

Definition of Q:

Line 157: What is $Q$? It hasn't been defined.

It is defined on Line 151 in Section 2.2 (of the original paper).

References

Bengio et al 2021

Buesing et al 2019

Experimental results, seeds and std deviation:

Some claims are not supported by the experiment results, which is exacerbated by poor empirical practices. For instance, line 411: "MCDS improves training compared to on-policy sampling" is supported by the learning curves in Figure 3. However, these curves are averaged over 3 runs, and the shaded regions (which are not visible) are standard deviations, a measure of dispersion and not confidence regarding the mean. There is little evidence that the difference is statistically significant. Standard deviation is not a measure of confidence and 3 seeds are insufficient to make such a general claim. The choice of hyperparameters is not justified. The reported metric is also not justified. The same criticism also applies to line 412: "Larger tree construction budgets providing a bigger improvement".

Our presentation of the Hypergrid results in Figure 3 are consistent with previous works (Bengio et al 2021, Madan et al 2022). Since Hypergrid is an easier environment, the difference in the learned L1-errors are quite low for our method. We don't expect the plots to change by adding more seeds or by using other metrics.

Hyperparameter selection:

In general, the hyperparameter choices for no experiment are justified, and all claims regarding differences between methods in Hypergrid and Blocksworld are not supported by statistical tests (e.g. confidence measures).

We have added more information about the hyperparameters to the Appendix, thank you for this suggestion.

Comparison with other baselines:

While the experiments show comparisons with on-policy sampling, there is no comparison with other baselines that attempt to improve GFlowNet training, such as epsilon uniform exploration, replay buffers, or local search. This limitation is acknowledged in the paper, however, I think their inclusion is required to evaluate MCDS.

We've added some experiments regarding epsilon uniform exploration in the Hypergrid to the appendix, see Section 7.4. Overall, MCDS still offers an improvement over on-policy.

Clarification on Figure 2:

What does the top plot in Figure 2 show? What do the lines represent? Why is this difference significant?

The top plot represents the approximation quality of $P_{D}(x)$ induced by flows $F_{D}(s,s')$ calculated with MCDS and maximum entropy MCTS. The line plot shows approximation error (measured using Jensen-Shannon divergence) with different tree construction budgets. The heat maps visualize the state flows $F_{D}(s)$ in green (darker = larger flow), where each state $s$ is represented as a state in the grid. The policies $P_{D}(s'|s)$ are visualized with blue arrows and red dots. In the visualization, the probability of each action is proportional to the size of its corresponding symbol.

In an 8x8 hypergrid, which has $|\mathcal{S}| = 65$ and $|\mathcal{A}| = 176$, a budget of 176 should theoretically be sufficient for exhaustive exploration with MCTS/MCDS. However, with our parallel implementation there is a stochastic component: when the number of workers is greater than one, it is possible for multiple workers to "EXPAND" the same node, potentially resulting in an incomplete tree even after the budget of 176 is exhausted. For this experiment we confirmed that budgets of $2^9$ and $2^{10}$ resulted in exhaustive exploration.

Both methods produce better approximations with larger budgets, but MCTS is unable to converge to the correct solution (i.e. JSD of zero) even with a very large budget. This happens because maximum entropy MCTS results in a flows that are biased towards terminating states with a larger number of trajectories. This is highlighted by the fact that in the four squares in the bottom-left corner, the "terminate" action receives very little probability mass in the MCTS visualization when compared to MCDS/optimal. These four states actually have high reward, but the MCTS solution assigns low probability of terminating because the trajectory lengths to get to them are very short (less than 2 steps).

To see this more clearly, we've modified Figure 2 in the updated paper to show the predicted terminating distributions $P_{D}(x)$ for the tree policies from MCTS and MCDS, instead of the state flows $F_{D}(s)$.

Note that in the hypergrid domain, every state is a terminating state. Also, this domain is small enough such that the optimal flows and policies, and the terminating distribution, can be calculated exhaustively in a short time. We also assume that there is a uniform backward policy $P_{B}(s|s')$ as is commonly used for GFlowNets, although this demonstration would also work with a different choices of $P_{B}(s|s')$.

References

Bengio et al 2021

Madan et al 2022