Optimizing Backward Policies in GFlowNets via Trajectory Likelihood Maximization

[Question 4] Finally, combining constraint reinforcement learning with TLM is a valuable direction for further work since it allows us to explicitly incorporate challenging types of constraints, e.g., the correctness of the molecules in the QM9 task, in a more principled way. However, to the best of our knowledge, such approaches are yet to be explored in GFlowNet literature. That is why we leave this question as an interesting direction for further studies.

References:

[1] Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, Pierre-Luc Bacon. Maximum entropy GFlowNets with soft Q-learning. AISTATS 2024

[2] Hyosoon Jang, Yunhui Jang, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Pessimistic Backward Policy for GFlowNets. arXiv preprint arXiv:2405.16012

[3] Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, Yoshua Bengio. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. NeurIPS 2021

[4] Moksh Jain, Emmanuel Bengio, Alex-Hernandez Garcia, Jarrid Rector-Brooks, Bonaventure F. P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, Yoshua Bengio. Biological Sequence Design with GFlowNets, ICML 2022

[5] Yihang Chen, Lukas Mauch. Order-Preserving GFlowNets. ICLR 2024

We thank the reviewer for their valuable suggestions and are happy to provide further details in response to their questions.

[Weakness 1] Firstly, we would like to respectfully disagree with the first stated weakness and note that we do directly compare in all experiments with the maxent backward approach proposed in [1] and discuss this approach in detail in our paper (Section 2.3). As for [2], we note that at the moment of the ICLR 2025 submission deadline, the paper [2] had not yet made its code publicly available, and the implementation details provided in the arXiv preprint were limited, making it challenging to reproduce the presented approach. After the ICLR 2025 submission deadline, the authors published the code for their method and experiments. Although it will be unfeasible to make a careful and faithful comparison during the time allocated for rebuttal, we are more than willing to include a comprehensive comparison to the approach of [2] in the camera-ready version of our paper.

[Weaknesses 1-3, Question 1] Secondly, we note that sEH is a challenging molecule generation environment, introduced in the seminal work on GFlowNets [3] and used as a benchmark in many GFlowNet papers, with the number of states $\approx 10^{16}$ and number of actions from $100$ to $2000$ (see [3], Section 4.2). There exists a simplified version of the environment in GFlowNet literature with a reduced number of fragments and thus reduced state space, see, e.g. [5]. However, we emphasize that this is not the setting considered in our submission. While we highlight in our paper that sEH is a more structured environment than QM9, it does not mean it is simple. No method obtains the correlation with a reward higher than $0.75$ on sEH in the presented experiments, meaning that the environment is far from being solved exactly.

However, we appreciate the reviewer’s comments on sEH task and the differences in TLM performance between different environments. We added a detailed discussion of the experimental results to the end of Section 4 to improve the clarity of our paper regarding these points. Concerning the poor performance on sEH, we argue that our experimental results align with the hypothesis by [1], which states that backward policy optimization is more beneficial in environments with less structure, like QM9. While for highly structured tasks like sEH, considering non-trivial backward approaches is of little utility and using a simple uniform backward policy generally shows strong performance. We refer the reviewer to the updated manuscript for a more detailed discussion.

[Weakness 4, Question 3] We appreciate the question about combining our approach with maxent backward [1]. Unfortunately, the direct combination is impossible because there exists a unique backward policy that maximizes the trajectory entropy, as it was proven in [1], and the essence of their approach is to obtain such a policy exactly. Our method focuses on training the backward policy to better align with the current forward policy in terms of trajectory distributions. However, a possible idea could be to add the entropy of the backward policy as a regularizer to our objective, promoting both diversity and alignment to the forward policy. We thank the reviewer for this valuable suggestion and would like to consider it as an interesting future research direction.

[Weakness 5, Question 2] Regarding the non-stationary reward problem, we would like to highlight that the TLM method introduces the non-stationary rewards to the RL formulation of the GFlowNet training problem, as does any backward policy optimization method. Indeed, the RL formulation of GFlowNet training considers the backward probabilities as a part of the intermediate reward functions. In our work, it was the only context in which we have mentioned the non-stationary RL problem, and we experimentally demonstrated that TLM handles it well. At the same time, the setting of the non-stationary GFlowNet reward, which is defined only for terminal states, was never explored in the literature directly; however, it naturally appears in training a reward proxy in an active learning setup, see [4], and we do think that it might be an interesting future extension of our work.

Dear Reviewer EZ4M, Could you please read the authors' rebuttal and give them feedback at your earliest convenience? Thanks. AC

As for the hyperparameter and neural network architecture choices, for the most part, we indeed utilize the choices established in previous literature [3, 4, 5] to follow their setups faithfully. Note that the goal of our experimental evaluation is not necessarily to obtain the best possible metrics in all of the tasks (which could be potentially improved in many ways, e.g., by taking larger neural networks) but rather to produce a direct comparison of various approaches to selecting the backward policy in a fair and controlled setup.

We thank the reviewer for the questions about plots with varying learning rates and their connection to the full plots in Appendix, which indeed needed to be discussed in the paper in sufficient detail. For bit sequence and molecule generation tasks, we consider the initial learning rate as a hyperparameter and study the dependence of the results for all methods on the choice. Figure 3 does not present the results for single training runs depending on the learning rate decay across the run, but the final results for different training runs with varying initial learning rates (depicted on the horizontal axis). Figure 2, on the other hand, illustrates how the number of discovered modes increases over the training duration for runs with fixed initial learning rates, while curves depending on the varying initial learning rates are presented in full plots in the Appendix in Figure 5 and Figure 6. This allows one to directly examine the robustness of various approaches with respect to the choice of learning rate, thus having a benefit over the type of presentation where one simply chooses the best learning rate for each method. We note that such a presentation was also used in some of the previous works [4].

References:

[1] Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, Pierre-Luc Bacon. Maximum entropy GFlowNets with soft Q-learning. AISTATS 2024

[2] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518, 2015

[3] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. NeurIPS 2022

[4] Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. ICML 2023

[5] Daniil Tiapkin, Nikita Morozov, Alexey Naumov, Dmitry Vetrov. Generative Flow Networks as Entropy-Regularized RL. AISTATS 2024

[6] Shangtong Zhang, Hengshuai Yao, Shimon Whiteson. Breaking the Deadly Triad with a Target Network. ICML 2021

[7] Hyosoon Jang, Yunhui Jang, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Pessimistic Backward Policy for GFlowNets. arXiv preprint arXiv:2405.16012

We want to express our gratitude to the reviewer for the very detailed review, the time and effort spent reading our paper, the positive feedback, and the valuable suggestions.

First, we would like to stress one of the main points of our work outlined in our Introduction: considering backward policy optimization in GFlowNets from the RL perspective and closing the gap related to it in the theory that connects GFlowNets and RL. The existing theory previously only considered the setting of a fixed backward policy. This point was never explored in the previous works, so the novelty and the motivation for our TLM algorithm directly arise from it. [1] utilizes RL as a tool to find their maxent backward policy, and the connection to RL is made only in the case of using this exact unique maxent policy. [7] does not consider the RL perspective at all. The scope of our work is to explain the simultaneous optimization of the forward and backward policies from the RL perspective and empirically show the benefits of our TLM algorithm, which is grounded in this perspective. We thank the reviewer for pointing out this weakness in the presentation. We added a corresponding clarification to Section 2.3.

We also highly appreciate the comments about unclear discussions and a lack of explanations of some experimental results. As suggested, we added a more detailed Discussion subsection to the end of Section 4 to improve the clarity regarding these points, where we consider two hypotheses. In regard to the poor performance on sEH, we argue that our experimental results are in line with the hypothesis put out by [1], which states that backward policy optimization is more beneficial in environments with less structure, like QM9. While for highly structured tasks like sEH, considering non-trivial backward approaches is of little utility and using a simple uniform backward policy generally shows strong performance. Regarding the performance differences when TLM is combined with various forward policy training objectives, we hypothesize that local GFlowNet training objectives defined on individual transitions, e.g., DB, benefit more from the combination with TLM. TLM propagates global information over whole trajectories, resulting in a good synergy with such local objectives. Meanwhile, TB, which is also defined on whole trajectories, shows fewer improvements when combined with TLM. We refer the reviewer to the updated manuscript for a more detailed discussion.

As for the metrics, reward correlation and the number of found modes are standard metrics that are widely utilized in GFlowNet literature, especially in molecule generation environments [1,3,4,5]. The use of two metrics is justified by the different aspects they are designed to measure:

1. Reward correlation measures how well the distribution over the target space induced by the trained forward policy matches the target reward distribution, i.e., how well the sampling problem is solved.
2. The number of found modes measures how many diverse high-reward objects the model discovers during training, which is the quantity we are interested in scientific applications but the one that does not directly measure how well we fit the target distribution, rather focusing on the exploration of the environment.

Regarding the similarities between our approach, the EM algorithm, and the Wake-Sleep Algorithm, we meant the structure of alternating minimization. This remark was made to underline that such a way to address the problem with a natural block structure is well-known in the literature. We would be happy to incorporate any changes the reviewer might suggest or even remove this remark.

As for the target networks for backward policies, it is a standard technique from reinforcement learning literature [2] utilized to introduce stability into the training procedure (which is required by our Theorem 3.1 and additionally explained in Appendix A.2). Online network is exactly the backward policy network that is trained via TLM objective. The simplest intuition behind using the target network in algorithms similar to Q-learning is that it prevents the target values from changing too fast, stabilizing the training process. Toward this aim, we maintain a separate, slowly updated copy of the backward policy network. Some theoretical results (see, e.g. [6]) provide theoretical support for the fact that “a target network stabilizes training” using a framework of the two-timescale stochastic approximation. However, we acknowledge that this fact is a common empirical wisdom rather than a theoretically established result.

Dear Reviewer A6Yj, Could you please read the authors' rebuttal and give them feedback at your earliest convenience? Thanks. AC

We thank the reviewer for their valuable suggestions and are happy to provide further details in response to their concerns.

We agree with the importance of comparative evaluation to previous works. However, we note that at the moment of the ICLR 2025 submission deadline, the pessimistic backward policy paper [1] had not yet made its code publicly available, and the implementation details provided in the arXiv preprint were limited, making it challenging to reproduce the presented approach. After the ICLR 2025 submission deadline, the authors published the code for their method and experiments. Although it will be unfeasible to make a careful and faithful comparison during the time allocated for rebuttal, we are more than willing to include a comprehensive comparison to the approach of [1] in the camera-ready version of our paper.

Second, we note that we do present an ablation study of different components of our method in Appendix A.2. We try to separately turn off various components of TLM, e.g. backward policy target networks, and experimentally verify that they have a positive effect on the performance.

Then, we would like to disagree about the presentation in Figure 1. DB, TB, and SubTB are indeed different existing GFlowNet objectives we utilize to train the forward policy, each having its own merits and performing better or worse between different tasks. One of the points of our experimental evaluation was to investigate how backward policy training with TLM would perform when paired with different forward policy training objectives, which can provide valuable insights for understanding our method and its utility. Just showing the best-performing forward policy training method would conceal some relevant information from the reader, and we note this as a limitation of the experimental evaluation provided in previous works on backward policy training [1, 2].

Finally, we highly appreciate the suggestion to provide a more detailed explanation for the performance of TLM with different choices of the GFlowNet training objective, which is related to our previous point presented above. We added a detailed discussion of our experimental results to the end of Section 4 of our paper. In summary, we hypothesize that local GFlowNet training objectives defined on individual transitions, e.g., DB, benefit more from the combination with TLM, which in turn propagates global information over whole trajectories, resulting in a good synergy. Meanwhile, TB, which is also defined on whole trajectories, shows fewer improvements when combined with TLM. We refer the reviewer to the updated manuscript for a more detailed discussion.

References:

[1] Hyosoon Jang, Yunhui Jang, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Pessimistic Backward Policy for GFlowNets. arXiv preprint arXiv:2405.16012

[2] Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, Pierre-Luc Bacon. Maximum entropy GFlowNets with soft Q-learning. AISTATS 2024

Dear Reviewer 6Sem, Could you please read the authors' rebuttal and give them feedback at your earliest convenience? Thanks. AC

We thank the reviewer for their positive feedback and are happy to provide further details in response to their concerns.

First, we would like to disagree with the comment about the toy nature of sEH and QM9 tasks. sEH is a challenging molecule generation environment, introduced in the seminal work on GFlowNets [1] and used as a benchmark in many GFlowNet papers, with the number of states $\approx 10^{16}$ and number of actions from $100$ to $2000$ (see [1], Section 4.2). There exists a simplified version of the environment in GFlowNet literature with a reduced number of fragments and thus reduced state space size, see, e.g. [2, 6]. However, we emphasize that this is not the setting considered in our submission. As for the QM9 task, it can be small and toy if the generation is done in a fragment-based fashion, as in [2]. However, we do the generation atom-by-atom and bond-by-bond, as in [3], which results in a more challenging task with a larger and less structured state space. We can approximate the number of states in it to be $\approx 4.5 \cdot 10^{14}$ and the number of actions to vary from $6$ to $100$.

In regard to the poor performance on sEH, we argue that our experimental results are in line with the hypothesis put out by [3], which states that backward policy optimization is more beneficial in environments with less structure, like atom-based QM9. While for highly structured tasks like sEH, considering non-trivial backward approaches is of little utility and using a simple uniform backward policy generally shows strong performance. We added a more detailed Discussion subsection to the end of Section 4 in our paper to improve the clarity regarding this point.

Second, we appreciate the suggestion to extend our results to other benchmarks. However, we would like to note that the suggested biological sequence generation environments from [4] are autoregressive, meaning that all possible actions are to append a token to the end of the sequence. This results in a tree-structured environment where each state has exactly one parent; thus, only one trivial choice of a backward policy exists, and its training is redundant. Other works on GFlowNets have made attempts at considering non-tree structured environments for biological sequence generation [5], but such environments are not well-established yet in the GFlowNet literature. We would like to highlight that the careful design of a new environment is a challenging problem, which is of independent interest and requires a lot of effort. That is why we consider it to be an interesting further research direction.

References:

[1] Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, Yoshua Bengio. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. NeurIPS 2021

[2] Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, Tommaso Biancalani. Towards Understanding and Improving GFlowNet Training. ICML 2023

[3] Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, Pierre-Luc Bacon. Maximum entropy GFlowNets with soft Q-learning. AISTATS 2024

[4] Moksh Jain, Emmanuel Bengio, Alex-Hernandez Garcia, Jarrid Rector-Brooks, Bonaventure F. P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, Yoshua Bengio. Biological Sequence Design with GFlowNets, ICML 2022

[5] Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. ICML 2023

[6] Yihang Chen, Lukas Mauch. Order-Preserving GFlowNets. ICLR 2024