Learning Energy Decompositions for Partial Inference in GFlowNets

Dear reviewer qpM8,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in $\color{blue}{\text{blue}}$.

In what follows, we address your comments one by one.

---

Q1-Method. Is LED-GFN also applicable for the tasks where the order of actions might matter?

Yes, to the best of our knowledge, all the GFlowNet frameworks (including LED-GFN) are applicable to scenarios where the order of actions matters. To this end, one can define the state to incorporate the order of actions, e.g., a state $s_{t}$ information may include the feature $[(a_0,0),\ldots,(a_{t-1},{t-1})]$ to ensure the order of actions-dependent GFlowNets.

---

Q2-1-Method. Why do we want to smooth potentials? In this case, can we just simply set each potential to $\frac{\mathcal{E}}{T}$?

As mentioned in Section 3.1, we want to smooth potentials to prevent heavily relying on a specific transition, e.g., $\mathcal{E}(s_T)=\phi_{\theta}(s_{T-1}\rightarrow s_{T})$. However, we do not set $\phi_{\theta}(s \rightarrow s') = \mathcal{E}(s_T)/T$ for a specific trajectory, as the potential is associated with multiple trajectories $\\{\tau | (s\rightarrow s') \in \tau\\}$. In other words, the potential $\phi_{\theta}(s \rightarrow s')$ should be designed to consider:

$

\mathcal{E}(s_{T})\approx\sum_{(s_{t}, s_{t+1}) \in \tau}\phi_{\theta}(s_{t}\rightarrow s_{t+1})\quad \text{for all } \tau=(s_{0},\ldots,s_{T}) \in \\{ \tau | (s\rightarrow s') \in \tau \\}.

$

Then, one can observe that the condition $\phi_{\theta}(s \rightarrow s') = \mathcal{E}(s_T)/T$ for a specific trajectory would not be optimal for another trajectory. In Figure 10 of our updated manuscript, we add such a baseline, i.e., set $\phi_{\theta}(s \rightarrow s') = \mathcal{E}(s_T)/T$ for a given trajectory at each step. One can see that our learning method still exhibits the most competitive performance.

---

Q2-2-Method. When $\mathcal{E}(s\rightarrow s')=0$, FL-GFN reduces to the DB, but is it common case in many tasks?

Yes, it can be a common case in real-world settings. Especially, when the energy evaluation of the intermediate state may be impossible, e.g., wet-lab evaluation for incomplete molecules, the change of energy would be evaluated as zero.

---

Q3-Method. LED-GFN might learn $P_F(s_{t+1}|s_t)\propto \phi_{\theta}(s_t\rightarrow s_{t+1})$, but can LED-GFN sample $x$ with probability proportional to $\exp{(-\mathcal{E}(x))}$?

We would like to clarify that LED-GFN learns a policy to be $P_F(s_{t+1}|s_t)\propto {F}(s_t\rightarrow s_{t+1})$ since our objective (Equation (4)) just reparameterizes the flow of the GFN with the potential: $\log{F}(s_t)=\log\tilde{F}(s_t)-\sum_{u=0}^{t-1}{\phi_{\theta}(s_{u}\rightarrow s_{u+1})}$. Hence, LED-GFN can sample $x$ with probability proportional to $\exp{(-\mathcal{E}(x))}$.

---

Q1-Experiment. In Figure 9(b), what is the number of calls, and why is the number of calls for LED-GFN the same as that for GFN?

The number of calls is the number of energy evaluations. LED-GFN has the same number of calls as the GFN since the training of the potential function does not require additional energy evaluation, as it learns from the trajectories and energies that have already been evaluated during GFN training (mentioned in the last paragraph of Section 3.2).

---

Q2-Experiment. Ablation studies on various dropout rates $1-\gamma$ for Figure 9(a) are missing. How does the dropout rate affect performance?

To address your comment, we incorporate experiments with dropout rates $10\\%$, $20\\%$, and $40\\%$ for the set generation and incorporate the results in Figure 9(a) of our updated manuscript. One can see that (1) applying dropout benefits credit assignment, (2) but too high a dropout rate may result in inefficient training.

---

Additional minor issues in writing.

Thank you for pointing this out! We update our manuscript to follow your valuable suggestions.

• We add $\log$ for $P_B$ in the DB objective.
• We clarify that we actually consider the subTB($\lambda$).
• We modify $\sum_{t=0}^{u}\phi(s_{t})$ and $\sum_{t=0}^{u+1}\phi(s_{t})$ to $\sum_{t=0}^{u-1}\phi(s_{t}\rightarrow s_{t+1})$ and $\sum_{t=0}^{u}\phi(s_{t}\rightarrow s_{t+1})$, respectively.
• We fix 'SubTB' to 'subTB' in Figure 4.
• We fix 'FL-GFN' and 'LED-GFN' in Figure (5) and (6) to 'FL-subTB' and 'LED-subTB', respectively.

W6. Plots with different markers or line styles would be helpful.

Thank you for your valuable suggestions! We modify all plots to incorporate different markers or line styles in our updated manuscript.

---

Q1-1. On alternative potential learning (Figure 10), why learning $\phi(s\rightarrow s')$ (ours) is better than learning an extension of $\mathcal{E}$ to nonterminal states (simple model)?

We first clarify that the 'simple model' in Figure 10 uses a model trained to predict the terminal energy from the terminal state but not trained with the intermediate states. Here, we assume that learning $\phi(s\rightarrow s')$ has shown better results, as the 'simple model' may return not informative values for intermediate states that are not considered during training.

---

Q1-2. On alternative potential learning (Figure 10), why do you need to learn a proxy model to predict terminal energies (simple model) instead of using the true energy function?

We consider such a baseline since it can be a solution to address one of our motivations: true energy evaluation for intermediate states can be expensive (described in Section 3.1). Since this baseline can predict the intermediate state energy as a model-based approach, there is no need to evaluate the true energy for the intermediate state.

---

Q1-3. Why do you not regress each $\phi(s\rightarrow s')$ to $\frac{1}{T}\mathcal{E}(s_T)$?

We do not set $\phi_{\theta}(s \rightarrow s') = \mathcal{E}(s_T)/T$ for a specific trajectory since the potential is associated with multiple trajectories $\\{\tau | (s\rightarrow s') \in \tau\\}$. In other words, the potential $\phi_{\theta}(s \rightarrow s')$ should be designed to consider:

$

\mathcal{E}(s_{T})\approx\sum_{(s_{t}, s_{t+1}) \in \tau}\phi_{\theta}(s_{t}\rightarrow s_{t+1})\quad \text{for all } \tau=(s_{0},\ldots,s_{T}) \in \\{ \tau | (s\rightarrow s') \in \tau \\}.

$

Then, one can observe that the condition $\phi_{\theta}(s \rightarrow s') = \mathcal{E}(s_T)/T$ for a specific trajectory would not be optimal for another trajectory. In Figure 10 of our updated manuscript, we add such a baseline, i.e., set $\phi_{\theta}(s \rightarrow s') = \mathcal{E}(s_T)/T$ for a given trajectory at each step. One can see that our learning method still exhibits the most competitive performance.

---

Q2. Can you apply your approach to the TB-based objective, e.g., reparameterize logits in the policy?

Thank you for the interesting suggestion! We follow your valuable suggestion by conducting experiments with modified TB that reparameterizes the logits of the policy with the potential. We incorporate corresponding results in Appendix E.3 of our updated manuscript. One can see that our LED-GFN with your suggested approach improves the TB. We assume that the local credits support assigning a high probability to the action responsible for the low terminal energy.

---

Reference

[1] Pan et al., Better Training of GFlowNets with Local Credit and Incomplete Trajectories, ICML 2023

Dear reviewer RMo5,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in $\color{blue}{\text{blue}}$.

In what follows, we address your comments one by one.

---

W1-1/W1-3. Rather than "partial inference of GFlowNets", "partial inference in/for GFlowNets" is more proper for what is being done in this paper.

We appreciate your detailed comments to improve our paper! In our updated manuscript, we incorporate your valuable suggestion by modifying "partial inference of GFlowNets" to "partial inference in GFlowNets".

---

W1-2. Section 2.2 attributes "partial inference" to FL-GFN [1], but that paper does not introduce this term.

We would like to clarify our terminology stems from FL-GFN [1] which states "partial inferences that do not contain a full specification of the latent variables that explain the given input" in the last sentence of paragraph six of the introduction section. To alleviate your concern, we modify our description to be more precise by stating "Pan et al. (2023a) first incorporate this concept for training GFlowNets".

---

W2. There are small errors in the exposition on GFlowNets.

Thank you for your insightful comments! We fix Section 2 by following all of your detailed comments.

• In Detailed balance of Section 2.1, we fix an error at the top of p.3 (which states $\exp(-\mathcal{E}(x))$ as the energy) by modifying it as the "exponent of the negative energy, i.e., the score of the object".
• In Sub-trajectory balance of Section 2.1, we clarify that we actually consider subTB$(\lambda)$ which interpolates between the DB and the TB.
• In Forward-Looking GFlowNet of Section 2.2, we clarify that FL assumes an extension of $\mathcal{E}$ to nonterminal states.

---

W3. Limitations/costs/failure modes of the proposed algorithm are not discussed.

We incorporate your valuable comments in our new discussion section (Appendix D).

As limitations, our method requires additional time and space costs to learn a potential function $\phi_{\theta}$. However, the increment is minor for real-world problems with prohibitively expensive energy evaluation, such as the drug discovery task that requires evaluating with docking tools or wet-lab experiments.

For detailed costs, we report the additional time costs of LED-GFN by comparing with the GFlowNets. In Appendix E.2 of our updated manuscript, we (1) provide time costs in molecular generation task, where the conventional subTB objective takes $5.80$ seconds for obtaining $100$ samples during training while our LED-GFN takes $6.61$ seconds, and (2) discuss the efficiency with respect to the time costs.

As failure modes, one can consider an extreme case where most actions are meaningless, i.e., they do not affect energies or the terminal state. In this case, the smoothed local credits are not meaningful. This extreme case can be avoided with well-designed actions and state spaces.

---

W4-1. What is the form of the model $\phi_{\theta}$? Does it share some parameters with the flow and policy models?

In all experiments, we design the potential functions $\phi_{\theta}$ to have the same architecture as the flow model, e.g., a feedforward network with the same hidden dimensions, where the input dimension is extended to consider the transition $s_t \rightarrow s_{t+1}$. We modify the first sentence of Appendix C to better highlight these points. Additionally, the potential model $\phi_{\theta}$ does not share any parameters with the flow and policy models.

---

W4-2. In order to be useful, how simple or complex the energy decomposition model needs to be, relative to the policy model?

Empirically, we find that the complexity of the energy decomposition model is sufficient when its complexity matches that of the flow (or policy) models. Intuitively, one can conclude that estimating the decomposition of an energy function is as challenging as learning the energy function itself, which is also necessary for the flow model.

---

W5. Some theoretical characterization of the optimal energy decomposition would be helpful.

Indeed, our work lacks theoretical aspects since we mainly focus on improving the performance of GFlowNets in practice.

As our best effort to incorporate your comments, we provide the theoretical characterization of the optimal energy decomposition in Appendix F of our new manuscript. To be specific, we derive the necessary conditions for the potential functions to perfectly minimize the energy decomposition loss. The result implies a positive correlation between a potential function evaluated on a state transition and the energies of the terminal states reachable from the state transition. This hints at the meaningful-ness of the local credits provided by the potential function.

Dear reviewer RMo5,

Thank you very much for the positive response. We think your comments were very helpful in improving our paper, and we are happy to hear that our efforts have addressed most of your concerns!

For additional comments, we agree that "new theoretical analysis lacks a clear explanation (but the theory is not essential for our paper)". Therefore, instead of incorporating the current analysis in the manuscript, we will consider the suggested analysis in our future manuscript as possible, while leaving the in-depth analysis as a future work (stated in the last sentence of Section 5 of the updated manuscript).

Dear reviewer q5Nd,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in $\color{blue}{\text{blue}}$.

In what follows, we address your comments one by one.

---

Q1-1. In Figure 10, is LSTM comparable to a GFlowNet-based model given its simplicity?

The 'LSTM' baseline in Figure 10 is comparable since it is also an energy decomposition-based GFlowNet model. It just uses the LSTM model to predict the potential $\phi_{\theta}(s\rightarrow s')$ to enable partial inference in GFlowNets.

Additionally, we would like to reclarify that Figure 10 shows the comparison with alternative energy decomposition methods for partial inference in GFlowNets, i.e., alternative methods to predict the potential $\phi_{\theta}(s\rightarrow s')$ for GFlowNets. In our updated manuscript, we modify the description of corresponding ablation studies (last paragraph of Section 4.5) to better reflect this point!

---

Q1-2. Why not test other baselines for energy decompositions?

We do not conduct extensive comparisons between energy decomposition methods since (1) there are no prior studies about energy decomposition-based GFlowNets and (2) our main focus is first to show the effectiveness of energy decomposition in GFlowNets.

Nevertheless, to further address your comment, we add an additional energy decomposition baseline by modifying a reward redistribution method in reinforcement learning [1] to decompose the energy for GFlowNets. In Figure 10 of our updated manuscript, one can see that our method still shows the most competitive performance.

If you provide information about an essential baseline, we would be happy to incorporate it!

---

Q2. Why not consider the other objectives beyond DB-based and subTB-based objectives?

The other conventional objective, i.e., TB, is not considered since TB is defined with the full trajectory but our goal is to enable partial inference for the local trajectory-based objectives, i.e., DB and subTB. However, it is worth noting that our new experiments in Appendix E.3 show that our energy decomposition even improves the TB-based objectives. For other objectives extending DB or subTB [2], combining LED-GFN with them can be an interesting future research direction.

---

Reference

[1] Gangwani et al., Learning Guidance Rewards with Trajectory-space Smoothing, NeurIPS 2020

[2] Pan et al., Generative Augmented Flow Networks, ICLR 2023

Dear reviewer xgzq,

We express our deep appreciation for your time and insightful comments. In our updated manuscript, we highlight the changes in $\color{blue}{\text{blue}}$.

In what follows, we address your comments one by one.

---

W1. There are no theoretical guarantees that the learned decomposition provides meaningful local credits in all settings.

Indeed, since our main focus is on improving the performance of GFlowNets in practice, our work lacks a theoretical guarantee. We provide our best efforts to alleviate your concerns in Appendix F of the updated manuscript and what follows.

To this end, we derive a fixed form condition for the potential functions that imply the meaningful-ness of our local credit via the correlation between potentials and the terminal energy (similar to [1]). To be specific, in Appendix F, we derive

$

\phi_{\theta}(s\rightarrow s') = K_{1}\sum_{\tau \ni (s\rightarrow s')}\mathcal{E}(s_{T}) + K_{2},

$

where the summation over $\tau$ is defined on trajectories that include $s\rightarrow s'$ as its intermediate transition and $s_{T}$ is the terminal state of the associated trajectory $\tau$. Furthermore, $K_{1} > 0$ and $K_{2}$ are constants with respect to the change in $\phi_{\theta}(s\rightarrow s')$ and $\mathcal{E}(s_{T})$.

Our result hints at how the potential function is likely to be positively correlated with the terminal energies $\mathcal{E}(s_{T})$ reachable after the transition $s\rightarrow s'$. This implies how the potential function provides meaningful local credit, as maximizing the potential function likely leads to the maximization of the terminal energy.

---

W2. There are no experiments on problems with long trajectories.

To alleviate your concern, we conduct additional experiments with relatively long trajectories: the set generation with the size of $80$. We incorporate this experiment as an additional ablation study in Appendix E.1 in our updated manuscript. Again, one can observe how LED-GFN significantly improves the GFN.

---

W3. The experiments do not consider tasks requiring expensive energy functions, so the significant computational advantage of LED-GFN is unclear.

To alleviate your concern, we provide additional experiments in Appendix E.2 of the updated manuscript that demonstrates a significant computational advantage of LED-GFN even when using moderately expensive energy functions (oracle used in molecular generation).

We compare (1) the detailed time costs and (2) the performance with respect to the time costs with the considered baselines, i.e., GFN and FL-GFN, in molecular generation. Our method (1) requires approximately $14\\%$ of the additional time required for FL-GFN to enable partial inference in GFlowNets, and (2) shows the time-efficient performances compared to the considered baselines.

---

W4. Why not consider a task to generate continuous data?

We focus on discrete objects due to their importance in applications such as drug discovery and protein design. Nevertheless, our framework can be naturally extended to continuous GFlowNets based on detailed balance and sub-trajectory balance objectives. This would be an interesting direction for future research.

---

Q1. Can you comment on the (1) scalability and (2) how much of the benefits LED-GFN provides for longer trajectories?

First, our LED-GFN introduces a small overhead compared to the GFlowNet training procedure, as it requires the evaluation of the potential function for each trajectory. In Appendix E.2, we provide the detailed time costs in molecular generation tasks. The conventional subTB objective takes $5.80$ seconds to obtain $100$ samples during training while our LED-GFN takes $6.61$ seconds.

Next, longer trajectories require deciding a large number of actions, which is more likely to make the local credit assignment more challenging for GFlowNets. The benefits of our LED-GFN are more significant in these scenarios by providing informative local credits for partial inference in GFlowNets, as supported by our new experiment in Appendix E.1.

---

Q2. Have you considered active learning with expensive energy evaluation and there is no intermediate energy signal?

Indeed, we did not conduct experiments with expensive energy evaluations, e.g., AutoDock. Nevertheless, we agree that the suggested experiments are interesting to highlight the benefits of our method: the efficiency with respect to the number of samples or energy evaluation. Unfortunately, we are unable to conduct such experiments during the rebuttal phase due to the expensive energy evaluation itself. We will consider this as a future work for evaluating our framework.

---

Reference

[1] Arumugam et al., An Information-Theoretic Perspective on Credit Assignment in Reinforcement Learning, NeurIPS 2020 Workshop on Biological and Artificial Reinforcement Learning

Dear reviewer xgxq,

Thank you for the response! We appreciate your comments to improve our paper.

We would like to clarify the following points to address your additional comments.

---

W1-1. The empirical performance is great, but the new theoretical analysis (Appendix F) is not clear, e.g., can not guarantee convergence.

We agree that our new analysis lacks a clear explanation due to the challenge of guaranteeing optimal conditions in all settings. However, as stated in response of W1, it is worth noting that our main focus is on improving empirical performance. Therefore, instead of incorporating the current analysis in the manuscript, we leave in-depth analysis as an interesting future work (stated in the last sentence of Section 5 of the updated manuscript).

---

W2-1/Q-1. In the new experiments for large set generation, LED-GFN shows similar performance compared to the FL-GFN, so the claim of LED-GFN seems to be weakened.

We would like to clarify that a large set generation is idealized setting for FL-GFN as a synthetic task, where the energy function is designed to perfectly identify the contribution of each action, i.e., ideal local credits (setting of the set generation discussed in Section 4.4).

In this task, achieving performance similar to FL-GFN still supports our claim by highlighting that the learned potentials can be as informative as ideal local credits (as extended results of Figure 7(a) for longer trajectories). We incorporate this clarification in Appendix E.1 of the updated manuscript.

---

W3-1. The advantage of time costs is promising, but experiments still do not consider expensive energy evaluation.

We are happy to hear that our new experiments address your concerns regarding the advantage of time costs!

For expensive energy evaluation, e.g., docking simulation, we do not consider such a setting which is unable to be conducted during the rebuttal phase due to the expensive costs itself, but we believe that this can be an interesting future direction (as stated in response of Q2).