Generative Flow Networks (GFlowNets; GFNs) are a family of energy-based generative methods for combinatorial objects, capable of generating diverse and high-utility samples. However, consistently biasing GFNs towards producing high-utility samples is non-trivial. In this work, we leverage connections between GFNs and reinforcement learning (RL) and propose to combine the GFN policy with an action-value estimate, $Q$, to create greedier sampling policies which can be controlled by a mixing parameter. We show that several variants of the proposed method, QGFN, are able to improve on the number of high-reward samples generated in a variety of tasks without sacrificing diversity.
Paper
Similar papers
Peer review
Summary
The paper "QGFN: Controllable Greediness with Action Values" introduces a novel approach to enhance Generative Flow Networks (GFNs) by incorporating action-value estimates (Q-values) to control the greediness of sampling policies. This method, called QGFN, includes three variants—p-greedy, p-quantile, and p-of-max—each designed to balance the generation of high-reward samples with the maintenance of diversity. Through comprehensive experiments on tasks like molecule generation and RNA design, the authors demonstrate that QGFN significantly improves the generation of high-utility samples while preserving diversity, providing a practical and effective solution for deploying safe and helpful LLMs.
Strengths
- **Innovative Approach**: The introduction of Q-values to control the greediness of sampling policies in Generative Flow Networks (GFNs) is a novel and creative solution to the challenge of balancing high-reward sample generation with diversity. - **Comprehensive Evaluation**: The paper includes thorough experiments across various tasks, such as molecule generation and RNA design, providing strong empirical evidence for the effectiveness of QGFN. - **Ablation Study**: The paper provides in-depth ablation studies on the hyperparameters of GFNs.
Weaknesses
- **QGFN variants matter**. This paper does not provide a method to select different QGFN variants. Different variants of QGFN have very different performances in different environments. In Section 5, the authors claim that p-of-max is suitable for small action spaces, and p-quantile is suitable for large action spaces. What if different states have different action spaces? For example, in the graph combinatorial optimization problems [1], earlier states have a much larger action space than later states, and the action space size will change (decrease) during the sampling. I think compared with a fixed action space, this setting probably requires the use of different QGFN variants on different states during the sampling. - **Lack of more complex environment**. Therefore, I wonder about the QGFN's performance on graph combinatorial optimization problems, such as MIA. [1] - There are some missing recent works that deal with sampling high-reward candidates with diversity, such as [2, 3] - I am happy to raise my scores if my concerns are resolved. [1] Zhang, D., Dai, H., Malkin, N., Courville, A., Bengio, Y., & Pan, L. (2023). Let the flows tell: Solving graph combinatorial optimization problems with gflownets. arXiv preprint arXiv:2305.17010. [2] Chen, Y., & Mauch, L. (2023). Order-Preserving GFlowNets. arXiv preprint arXiv:2310.00386. [3] Jang, H., Kim, M., & Ahn, S. (2023). Learning Energy Decompositions for Partial Inference of GFlowNets. arXiv preprint arXiv:2310.03301.
Questions
See **Weaknesses**
Rating
5
Confidence
3
Soundness
3
Presentation
3
Contribution
3
Limitations
See **Weaknesses**
Thanks for your response. Regarding the additional $Q$, I think training $Q$ and GFN is only 5-10% more compared with standard GFN is very important to justify the inclusion of the additional $Q$. Because modified GFNs also require more training time than standard GFNs. I raise the score to 5. The reason I do not raise to the higher score is the following. > Even if QGFN took twice as long to train, if it means 3x more high-reward candidates. But the problem is that it does not have so much performance gain, especially when the lack of a stronger baseline. In the last part, I actually mean that this paper should contain more baselines like local search GFN, LED GFN, and compare QGFN directly with them (not incorporate $Q$ into them). Since they are the most up-to-date GFN variants and do not require an additional $Q$. The claim that $Q$ can improve these GFN variants is additionally from your rebuttal, not from the reviewers. Regarding these two claims, that "Q is better than GFN variants" and "Q can improve the GFN variants", the only result is that FL+Q is slightly better than FL. Again, I understand the limited rebuttal time can be an issue, but the first claim should not require much coding to verify given these methods are all open-source.
Summary
The paper focuses on improving high-reward sample collection, i.e., exploitation, in training GFlowNets. The motivation stems from the fact that the flow may pursue states that lead to many low-reward states rather than focusing on states that lead to high-reward states. To this end, the authors propose incorporating Q-value estimation, i.e., expected future reward, and interpolating between flow and Q to make a transition. Experimentally, the proposed method shows promising performance in fragment and atom-wise molecule generation and RNA-sequence generation.
Strengths
- The paper is well-written and easy to follow, with a well-developed motivation and idea. - It seems interesting that the compositional nature of the generation model can be considered to improve the performance of GFlowNets for practical applications. - This paper discusses various design choices for incorporating Q-value estimation, which offer different perspectives on "greediness." - Overall, the experiments are well-done, and the reported results demonstrate the effectiveness of the proposed model.
Weaknesses
No major weaknesses. See questions below.
Questions
- Are there failure scenarios where QGFN shows poor performance compared to GFN? (opposite of Figure 1) - I am curious if QGFN performs better than a policy interpolated between GFlowNets trained with low temperature and high temperature (with $p$). - Is there an analytic form of the sampling distribution of terminal states for the proposed methods? - It would be interesting to illustrate the sample efficiency for discovering modes in a toy example, e.g., hyper-grid. - Line 160: "good.." => "good."
Rating
7
Confidence
4
Soundness
3
Presentation
4
Contribution
3
Limitations
The authors have acknowledged the limitations of their approach.
Summary
The paper proposes jointly learning a $Q$ function and a policy network $P_F$ to improve the search for high-valued states when training GFlowNets. To achieve this, the authors develop three sampling strategies for composing $Q$ and $P_F$, $p$-greedy, $p$-of-max, and $p$-quantile, and show that the resulting algorithm, termed QGFN, often leads to faster discovery of modes relative to GFlowNet baselines. In spite of QGFN’s notable performance, I believe the paper would greatly benefit from a more clearly described solution and from the inclusion of stronger, optimization-focused baselines. I will be happy to increase my score if the enumerated weaknesses and questions below are properly addressed during the rebuttal period.
Strengths
1. **Well-described problem**. Section 3 highlights the issues of an exclusively GFlowNet-oriented approach to the search of high-valued states with clarity, namely, the large probability mass associated with low-reward states. 2. **Intuitively sound method with strong empirical results**. When the $Q$ function is accurately learned, QGFN — which interpolates between a GFlowNet (p = 0) and a DQN (p = 1) — should yield samples with higher rewards for larger $p$. This behavior is experimentally confirmed. However, the distribution from which QGFN samples, even when $Q$ is perfectly estimated, is mostly unclear. See weakness 4 below. 3. **Extensive experimental campaign**. (however, important baselines are missing) Experiments include four commonly used benchmark tasks for GFlowNets, and QGFN often outperforms alternative methods by a large margin.
Weaknesses
1. **Missingness of important baselines**. FL-GFlowNets (Pan et al., 2023) and LED-GFlowNets (Jang et al., 2024) exhibited strong performance for the optimization-focused tasks of molecule and set generation. Notably, contrarily to the greediness $p$ of QGFN, these methods do not require tuning of a hyperparameter that may unpredictably modify the sampling distribution. Authors should include FL-GFlowNets and LED-GFlowNets in Figures 2, 3, and 4. Have the authors considered increasing the sharpness of the distribution by modifying the temperature in LSL-GFN? 2. **Training of the $Q$-network is insufficiently described**. Lines 163-175 are hard to follow and Algorithm 1 in Section B does not provide sufficient details for understanding the training of QGFN. If I understand correctly, the adoption of multi-step returns with the horizon size set as the maximum trajectory length implies that the each $R(s_{t})$ in Equation (2) corresponds to the expected reward among $s_{t}$’s children under the chosen sampling policy. In any case, an unambiguous equation representing the learning objective for $Q$ should be included in Section 4. 3. **Difference between each sampling strategy is unclear**. In practical applications, a sampling strategy would have to be chosen, as selecting among the proposed approaches can be excessively time-consuming — even without retraining the model. However, $p$-greedy, $p$-of-max, and $p$-quantile seem to perform relatively different when the benchmark task is modified. A general approach for choosing a sampling strategy should, then, be proposed by the authors. Otherwise, the method is hardly usable in real applications. 4. **QGFNs lack theoretical guarantees**. Important questions are left unanswered. What if $Q$ is not properly trained? How do we know that $Q$ is inaccurately learned? What is the sampling for a given $p$ and sampling strategy? To improve QGFN’s reliability, a theoretical analysis should be included. Does a sufficiently trained QGFN finds, in average, more rewarding than standard GFlowNets? 5. **Rationale for choosing $p$ is not well-defined**. Experiments are filled with cryptically chosen numbers. In Table 1, the $p$ value for $p$-of-max and $p$-quantile are respectively set to $0.9858$ and $0.93$; in Figure 3, a “mode” is defined as a molecule with Tanimoto similarity smaller than $0.70$ and a reward larger than $1.10$. What is the rationale behind such numbers? How can a practitioner choose $p$? Also, if the choice of $p$ relied on a grid search over $N$ values in $[0, 1]$, a fair comparison with GFN-TB and GFN-SubTB in Table 1 should allow these methods to sample $N$x times more trajectories than QGFN. However, it is unclear whether this observation was considered for the experiments defining Table 1. 6. [Minor] Typos. Algorithm 1 refers to $\mathcal{L}\_{flow}$ and $\mathcal{L}\_{huber}$, which are undefined. I assume $\mathcal{L}\_{flow}$ is $\mathcal{L}\_{TB}$, but I could not find a definition for $\mathcal{L}\_{huber}$.
Questions
* Why does QGFN provide more diverse samples than, e.g., GFN-TB in Table 1 even when $p$ is set as high as $p = 0.9858$? Do we need to train a GFN? Given the large values of $p$ (above 0.9), it seems that the GFN is not very helpful. I wonder whether an untrained $P_F$ would produce similar results. * What is $k$ in line 230? * Figures 9 and 11 are quite difficult to grasp. Please consider visually encoding the greediness of the method as the size or opacity of points in a scatter plot. * Also, the convergence speed of GFlowNets via TB minimization is known to be dramatically affected by the learning rate of $\log Z$, which the authors set to $10^{-3}$. I wonder whether the relative results would remain the same if such learning rate was increased to $10^{-2}$ or $10^{-1}$.
Rating
4
Confidence
5
Soundness
2
Presentation
2
Contribution
2
Limitations
Authors claim that the computational overhead induced by learning a Q network and the sensitivity of sampling to the accuracy of $Q$ are the main limitations of their work.
Thank you for the detailed clarifications. I have no further comment.
Thank you for the clarifications. However, I think the question, "How to choose the best QGFN variants given different environment" is still unsolved. Given a state $s$ and the action space $A(s)$ at this state, there should be more detailed instruction on how to choose the QGFN variants and $p$. Regarding the best $p$, I have additional concerns below from the question from the first reviewer. From your rebuttal, the choice of $p$ in Table 1 is based on Figure 6. the chosen $p$ of p-greedy, p-of-max, and p-quantile QGFN are 0.4, 0.9858, and 0.93 respectively, which is very different from each other (and very strange numbers). According to Figure 6, the p-of-max close to the "Greediest" ($p=1$) will have inferior performance than "Least Greedy" ($p=0$). Therefore, the choice of an appropriate $p$ is non-trivial. (Also, The color lines for the left two and right two subfigures seem to be inconsistent, and for the 3rd subfigures, no numbers on the x-axis. ) I think choosing training $p$ using Figure 6 is not appropriate. if you are given a trained $P_F$ and $Q$, I think you can enumerate different $p$s at inference to select the best $p$ to combine them together at inference. However, as I reread the paper, the different $p$s are chosen in training, I do not think it is acceptable to enumerate all different $p$ in training since it will be very computationally expensive to retrain the model for every $p$ and not fair to other methods. I temporarily decrease my score to 4. If the above concerns are resolved. I am willing to raise my scores again.
Response
Thank you for engaging with us. We apologize for a mistake in our rebuttal to bgQq, inference time $p$ values are chosen based on Figure 9, not Figure 6. Figure 6's left two figures use these inference-time $p$ values, while Figure 6's 3rd subplot show the effect of different (fixed) $p$ values used for training. We understand your perspective, but we believe it comes from us not properly communicating the facts rather than what appears to be (and would be very concerning!) _strange numbers_. Below are our detailed responses to your questions. We hope these address your concerns. We are happy to answer any further questions you may have. Let's start with the choice of $p$. We make two claims: - A, $p$ can be treated like a hyperparameter, set once during training, - B, $p$ (and $\mu$) can be changed after training to _even further_ improve the reward-diversity quality of samples. Figures 2,3,4,5, and 6 (3rd subplot) show A. Treat $p$ as a hyperparameter, employ standard hyperparameter search. Even in this setting, QGFN shows promise. We hope this clarifies our approach to the reviewer. In a desire to show how malleable our method is, we went further and claimed B. This is shown in Figures 6 (two left figures) and 9 and Table 1. Without any retraining, we can get even more performance out of the models by tweaking one single parameter, $p$. **What is this "Greediest" to "Least Greedy" spectrum?** We chose this as an axis rather than 0-1 because 0-1 makes little sense for p-of-max. Instead (and we wrote this in l290-292 but clearly need to make this more prominent in the text), we use 0.9-1 for p-of-max and 0-1 for p-quantile and p-greedy. Below 0.9, p-of-max in practice samples from $P_F$ so it would have been a waste of space to show the 0-0.9 values. To overlay all the results on a single plot, it made more sense to us to simply show the evolution of behavior as a function of "greediness" rather than put too many details and multiple overlapping axes. This in retrospect may have been a confusing choice for the reader, and we will change this to make it clearer. Next, while 0.9853 may seem unusual and arbitrary, it is computed during inference with $p$ chosen from `np.linspace(0.9, 0.999, 16)`. Similarly, 0.93 is the second to last value of `np.linspace(0, 1, 16)`. > I think choosing training $p$ using Figure 6 is not appropriate The 3rd subplot of Figure 6 does **not** show $p$ determined **after** training, it shows the result of a hyperparameter search for $p$. Why would it not be appropriate to chose a training hyperparameter from the result of this? > I do not think it is acceptable to enumerate all different $p$ in training We believe we perform standard hyperparameter search, $p$ is chosen beforehand before any training occurs. This means $p$ is treated like any other hyperparameter, like learning rate, or batch size. This is a common practice in method papers. We want to clarify that we are **not** adjusting $p$ during training based on inference performance. Claim B is entirely about what happens _after_ having trained models during inference. This distinction between Claims allows our approach to be computationally feasible and fair when compared to other methods. Note that for baselines, we _also_ do a hyperparameter search on greed-controlling parameters, notably $\beta$. Part of the results of this search are shown in Figure 6's left two subplots. This is fair standard practice. Finally on this point, we apologize for the color confusion of Figure 6; we will correct this in the revision. > there should be more detailed instruction on how to choose the QGFN variants and $p$ We did not explicitly state this, as our focus is on showing the usefulness of the _existence_ of $Q$-$P_F$ mixtures, which we think is a valuable NeurIPS-level scientific contribution in itself, rather than advising on a specific approach. However, we understand that readers may be curious about which variant to use, and we reiterate what we write in our rebuttal, and will emphasize in our paper that starting with $p$-greedy (with $p=0.5$) is a reliable option. We speculate in our paper that variants have effects which partially depend on the action space because of the nature our tasks. These speculations are based on evidence, but instructions for the general case would require vastly more evidence, more compute, on many more types of action spaces than what we have. We hope to impress upon the reviewer that this would be worthy of an entirely different 9-page paper. We have introduced a novel method, shown that it works, and that it works because of the hypothesized mechanism, and went beyond that and showed that many variants of the idea are both sensible and capable of producing interesting results. We agree with the reviewer that this paper would be stronger if we had a general recipe, but providing such a comprehensive solution would be beyond the intended scope and contribution of this paper.
Thanks for your response. However, although $p$ can be seen as a hyperparameter in a standard hyperparameter search, it is still an additional hyperparameter compared with other common hyperparameters shared with all GFN methods, such as batch size, learning rate and $\beta$. Also, for each iteration of search on $p$, one GFN network and one Q network needs to be trained, it is additional burden compared with GFN only. Comparing with modified GFN, such as local search GFN, LED-GFN, QGFN seems to be more computationally expensive by an additional Q, and the paper and the rebuttal do not provide a comparison in the experiments (although this paper cited local search GFN). In the rebuttal to the Reviewer bgQq, the authors said that Q function could help improve these modified GFN but no experiments except for FL-GFN, a relatively weaker baseline, either. Therefore, I maintain my current score.
Response
Thank you again for engaging with us. We'd like to reiterate that our research methodology aligns with standard practices in the ML field. We understand your concern, and you are correct that $p$ is an additional hyperparameter, but we would like to emphasize that the vast majority of new methods that the field adopts come with new hyperparameters. Additionally, please note that we do tune the relevant hyperparameters of our baselines, and perform extensive investigations into the effects of the new aforementioned hyperparameter. > for each iteration of search on $p$, one GFN network and one Q network needs to be trained, it is an additional burden compared with GFN only. We agree that training both $Q$ and GFN is a valid concern, in principle, but as we've noted here and in our paper's _Limitations_, training $Q$ and GFN does not take much more time; specifically only 5-10% more, which did not seem worth reporting precisely. This negligibly increased training time may be well worth it given the extra performance. We again would like to note that this is standard practice in ML, to pay more compute to get more performance. Even if QGFN took twice as long to train, if it means 3x more high-reward candidates, it seems to us like a very reasonable trade-off. > Comparing with modified GFN, such as local search GFN, LED-GFN, QGFN seems to be more computationally expensive by an additional Q, and the paper and the rebuttal do not provide a comparison We take note of your point that we do not report runtimes in the paper, and will be sure to include this in our revision. > Q function could help improve these modified GFN but no experiments except for FL-GFN, a relatively weaker baseline, either We ran FL-GFN as per your suggestion for testing on MIS. Given the limited time available, we used the public codebase which contains FL-GFN, and we did see improvement. This is additional evidence (we e.g. ran QGFN with SubTB and FM as shown in Figure 12 and 13) that QGFN can apply to any GFN. Unless there is a specific reason for which you believe that QGFN would not improve GFN variants such as LED-GFN, it would be beyond the scope of this paper to expect an implementation of all the now numerous GFN variants. A priori, QGFN would improve them, and we've demonstrated evidence of the mechanisms through which QGFN would achieve that. LED-GFN, for example, is not immune to the example given in Figure 1 of our paper (which is analogous to real-world dense-reward settings). We urge you to reconsider your evaluation in light of QGFN's contributions and the evidence presented in this paper.
Clarification on $p$ and other values
We apologize for a mistake in our rebuttal, inference time $p$ values are chosen based on Figure 9, not Figure 6. Figure 6's left two figures use these inference-time $p$ values, while Figure 6's 3rd subplot show the effect of different (fixed) $p$ values used for training.
Thank you for the responses and additional experiments. Also, sorry for the delayed response. > 0.7 has been used Thank you for the clarification. Some works (e.g., [1, 2]) do not seem to explicitly report this when training a GFlowNet on the QM9 environment. A reference would be helpful here. > This is incorrect, because we are reporting averages. As a thought experiment, if we sampled infinitely many samples for them, the average reward and similarity would converge to some values. Indeed, for infinitely many samples, all reported metrics would perfectly reflect the trained model’s underlying distribution. Your results, however, only cover the finite case representing a small fraction of the state space. For the average reward in Table 1, the difference may be harmless; for the number of modes in Figure 4, the difference is critical. > We are unsure what the reviewer means My question is related to the comment above. If I implement, e.g., a grid search with $N$ candidates to find the best $p$ for a given task, and in each run I sample $T$ trajectories from the trained model, then I would have sampled $NT$ trajectories in total, right? When my interest lies in finding high-valued and diverse samples, will this approach be more effective than directly sampling $NT$ trajectories from a standardly trained GFlowNet? The results in the paper do not seem to provide evidence for this. In other words, the question I ask is: if $p$ cannot be generally chosen prior to inference, is QGFN beneficial given the extra computation cost of hyperparameter selection? > Style of Fig. 9&11 Thank you. In my view, the figure attached to the rebuttal PDF is significantly better than the ones reported in the main text. While I understand that the discussion period is near the end and that there is no time remaining for further discussion, I believe that crucial aspects of the author’s work, e.g., fair comparison to baselines and clarity of presentation, were not properly addressed. [1] Learning to Scale Logits for Temperature-Conditional GFlowNets. Kim et al. ICML 2024. [2] Towards Understanding and Improving GFlowNet Training. Shen et al. ICML 2023.
Decision
Accept (poster)