Paper: Learning Energy Decompositions for Partial Inference in GFlowNets (ICLR-2024)

Reviewer comment (Theory):
The paper considers GFlowNets whose reward function \( R(s_{n}) \) corresponds to a potential function that is additive over transitions, i.e., \( - \log R(s_{n}) = \mathcal{E}(s_{n}) = \sum_{t=0}^{n-1} \phi_{\theta} (s_{t} \rightarrow s_{t+1}) \). This makes sense for set GFlowNets, where \( s_{n} \) contains information about all the transitions but not about their order, as \( s_{n} = x \) is the set of elements that have been added at each transition. Is it also applicable for the tasks where order might matter?

Author response:
**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.

Reviewer comment (Theory):
In terms of Figure 2, why do we want to have (b) for potentials leading to approximately the same energies to minimize variance (more smooth transitions)? In this case, should we just simply set \( \frac{1}{3} \) if \( \mathcal{E} = 1 \)?

Author response:
**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:
\begin{equation*}
    \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 \\}.
\end{equation*}
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.

Reviewer comment (Theory):
In terms of Figure 2, when \( e^{- \mathcal{E}(s)} = e^{- \mathcal{E}(s^{\prime})} \Rightarrow \mathcal{E}(s \rightarrow s^{\prime}) = 0 \), it reduces to the DB constraint, such that we cannot take advantage of intermediate rewards or energies to learn from incomplete trajectories. Is this a common case in many tasks?

Author response:
**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.

Reviewer comment (Evaluation):
In terms of Figure 9(b), how should "number of calls" be understood? Assuming we have 16 trajectories now, LED-GFlowNets compute 16 terminal rewards, while FL-GFlowNets need to compute all intermediate and terminal rewards (should be \( \sum_{i=1}^{16} n_{i} \), where \( n_{i} \) is the number of states, except \( s_{0} \), in the \( i \)-th trajectory). Thus, does this correspond to 16 calls vs. \( \sum_{i=1}^{16} n_{i} \) calls? However, shouldn't we have more than 16 calls for LED-GFlowNets, since we need to train the potential function?

Author response:
**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**).

Reviewer comment (Experiments):
For experimental details, the authors state "We set the dropout probability as 10% for tasks with a trajectory length less than 10 and 20% for others." Did they try other proportions? How does \( \lambda \) affect performance? Ablation studies on \( \lambda \) for Figure 9(a) are missing.

Author response:
**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.

Reviewer comment (Writing):
\( \log \) is missing for \( P_B \) regarding the DB loss, as well as in the Appendix.

Author response:
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.

Reviewer comment (Novelty):
As far as I understand, it should be: SubTB(\( \lambda \)) [Madan et al.] is a practically useful interpolation between TB and DB losses.

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

- We clarify that we actually consider the subTB($\lambda$).

Reviewer comment (Theory):
Below equation (4), "... by replacing \( \mathcal{E}(s) \), \( \mathcal{E}(s^{\prime}) \) with" should be \( \sum_{t=0}^{u} \phi_{\theta}(s_{t} \rightarrow s_{t+1}) \) and \( \sum_{t=0}^{u+1} \phi_{\theta}(s_{t} \rightarrow s_{t+1}) \), respectively.

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

- 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.

Reviewer comment (Writing):
Figure 4: "SubTB-based objectives" should be "subTB-based objectives".

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

- We fix 'SubTB' to 'subTB' in Figure 4.

Reviewer comment (Presentation):
For Figure 4, it is clear that DB or subTB is considered. But for Figures 5 & 6, it's just LED-GFN or FL-GFN. Which training objective is used? Although "For FL-GFN and LED-GFN, we consider a subTB-based implementation" is mentioned for molecule generation, maybe just use LED-subTB for the figures to avoid such confusion.

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

- We fix 'FL-GFN' and 'LED-GFN' in Figure (5) and (6) to 'FL-subTB' and 'LED-subTB', respectively.

Reviewer comment (Writing):
I take issue with the use of "partial inference of GFlowNets" in the title and "partial inference" elsewhere and would strongly suggest for this to be revised. "Inference of X" means, given some information related to X, producing an instance of X or a distribution over X. So "partial inference of GFlowNets" should mean learning part of a GFlowNet, or learning it from incomplete information. This is not what is done in this paper: the whole GFlowNet is learned, but some terms in the objective can be computed using partial trajectories. The third paragraph of the introduction and the second paragraph of 2.2 attribute "partial inference" to [Pan et al., 2023a], but in fact that paper does not introduce this term and does not ever use it. The minimal fix would be to make the title "partial inference in/for GFlowNets", which would more accurately describe what is being done.

Author response:
**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".

Reviewer comment (Writing):
Small errors in exposition on GFlowNets: Error on the top of p.3: $\exp(-{\cal E}(x))$ is the reward, not the energy. End of 2.1: SubTB as written does not "interpolate between DB and TB". It only interpolates if one uses the $\lambda$ parameter from [Madan et al., 2023], in which case it indeed interpolates between DB ($\lambda\to0+$) and TB ($\lambda\to+\infty$). In 2.2: As written, (2) simply does not make sense: $\cal E$ is defined to be a function with domain $\cal X$, but then it is applied to nonterminal states $s$! This can be fixed by writing that FL **assumes** an extension of $\cal E$ to nonterminal states, and that the freedom we have in the choice of this extension is a starting point for this paper (plus briefly discuss possible sources of the partial energies).

Author response:
**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.

Reviewer comment (Evaluation):
Limitations, costs, and failure modes of the proposed algorithm are not discussed.

Author response:
**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.

Reviewer comment (Reproducibility):
Missing details on the form of the model that predicts $\phi$ (not found in the appendix). It is unclear whether it shares some parameters with the flow and policy models. It is an interesting question how simple or complex the energy decomposition model needs to be, relative to the policy model, in order to be useful.

Author response:
**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.

Reviewer comment (Theory):
Some theoretical characterization of the optimal energy decomposition would be helpful (even in the case of a tree-shaped state space).

Author response:
**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.

Reviewer comment (Presentation):
In all plots with curves, please use markers or line styles, and not just colour, to distinguish the curves.

Author response:
**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.

Reviewer comment (Theory):
Learning the decomposition of the terminal potentials is interesting, but there are no theoretical guarantees that the learned decomposition provides meaningful local credits in all settings.

Author response:
**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 
\begin{equation*}
\phi_{\theta}(s\rightarrow s') = K_{1}\sum_{\tau \ni (s\rightarrow s')}\mathcal{E}(s_{T}) + K_{2},
\end{equation*}
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.

Reviewer comment (Experiments):
In terms of the empirical results, while the method performs quite well on a variety of tasks, there is an important caveat to note regarding the size of the problems. The trajectory length in the problems considered is quite small. There are no experiments on problems with long trajectories, which is where the local credit assignment would be critical and thus demonstrate the efficacy of the approach.

Author response:
**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.

Reviewer comment (Experiments):
Another motivation for the approach even in the presence of ideal intermediate signals is that the true intermediate energy function can be expensive to compute. However, all the experiments consider tasks where this is not the case, so it is unclear whether there is a significant computational advantage.

Author response:
**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.

Reviewer comment (Experiments):
Another important caveat of the empirical analysis is that it focuses on discrete problems and does not consider the continuous case.

Author response:
**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.

Reviewer comment (Experiments):
Can you comment on the scalability of the approach and how much of the benefit it provides for longer trajectories?

Author response:
**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**.

Reviewer comment (Experiments):
Have you considered combining the LED-GFN approach in an active learning setting where the terminal reward is expensive to compute and there is no intermediate energy signal? (This seems to be the actual useful case in the molecule generation case)

Author response:
**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.