Learning Energy Decompositions for Partial Inference in GFlowNets
Authors: Hyosoon Jang
Minsu Kim
Sungsoo Ahn
Venue: ICLR 2024 oral (2024)
Keywords: Generative flow networks
reinforcement learning
generative models
Primary area: generative models
TL;DR: We investigate a learning-based approach to produce informative local credits that facilitate the partial inference of GFlowNet

Abstract: This paper studies generative flow networks (GFlowNets) to sample objects from the Boltzmann energy distribution via a sequence of actions. In particular, we focus on improving GFlowNet with partial inference: training flow functions with the evaluation of the intermediate states or transitions. To this end, the recently developed forward-looking GFlowNet reparameterizes the flow functions based on evaluating the energy of intermediate states. However, such an evaluation of intermediate energies may (i) be too expensive or impossible to evaluate and (ii) even provide misleading training signals under large energy fluctuations along the sequence of actions. To resolve this issue, we propose learning energy decompositions for GFlowNets (LED-GFN). Our main idea is to (i) decompose the energy of an object into learnable potential functions defined on state transitions and (ii) reparameterize the flow functions using the potential functions. In particular, to produce informative local credits, we propose to regularize the potential to change smoothly over the sequence of actions. It is also noteworthy that training GFlowNet with our learned potential can preserve the optimal policy. We empirically verify the superiority of LED-GFN in five problems including the generation of unstructured and maximum independent sets, molecular graphs, and RNA sequences.

Published as a conference paper at ICLR 2024
LEARNING ENERGY DECOMPOSITIONS FOR PARTIAL INFERENCE IN GFLOWNETS
Hyosoon Jang', Minsu Kim², Sungsoo Ahn
POSTECH KAIST
{hsjang1205, sungsoo. ahn ac. min-suckaist
ABSTRACT
This paper studies generative flow networks (GFlowNets) to sample objects from the Boltzmann energy distribution via sequence of actions. In particular, we focus on improving GFlowNet with partial inference: training flow functions with the evaluation of the intermediate states or transitions. To this end, the recently developed forward-lookie looking GFlowNet reparameterizes the flow functions based on evaluating the energy of intermediate states. However, such an evaluation of intermediate energies may (i) be too expensive impossible to evaluate and (ii) even provide misleading training signals under large energy fluctuations along the sequence of actions. To resolve this issue, propose learning energy decompo- sitions for GFlowNets (LED-GFN) Our main idea is to (i) decompose the energy of object into learnable potential functions defined on state transitions and (ii) repa- rameterize the flow functions using the potential functions. In particular, to produce informative local credits, we propose regularize the potential change smoothly over the sequence of actions. It is also noteworthy that training GFlowNet with our learned potential can preserve the optimal policy, We empirically verify the superiority of LED-GFN in five problems including the generation of unstructured and maximum independent sets, molecular graphs, and RNA sequences.
INTRODUCTION
Generative Flow Networks (Bengio et al., 2021a, GFlowNets or GFNs) are frameworks to sample objects through sequence of actions, e.g., iteratively adding nodes to a graph. Their key concept is training policy that sequentially selects the actions to sample the object from the Boltzmann distribution (Boltzmann, 1868). Such concepts enable discovering diverse samples with low energies, ie., high scores, as an alternative to reinforcement learning (RL))based methods which tend to maximize the return the sampled object (Silver et al., 2016; Sutton & Barto, 2018).
To sample from the Boltzmann distribution, GFlowNet trains the policy to assign action selection probability based on energy of terminal state (Bengio et al., 2021 Malkin et al., 2022), e.g., high probability the action responsible for the low terminal energy. However, such training has fundamental limitations in credit assignment, as it is hard to identify the action responsible for terminal energy (Pan et al., 2023a). This limitation stems from solely relying on the terminal energy associated with multiple actions, lacking the information to identify the contribution of individual actions, akin to challenges in RL with sparse reward (Arjona-Medn Medina al., 2019; Ren al., 2022).
An attractive paradigm to tackle this issue is partial inference (Pan et al., 2023a) that trains flow functions with local credits, e.g.. evaluation of the intermediate states or transitions. Such local credit identifies individual action contributions to the terminal energy before reaching the terminal state. To this end, Pan et al. (2023a) proposed a forward-looking GFlowNet (FL-GFNN), which assigns the local credit based the energy of incomplete objects associated with intermediate states.
However, FL-GFN crucially relies on two assumptions that may not hold in practice. First, FL-GFN requires evaluating the energy of the intermediate state in the trajectories. However, the energy function can be expensive or even impossible to evaluate. Next, FL-GFN assumes the energy of intermediate states to provide useful hints for the terminal energy. However, this may not be true when the intermediate energy largely fluctuates along the sequence of states, e.g., low intermediate energies may lead to terminal state with high energy. We illustrate such a pitfall in Figure 1.

---

Published as a conference paper at ICLR 2024
action
bag state
X
high score
cumu alive credits
LED-GFN
FL-GFN
time step
Figure I: The local credit evaluation in bag generation (Example 1). (first row) The task is to generate bag of entities, where the seven same entities yield high score. (second row) Lefi-to-right indicates the state transitions over given trajectory. (third row) The energy-based evaluation fails to produce informative local credits since every intermediate state has zero energy, whereas our potential function produces informative credits by enforeing the potentials to be uniformly distributed.
Contribution. We propose learning energy decomposition for GFlowNet,, coined LED-GEN. Our key idea is to perform partial inference by decomposing terminal state energy into sum of learnable potentials associated with state transitions and use them as local credits. In particular, we show how to regularize the potential function to preserve the ground-truth terminal energy and minimize variance over the trajectory to yield informative potentials. Figure I highlights how LED-GFN provides informative local credits compared the existing approach.
To be specific, our energy decomposition framework resembles the least square-based return decom- position for episodic reinforcement learning (Efroni et al., 2021; Ren et al.. 2022). We parameterize potential function with a regression model that is constrained to be equal to the terminal energy when aggregated over the entire action sequence. We also regularize the potential function to minimize the variance along the trajectory, so that the potential function provides dense local credits in training GelowNet. Such potentials associated with intermediate transitions provide informative signals, as each of them is enforced to be correlated with the terminal energies. The training of potential function is online, which uses samples collected during GFlowNet training.
We extensively validate LED-GFN on various tasks: set generation (Pan et al., 2023a), bag generation (Shen et al., 2023), molecular discovery (Bengio et al., 2021b), RNA sequence generation (Jain et al., 2022), and the maximum independent set problem (Zhang et al., 2023). We observe that LED-GEN (1) outperforms FL-GGN when the assumption of intermediate energy does not hold, (2) excels in practical domains compared to GelowNets and RL-based baselines, and (3) achieves similar performance to FL-GFN even when intermediate energy provides the "ideal" local credit.
2
PRELIMINARIES
In this section, we describe generative flow networks (Bengio et al., 2021a, GFlowNets or GFNs) and their partial inference algorithm. We describe additional related works in Appendix A.
2.1 NETS
GFlowNets sample from discrete space X' through a sequence of actions from the action space A that make transitions in the state space S. For each complete trajectory r = (so, $1,.... sT), the terminal state is the object æ = ST E X' to be generated. The state transitions are determined by the action sequence (a,,... 1), 1), e.g., a, determines - 8141- The policy Pr(s'|s) selects the action a to
transition from the current state 8 to the next state s' and induces a distribution over the object x.
The main objective of GelowNet is to train the policy Pr(-).) that samples objects from the Boltzmann distribution with respect to given energy function E : X + as follows:
(1) where PJ (x: is the distribution of sampling object induced from marginalizing over the trajectories conditioned on I = sy. We omit the temperature for simplicity. To this end, GFlowNet trains with auxiliary objectives based on state transition, trajectory, sub-trajectory information.
(x) x exp(-E(x:),
2

---

Published as a conference paper at ICLR 2024
Detailed balance (Bengio et al., 2021b, DB). The DB utilizes the experience of state transitions to train GFlowNet. It trains the GelowNet with a forward policy model (s a backward policy (ss(ss'), and state flow estimator F(-) S It by minimizing the following loss function:
(DB(s, s') = (log (s) + log (s log F(5') log/ log?p(s(s'))"
where the flow Fs) for the terminal state $T is defined to be identical to the exponent of the negative energy exp (x)), i.e., the score of the object.
Trajectory balance (Malkin et al., 2022, TB). The TB aims to learn the policy faster by training on full trajectories. To this end, TB requires a forward policy model ss) a backward policy (s]s'')) and learnable scalar to minimize the following loss function:
T-1
log Pr($441(8)) - E(x) Σ log Pa(silees))
T-1
2
LTB = log Z +
Σ t=0
This objective is resilient to the bias from inaccurate flow estimator F(.) used in DB, since it directly propagates the terminal energy train on intermediate states. However, the TB suffers from the high variance of the objective over the collected trajectories (Malkin et al., 2022).
Sub-trajectory balance (Madan et al., 2023, subTB). The subTB trains forward and backward policies (s'(ss, (B(s'(s), flow function F-)), and a learnable scalar Z. It trains on flexible length of sub-trajectory su + SU+1 S to minimize the following loss function:
U+L-1
U+L-1
2
S -oog log F(s))
+
Σ t=
log Pf($c41l8)) log F(SU+L)
Σ
log Pa(silsett)
In addition, one can consider the weighted average of LruhTn over all possible lengths I, = 1.....,?? of the sub-trajectories with the weight X4, which enables the interpolation between the DB and the TB (Madan et al., 2023). In this paper, we consider this approach as subTB for simplicity.
2.2 PARTIAL INFERENCE IN GELOWNETS
The GFlowNet training is often challenged by limitations in credit assignment, i.e., identification and promotion of the action responsible for the observed low energy. This limitation stems from relying solely on the terminal state energy as the training signal. The terminal energy lacks information to identify the contribution of individual action, akin how reinforcement learning with sparse reward suffers from credit assignment (Arjona-Meding et al., 2019; Ren et al., 2022).
Partial inference a promising paradigm to resolve this issue by learning from local credits, where Pan et al. (2023a) first incorporate this concept in training GFlowNets. Specifically, the partial inference aims to evaluate individual transitions sub-trajectories, i.e., local credits, and provide informative training signals for identifying the specific contributions of actions. To this end, Pan et al. (2023a) proposed forward-Lookinn GFlowNet (FL.GGNN), which evaluates intermediate state energy as local credit signal for partial inference.
Forward- Looking GFlowNet (Pan et al., 2023a, FL-GFN), To enable partial inference, the FL-GFN defines new training objective that incorporates energy function & : X + R for intermediate states. To this specific, FL-GFN modifies the DB as follows:
CF)(s, = (log F(s) + log (s - Ess) + E(s') log F(s') log PB(s(s)))?
(2)
where F(s) exp (E(s)) is the re- re-parameterized flow function and E(s') {(s) is the energy gain associated with the transition from s to s'. Note that the energy function defined only on the terminal state space X: hence, FL-GFN assumes extension of the energy function & to non terminal states, e.g., evaluating energy of an object associated with the intermediate state s. Pan et al. (2023a) shows that optimum of Equation (2) induces a policy Pe (p(-l-) that samples from Boltzmann distribution. While Equation (2) is associated with DB, FL-GFN is also applicable to subTB with energy gains in the level of sub-trajectory.
3

---

Published as a conference paper at ICLR 2024
[e]
potential (energy decomposition)
GFlowNet (minimizing 4am))
:object energy
potentials
Illustration of energy decomposition for partial inference in GFlowNet. LED-GEN enables partial inference with potentials which (a) approximate the object energy via summation, and (b) minimize variance along the action sequence.
Figure 2:
LEARNING ENERGY DECOMPOSITION FOR GFLOWNET
While the FL-GFN is equipped with partial inference capabilities, it relies the energy function to assign local credits, which can be expensive to evaluate, or lead to sub-optimal training signals (details are described in Section 3.1). In this paper. propose learning energy decomposition for GFlowNets (LED-GFN) to achieve better partial inference. In what follows, we describe our motivation for better partial inference (Section 3.1) and the newly proposed LED-GEN (Section 3.2).
3.1 MOTIVATION FOR BETTER PARTIAL INFERENCE
Our motivation stems from the limitations of FL-GGNN which performs partial inference based on evaluating the energies of intermediate states with respect to single transition (for DB) or sub-traje trajectory (for subTB). In particular, the energy gain may yield a sub-optimal local credit signal due to the following pitfalls (see Figures and 9(b)).
Negative error
inspired by how
0.6
LED-GFN
FL-GFN
GFN
A. The energy evaluation for an incomplete object, i.e., interme- diate state, can be non-trivial. In addition, the of energy evaluation can be expensive, which can bottleneck efficient training when called for all visited states.
165
10
10>
Number of samples (x2.5)
Figure Negative relative mean еπто (1) for estimating the true Boltzmann distribution on Exam- ple 1-rype task (Bag).
B. The energy exhibit sparse high variance interme- diate states within trajectory, even returning zero for most states, which is -nnioorrrtttve for partial inference.
We further provide a concrete example for the pitfall B. of FL-GFN, which is illustrated in Figure and Figure 3' for both conceptual and empirical purposes, respectively:
Example 1. Consider adding objects from {A, B, C, D, F) to a bag with maximum capacity of nine. Define the energy as -1 when the bag contains seven identical objects and O otherwise.
For Example 1, the intermediate energy (which is always 0) does not provide information for the terminal energy. However, the number of the most frequent elements in the bag is informative even intermediate states since greedily increasing the number leads to the best terminal state.
Our observation in Example hints at the existence of partial inference algorithm to provide better local credit signals. We aim to pursue this direction with a learning beamningbbssed approach. That is, we parameterize the class of potential functions that decompose the terminal energy to provide local credit signals for GFlowNet training. Our key research direction is to understand about what conditions of the potential functions are informative for partial inference.
3.2 ALGORITHM DESCRIPTION
In this section, we describe our framework, coined learning energy decomposition for GFlowNet (LED-GFN), which facilitates partial inference using learned local credit. To this end, propose to decompose the terminal energy into learnable potentials defined on state transitions. Similar to FL-GFN, we reparameterize the flow model with local credits, i.e., potentials. In contrast to FL-GFN,
'Tee detailed experiment settings described in Section 4.1.
4

---

Published as a conference paper at ICLR 2024
Algorithm Learning energy decomposition for GFlowNet
1: Initialize the buffer B, forward and backward policy Pr, PB. state flow F, and model do.
2: Update the model do to minimize fED if the generation trajectories given in advance.
3: repeat
4: Sample a batch of trajectories (no)mm from forward policy Pp.
5: Update buffer B <- B U (T))mr
6: for 7 = ,,,,.,,N. .NV do
Energy decomposition learning
7: Sample batch of trajectories
) {n.)mm from the buffer
8:
Update the model deg to minimize LL with (T) (T))mm)-
9:
end for
10: Compute intermediate energy gains (0 0;) for all (5;, a4) ET.
11: Update the GFlowNet Pr, Pn, F to minimize LLED with T and do a:) for all (st,
12:
until converged
we optimize the local credits to enhance partial inference by minimizing the variance of potentials along the action sequence. See Figure 2 for illustration of LED-GFN.
Training with potential functions. To be specific, we decompose the energy function associated with the terminal state into learnable potential functions de follows:
T-1
E(z) (e((T)
Σ do(ss stt1),
where т = (so, $1..... ST), x = sT, and the potential functions are defined on state transition
+ St+1. Similar FL-GFN, we use the potential function train the forward and backward GFlowNet policies PF, PB and flow model F to minimize the following loss:
(LE = (log F(s) + log (=(s'(s) (ss + ss + s') log F(s') log PB(s(s'))
(4)
Given a sub-trajectory (so...., Su = 8, Sa+1 s'), one can derive this objective from Equation (2) by replacing {s), with Et= Φo(s; S(+1) and Σ ( - St+1), respectively. That is, it is evident that Equation (4) preserves the optimal policy of GFlowNet when E(x) = (T) is satisfied for all trajectories terminating with
Our objective becomes equivalent to that of FL-GFN when de(s -' E(s') Ess. but our key idea is to learn the potential function de instead of the energy gain which can be expensive and may exhibit sparsity or high variance, as pointed out in Section 3.1. Note that one can also introduce an approximation error E(x) (o(T) as an additional correction term preserve the optimal policy of GFlow even when the potential function de is inaccurate. In Appendix B. we describe how LED-GFN consistently induces the optimal policy that samples from the Boltzmann distribution.
Training potentials with squared loss. In training the potential function, our key motivation is twofold: (a) accurately estimating the true energy through summation and (b) providing dense and informative training signals by minimizing variance along the action sequence.
To this end, given a trajectory T (Su...., ., ST = 2)), we train the potential functions de minimize the loss function for (a) achieving {zx) ~ (b((T) with (b) dropout-basee regularization:
( -
27
-
(T) Bew-Ber(>))
-
Σ
2,00 + s1+1)
1=0
The T-length random vectors z promotes the dropout, where = 0 sampled from the Bernoulli distribution with probability y-y. Dividing by T and C = TT 2, aligns the scales compensate for the scale reduction induced by dropout. When 7 = 1 and the loss function is minimized, i.c., (T) 0 for all r, the potential function decomposes the energy function without error. When y < 1, dropout prevents heavy reliance on specific potentials to satisfy Equation (3), thereby reducing the variance and sparsity of the potentials along the action sequence." Note that our intuition is
RRen et al. (2022) provide formal proof that Equation (5) serves surrogate objective satisfy Equation (3) while reducing the variance and sparsity of the potentials along the action sequence.
5

---

Published as a conference paper at ICLR 2024
3.0
3.0
Aug top- top-1100 (x0109
2.5
2.5
7.0
2.0
7.0
1.5
1.3
3.5
1.5
1.0
Number
1.0
LEC-55
LeB
1LD8
0.5
0.5
15"
10
101
201
10t
10
10:
16f
15
Number samples (X2.5)
Number of samples (x2.5)
Number samples (X2.5)
Number of samples (x2.5)
(a) DB-hased objectives
(h) subTB-based objectives
Figure 4: The performance on bag generation. The solid line and shaded region represent the mean and standard deviation, respectively. The LED-GFN shows superiority on both DB and subTB.
similar to recent works on learning return decomposition to alleviate sparse reward problems in reinforcement learning (Arjona-Medina et al., 2019; Gangwani al., 2020; Ren al., 2022).
To train the potential function, we define its training as online learning within GFlow Net training, i... learning from trajectories obtained during GFlowNet training. We describe the overall algorithm in Algorithm 1, Such alternative training of the potential function and the policy is similar to model-based reinforcement learning algorithms (Luo et al., 2018; Sun et al., 2018; Janner et al., 2019) for monotonic improvement of policies. Note that LED-GFN can also be implemented with subTB, where the details are described in Appendix B.2.
4 EXPERIMENT
We extensively evaluate LED-GEN on various domains, including bag generation (Shen et al., 2023), molecule generation (Bengio al., 2021a), RNA sequence generation (Jain et al., 2022), set generation (Pan al., 2023a), and the maximum independent set problem (Zhang et al., 2023). Following prior studies, we consider the number modes, i.c., samples with energy lower than specific threshold, and the average top-100 score the base metrics, which are measured via samples collected during training. We report all the performance using three different random seeds.
4.1
BAG GENERATION
First, we consider a bag generation task (Shen et al., 2023). The action is adding an object from seven distinct entities bag with maximum capacity of 15. The bag exhibits low energy when including seven or more repeated entities of the same type. In this task, we compare our method with GFN and FL-GFN. We consider both DB-based (Bengio et al., 2021b) and subTB-based (Malkin et al., 2022) implementations. The detailed experimental settings are described in Appendix C.
Results. We present the results Figure 4. Here, one can observe that our method excels bag generation compared to GFN and FL-GFN on both DB and subTB. In particular, FL-GFN fails to make improvements on the subTB- based implementation, since most states do not provide informative signals for partial inference (as illustrated in Figure 1). In contrast, LED-GF consistently improves performance by producing informative potentials enhance partial inference.
4.2 MOLECULE GENERATION
Next, we validate LED-GFN in a more practical domain: the molecule generation task (Bengio et al., 2021a). This task aims to find molecules with low binding energy to the soluble epoxide hydrolase protein. this task, molecule generated by constructing junction trees (Jin al., 2018), with the actions of adding molecular building blocks. The binding energy between the molecule and the target protein computed using a pre-trained oracle (Bengio et al., 202
In this experiment, consider PPO (Schulman et al., 2017), and three GFN models: DB, TB, and subTB (Madan et al., 2023) as the baselines. Additionally, we compare our approach with GAFN (Pan al., 2023b) and FL-GFN. For FL-GFN and LED-GFN. consider subTB- implementation. The overall implementations and experimental settings follow prior studies (Bengio et al., 202 la; Pan al., 2023a), which described in Appendix C.

---

Published as a conference paper at ICLR 2024
Number of modes (X10
LED-subTB
7.0-
ANA 07 1
7.5
FL-subTB
7.0
subTB
GAFN
1.0
0.0
1,0
1.5 Number samples (x<105)
Number samples (x11%)
Figure 5: The performance on molecule generation. The solid line and shaded region represent the mean and standard deviation, respectively. The LED-GE GFN shows superiority compared the considered baselines in generating diverse high reward molecules.
3.0
1x10-1
LED-subTB
2.5
9.6
FL-subTB
Number of modes
Avg of topp 100
9.4
SubTB
1.5
9.2
GAFN
0.5
9.0
10
10a
10
10"
10
Number samples (x2.5)
Number samples (x2.5)
Figure 6: The performance on RNA sequence generation. The solid line and shaded region represent the mean and standard deviation, respectively. The LED-GFN shows superiority compared to the considered baselines in generating diverse high reward RNA sequences.
One can see that LED-GFN outperforms the considered baselines in enhancing the average score of unique top-100 molecules and the number of modes found during training. These results highlight that LED-GFN also beneficial for real- world generation problems with a large state space.
Results. The results presented in Figure 5.
4.3 RNA SEQUENCE GENERATION
We also consider RNA sequence generation task for discovering diverse and promising RNA sequences that bind to human transcription factors (Barrera et al., 2016; Trabucco et al., 2022; Jain et al., 2022). The action is appending prepending energy is pre-computed based on wet-lab measured DNA binding activity to Sine Oculis Homeobox Homolog 6 (Barrera et al., 2016). We consider the same baselines in the molecule generation task. For FL-GFN and LED-GEN, we consider SubTB-based implementation.
amino acid the current sequence. The
The results presented in Figure 6. One observe that LED-GEN outperforms the considered baselines. Furthermore, FL-GFN only makes minor differences compared to GFN, while LED-GFN makes noticeable improvements. These results highlight that energy-based partial inference can fail to improve performance in practical domains, while the potential learning- approach consistently leads to improvements.
Results.
4.4 COMPARISON WITH IDEAL LOCAL CREDITS
In these experiments, we demonstrate that LED-GEN can achieve similar performance compared to FL-GFN, even when the intermediate state energy is sufficient to identify the contribution of the action, i.e., ideal local credit (Zhang et al., 2023). Note that this tasks are idealized, since designing such energy function requires complete understanding of the domain. Especially, focus two tasks: set generation (Pan et al., 2023a) and the maximum independent set problem (Zhang et al., 2023). For these tasks, compare our method with GFN and FL-GFN.
Set generation. The set generation task is similar to the bag generation. The actions are adding an objects from 30 distinct objects to with maximum capacity of 20. The energy is evaluated

---

Published as a conference paper at ICLR 2024
185
20.0
11..0
Aug of topo1
2.0
Metric
37.0
Mode
10.0
166
2.0
LED
16.0
LED 08
1.5
PL-28
PL-3B
3.0
08
DB
0.0 02 2.4 1. 3.2 1.6 Number samples (x10')
05 02 0.4 9.6 0.5 14
04
2.0
4.0
6.2
8.3
0.2
2.0
4.0
Number of samples (x10
Step (x10*)
Step (x20
(a) Set generation
(b) Maximum independent set problem
Figure 7: The performance comparison with ideal local credits. The solid line and shaded region represent the mean and standard deviation, respectively. The LED-GFN shows similar performance to the FL-GFN, even when the intermediate energy is sufficient to identify the action contribution.
0.48
0.2
GRN
LED
Fl-subrB
Negative etror
subiis
0.4
1-Similarity
Negative
0.1
0.4
0.36
0.0
SUbTB
DB
0.0
subTB
DB
7.0
7.4
7.8
6.2
Method
Method
Avg of top-100
(a) Bag generation (L)
(b) RNA sequence generation (L)
(c) Molecule generation
Figure 8: (a-b) The negative relative mean error comparison (lower is better), Our LED-GFN better approximates the target distributions. (c) The Tanimoto similarity with respect to the average scores (upper-right is better). Our LED-GFN produces more diverse and promising molecules.
by accumulating the individual energy of each entity, so the intermediate energy gain has complete information to identify the contribution of each action (Pan et al., 2023a). We describe the detailed experiment settings in Appendix B.
Maximum independent set problem. This task aims find the maximum independent set by selecting nodes, and the energy is evaluated based on the current size of the independent set (Zhang al., 2023). Here, we compare the performance on validation graphs following Zhang al. (2023). At each step, we sample 50 solutions for each validation graph to measure the average scores and the number of mode founds (greater than 18.5). The overall implementations and hyper-pann parameers follow prior studies (Zhang et al., 2023).
Results. illustrated in Figure one can see that our approach achieves similar performance to FL-GFN, even though the intermediate state energy provides ideal local credit for partial inference. These results highlight that the potentials of LED-GFN can be as informative as ideal local credit, which provides complete identification of the action contributions.
4.5 ABLATION STUDIES
Goodness-of-fi to the true Boltzmann distribution. We show how well our algorithm finds good samples sampled from the target distribution, i.e.. Boltzmann distribution. Following Shen al. (2023), we measure the relative mean error between the target distribution and empirical samples obtained during training. Figures 8(a) and 8(b), one can observe that LED-GFN achieves better approximation target distribution compared to considered baselines.
Diversity vs. high score. Next, we further verify that our algorithm not only generates high-scoring samples but also diverse molecules. Specifically, score of the t0p-100 samples and the diversity these samples. To measure diversity, we compute the average pairwise Tanimoto similarity (Bajusz et al., 2015). In Figure 8(c), we illustrate the Tanimoto similarities with respect to the average of top-100 scores. Here, one can observe that LED-GFN achieves better diversity with respect to the average scores.
analyze the trade-off between the average
8

---

Published as a conference paper at ICLR 2024
Number of modes (K300)
9.0
8.0
8.0
Number of modees
7.00
of 1000100
2.5
Avg of 2op 100
2.3
6.5
LEDaubrB
y=
RUSUbTB
wb regubrization
6.0
SUbTS
00 0.2 04 85 08 1.0 Number of samples IK10")
0. 04 066 3.5.. 12112414
0.0
2.0 Number oradle calls (x10*)
03
T5
20
03
13
2X
Number of samples (x10')
Number oracle calls (x10%)
(a) Set generation
(b) Molecule generation
Figure 9: (a) The benefits of regularizing variance of potentials. The regularization improves the performance. (b) The performance over the number of energy evaluation. The FL-GFN can not
make improvements with respect to the number of energy evaluation.
eumb of modes (x10))
Number of mootes (x109)
Weg of t 100 (x10
2.5
10g attop-100
7.5
3.0
7.0
2.0
6.54
Simple Smoothing
Simple Smoothing
Number samples 1x10
samples Ex10')
Number samples
samples (x10F)
(a) Set generation
(b) Molecule generation
Figure 10: Comparison between energy decomposition methods for partial inference in GFlowNets. Our energy decomposition learning scheme shows most promising results.
also analyze how reducing the variance of potentials benefits the improvement in performance. In the set generation, we compare LED-DB with various dropout rates 1 - = 0.1, 0.2, 0.4 and its counterpart without regularization - = in Equation (5). In Figure 9a), can that applying regularization yields more promising performance, although too high dropout rate may result in inefficiency.
Regularized vs. non-regularized potentials.
Number of energy evaluation performance We analyze the performance with respect to the number of energy evaluation, which can be expensive. In Figure 9(b), one can see that FL-GFN can improve performance since it requires evaluating the energy for every visited states. In contrast, LED-GEN uses a potential function without energy evaluation for intermediate states.
Comparison with alternative energy decomposition methods. We further investigate the following alternative energy decomposition schemes to enable the partial inference in GFlowNets.
One may train simple model Φo X' + R to predict the terminal energy, and utilize it to compute potentials (s - s') = (s') (s). This approach can be interpreted as extension of proxy model-based GFlowNet (Jain et al., 2022) for partial inference.
Based on the LSTM-based decomposition method (Arjona-Medine Medina et al., 2019), one can de- sign the potential og - S++1) as the difference between two subsequent predictions for (5o, ........ St, (e) and (so, do. St41, at+1) using an LSTM.
Based on the return smoothing over the trajectory (Gangwani et al., 2020). one can simply set Φe(s; + 8(41) ( for given trajectory 7 at each step.
In Figure 10, we compare each method molecule and set generation tasks with a DB-based implementation. Here, performance due to its capabilities in producing dense and informative potentials.
see that the least square-based approach shows the most competitive
5 CONCLUSION
In this paper, we propose learning energy decomposition for GFlowNets LFND. Experiments various domains show that LED-GFN is promising compared to existing partial inference methods for GFlowNet. An interesting avenue for future work is developing new partial inference techniques using learnable local credit, theoretically analyzing learned potentials and optimal local credits.

---

Published as a conference paper at ICLR 2024
Reproducibility. We describe experimental details in Appendix C, which provides the base imple- mentation references, enviromments, and detailed also include the codes for molecule generation tasks based on the official implementation codes of the prior study (Pan et al., 2023a).
In the supplementary materials,
ACKNOWLEDGEMENTS
This work partly was supported by Institute of Information communications Technology Planning & Evaluation (ITP) grant funded by the Korea governmen((MSIT) (No. IITP-2019-0-01906, Ar- tificial Intelligence Graduate School Program(POSTEECH), the National Research Foundation of and Basic Science Research Program through the National Research Foundation of Korea(NRF) funded by the Ministry of Education (2022R) A6A 1 A030529511)))
Korea(NRF) grant funded by the Korea (OVe) (No.
REFERENCES
Jose A Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, Johannes Brand- stetter, and Sepp Hochreiter. Rudder: Return decomposition for delayed rewards. Advances in Neural Information Processing Systems, 32, 2019.
Dávid Bajusz, Anita Rácz, and Károly Héberger. fingerprint-based similarity calculations? Journal of cheminformaties, 7(1):1-13, 2015.
Why tanimoto index an appropriate choice for
Luis A Barrera, Anastasia Vedenko, Jesse V Kurland, Julia M Rogers, Stephen Gisselbrecht, Elizabeth Rossin, Jaie Woodard, Luca Mariani, Kian Hong Kock, Sachi Inukai, et al. Survey of variation in human transcription factors reveals prevalent dna binding changes. Science, 351 16280):1450-1454 2016.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381-27394 2021a.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. arXiv preprint arXiv:2111.09266 2021b.
Ludwig Boltzmann. Studien uber das gleichgewicht der lebenden kraft. Wissenschafiliche Abhand- lungen, 1:49-96, 1868.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2019. URL https: //openreview n wttp://wwww....et.e.eeeeeee-ee Hormm?id-HHuUnne
Yonathan Efroni, Nadav Merlis, and Shie Mannor. Reinforcement learning with trajectory feedback. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 7288-7295, 2021.
Tammay Gangwani, Yuan Zhou, and Jian Peng. Learning guidance rewards with trajectory-space smoothing. Advances in Neural Information Processing Systems, 33:822-832, 2020.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference machine learning, pp. 1272-1272. PMLR, 2017.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure FP
Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning, pp- 9786-9801. PMLR, 2022.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-basel policy optimization. Advances in neural information processing systems, 32, 2019.
10

---

Published as a conference paper at ICLR 2024
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning, pp. 2323-2332 PMLR. 2018.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorith- mic framework for model-based deep reinforcement learning with theoretical guarantees. In International Conference on Learning Representations. 2018.
Kanika Madan, Jarrid Rector-Brooks Maksym Korablyov, Emmanuel Bengio, Moksh Jain, An- drei Cristian Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning gflownets from partial episodes for improved convergence and stability, International Conference on Machine Learning, pp. 23467-23483. PMLR, 2023.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. Advances in Neural Information Processing Systems, 35: 5955-5967, 2022.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of gflowntts with local credit and incomplete trajectories. arXiv preprint arXiv:2302.01687, 2023a.
Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative augmented flow networks. International Conference on Learning Representations, 2023b. URL https: /pwpnrevieww...ee///ee .nt. /eret/foruue forum? iduurrduFFro _
Zhizhou Ren, Ruihan Guo, Yuan Zhou, and Jian Peng. Learning long-term reward redistribution via randomized return decomposition. In International Conference on Learning Representations, 2022. URL https: https://pennzeiiww.ccem wtt.://teee..ttm...oo...occoooee
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1777v:1707.06347, 2017.
Max Walt Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towards understanding and improving gflownet training. In Proceedings of the 40th International Conference Machine Learning. Proceedings of Machine Learning Research, PMLR, 2023.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Joannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
Wen Sun, Geoffrey J Gordon, Byron Boots, and J Bagnell. Dual policy iteration. Advances in Neural Information Processing Systems, 31, 2018.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bene bench: Benchmarks for data-driven offline model-based optimization. In International Conference on Machine Learning. pp.
PMLR, 2022.
Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron Courville, Yoshua Bengio, and Ling Pan. Let the flows tell: Solving graph combinatorial optimization problems with gflownets. arXiv preprint 205. 2023.

---

Published as a conference paper at ICLR 2024
A RELATED WORKS
Generative augmented flow network (Pan et al., 2023h, GAFN). The GAFN is learning frame- work that incorporates intermediate rewards for exploration purposes. Specifically, GAFN measures the novelty score of a given state, which provides an intrinsic signal to facilitate better exploration towards unvisited states. To compute the novelty score, this method also incorporates online training of random network distillation (Burda et al., 2019), which assigns lower scores to unseen states compared to more frequently observed states.
Proxy model-based GFlowNet. Jain et al. (2022) propose proxy model-based training of GFlowNet for discovering diverse and promising biological sequences. They train a proxy model of the energy function to mitigate the expensive cost of evaluating biological sequences, such as wet-lab evaluation. Additionally, they introduce an active learning algorithm for the model-based GFlowNet, leveraging the epistemic uncertainty estimation of the model to improve exploration.
Return decomposition learning. Our LED-GFN approach is inspired by return decomposition learning of reinforcement learning in sparse reward settings. Their goal is to decompose the return into step-wise dense reward signals (Arjona-Medina et al.. 2019; Gangwani et al., 2020; Ren et al., 2022). They have studied various return decomposition methods. First, Arjona-Medina et al. (2019) utilize an LSTM-bassd model to produce step-wise proxy rewards. Next, Gangwani et al. (2020) propose simple approach that uniformly redistributes the terminal reward over the trajectory. Ren et al. (2022) train a proxy reward function using randomized return decomposition learning which is contrained to produces dense and informative proxy rewards.
12

---

Published as a conference paper at ICLR 2024
B DETAILS OF LED-GFN
B.1 PRESERVING OPTIMAL POLICY OF GFLOWNET
Although the potential function inaccurate, we show that optimum of LLDD can induce an optimal policy that samples from Boltzmann distribution. We give simple proof by reduction to the optimum of DB objective (Bengio al., 2021b). Suppose the parameterization Φo - s') 0o(s') Φo(s).
and log (s) = ( + log F()). Then, we can reformulate LEDD as follows:
(S s') log (s) + log (p log ((')) -
logPs(s)s'))
we add an additional correction term [(x) (T( -E(x) 60 to the terminal flow such that log F(r) = -E(r), this objective becomes to be equivalent to the reparameterization of DB, where the optimum induces policy sampling from Boltzmann distribution (Bengio al., 2021b). Therefore, the optimum of LED-GFN can still induce the policy that samples from the Boltzmann distribution. We refer to this correction-based approach as LED-GFN*.
However, our implementation follows prior study in return decomposition learning (Ren et al., 2022), which uniformly redistributes the decomposition trajectory (we denote this approach as LED-GFN in experiments). In Figure 11, we empirically observe that this approach further improves the training of GFlowNets. We assume that uniformly redistributed decomposition error partially provides more dense and informative local credit signals correlated with future energy.
the transitions within the given
Number of modes ix10
3.0
2.55
2.5
2.0
7.5
LED-DB
LED-DB"
FL-DB
0.5
DB
0
Number of samples (.5)
Number of samples (x2.5)
Figure 11: The performance on bag generation.
B.2 TRAINING ON SUBTB
The LED-GEN can also be implemented on subTB by incorporating the potentials within sub- trajectories. To this specific, one modify subTB as follows:
U+L-1
U+L-1
LLED- (log F(su)+
Σ
log Pr(S
+
Σ
Φo(ss - 841)
U+L-1
log F(SU+L-1)
Σ
log P
1=
which based on sub-trajectory sU + SU+1 ... SU+L-1. This objective is equivalent to FL-GFN on the subTB when (s(s - -s E(s') E(s) (Pan et al., 2023a), but we replace it with the potentials for better credit assignment. Note that we actually consider the weighted average of all possible lengths T. 1 1,..., T of the sub-trajectories with the weight 2
LLED
Tb L
13

---

Published as a conference paper at ICLR 2024
C EXPERIMENTAL DETAILS
In all experiments, we design the potential functions to have the same architecture as the flow model, e.g., feedforward network with same hidden dimensions, where the input dimension is extended to consider the transition 8t + 8t+ In Equation (5), we set the dropout rate y as 10% for tasks with a trajectory length less than 10 and 20% for others.
Bag generation (Shen et al., 2023). The experiment settings, implementations, and thyer-parameee are based on prior studies (Shen et al.. 2023). The bag generation task is to generate a bag with maximum capacity of 15. There are seven types of entities, and each action includes one of them in the current bag. If contains seven or more repeats of any items, it has reward 10 with 75% chance, and 30 otherwise. The threshold for determining the mode is 30.
In each round, we generate B 32 bags from the policy. The GFlowNet model consists of two hidden layers with 16 hidden dimensions, which is trained with a learning rate of le-4. We use an exploration epsilon 0.01. this task, we incorporate buffer for energy decomposition learning. The mini-batch, size B2 is same B1. We set the number of iterations in energy decomposition learning = 8 for each round. Note that reducing still leads to promising results compared to baseline.
Molecule generation (Bengio et al., 2021a). The experiment settings, implementations, and hyper- parameters are based prior studies (Bengio et al., 2021a; Pan al., 2023a). The maximum trajectory length is 8, with the number of actions varying between around 100 and 2000 which is depending on the state. The threshold for determining the mode is 7.5.
In each round, we generate four molecules, i.e., B1 = 4. The model consists of Message Passing Neural Networks (Gilmer et al., 2017) with ten convolution steps and 256 hidden dimensions, which is trained with a learning rate of 5e-4. We rescale the reward so that the maximum reward is close one, and the exponent of set to 8.0. We use an exploration epsilon of 0.05. In energy decomposition learning, we do not use buffer and immediately use the molecules that are sampled in each round. For PPO, we set the entropy coefficient to 1c-4 and do not apply the reward exponent because it causes a gradient exploding.
RNA sequence generation (Shen al., 2023). The experiment settings, implementations, and lyper-parameters are based on prior studies (Shen et al., 2023). The action is defined as prepending or appending an amino acid to the current sequence. The maximum length is 8 and the number of actions is 4. The mode is determined based on whether it is included in predefined set of promising RNA sequences (Shen et al., 2023).
In each round, generate B1 16 sequences. The GFlowNet model consists of two hidden layers with 128 hidden dimensions, which is trained with a learning rate of 1e-4. The reward exponent set to 3.0. We an exploration epsilon of 0.01. In this task, we incorporate a buffer for energy decomposition learning. The wyper-parameterr for energy decomposition learning are the same as that of bag generation. For PPO, we the entropy coefficient to 1c-2.
Set generation (Pan et al., 2023a). The experiment settings, implementations, and hyper-paramente parameters based prior studies (Pan et al., 2023a). The set generation task is generate set that involves 20 entities. There are 30 types of entities, where each action includes one of them in the current set. We set the threshold for determining the mode as 0.25.
In each round, we generate B1 16 sets. The GFlowNet model consists of two hidden layers with 256 hidden dimensions, which is trained with learning rate of 0.001. The hyper-parameters for energy decomposition learning are the same as that of molecule generation.
14

---

Published as a conference paper at ICLR 2024
2.0
Number of modes (x104)
Aug of of (x105
8.0
Number of modns (x105)
1.0
8.0
Avg of tott-p-110 100
4.0
2.0
LED DB
3.0
5.5
LED-ubTT
88 ds u. Number samples (x1.0P)
0.83 04 0.5
0.0 0. 0.4 56
3.6
0.0 02 6.6 6.6
Number samples (x107)
Time costs (secx10*)
Time costs (secx10f)
(a) Set generation
(b) Molecule generation
Figure 12: (a) Large generation. The partial inference significantly improves the performance. (b) The performance with respect to the time costs. The LED-GFN shows most promising results.
Number of modes (x20")
8.0
100
7.5
2.0
TB
0.3
0.5
0.3
3.5
200
2.5
0.0
03
100
Number of samples (x10%)
Number samples (x102)
Figure 13: The performance on molecular generation. The energy decomposition learning also benefits the TB-based objective.
D ADDITIONAL ABLATION STUDIES
D.1 LARGE SET GENERATION
We conduct additional experiments with relatively long trajectories: the set generation with the size of 80 where the actions are adding an objects from 100 distinet objects (Pan et al., 2023a). This task more likely to make the local credit assignment more challenging with the terminal energy. Note that this setting is an idealized setting for FL-GFN, where intermediate energy gain has complete information identify the contribution of each action (setting discussed in Section 4.4). In Figure 12(a), one can observe that LED-GFN (I) significantly outperforms the GFN and (2) shows similar performance compared to FL-GFN with ideal local credits. This result highlights that learned potentials can be informative as ideal local credits.
lot of actions associated
D.2 TIME COST ANALYSIS
We also analyze time costs. Especially, we consider molecular generation tasks and compare LED-GFN with GFN and FL-GGNN on subTB-based objectives. We use single GPU of NVIDIA A5000) for this experiment. In Table Table 1, one can see that LED- GFN only incurs small overheads (approximately 10%) compared to GFN. In Figure Figure 12(b), we also compare the performance with respect to the time costs. One can see that LED-GFN achieves the most promising results with respect to the time costs.
Table Time costs (sec) analysis. The LED- GFN does not incur significant overheads.
Method
Time cost
subTB
5.80 (-)
FL-subTB
11.43 (1 5.63)
LED-subTB
6.61 (* 0.81)
D.3
TB- TB-BASED OBJECTIVE
We also combine energy decomposition learning to TB-based objectives by reparameterizing the see that energy
logits of the policy with the potential (g + s'). In Figure Figure 13, decomposition even improves on the TB. We assume that local credits also benefit the TB since they support assigning high probability to the action responsible for the low terminal energy.
15


Reviews:
[{"review_id": "s01SLdaAjQ", "reviewer": "Reviewer_xgzq", "summary": "The paper studies the problem of credit assignment in the context of GFlowNets. GFlowNets are amortized samplers which are trained to sample from a target energy function. Typically, GFlowets are only trained using this energy as the terminal reward for the generated object.  Even cases where some intermediate energy signals are available, using them can be computationally expensive. Similar challenges are also studied in the context of reinforcement learning in environments with sparse rewards, which includes approaches such as RUDDER which models the decomposition of returns. This paper proposed LED-GFN a method which learns a decomposition of the energy function to enable partial inference in the context of GFlowNets. The method leverages the flow reparameterization from forward-looking GFlowNets to utilize these partial energy estimates. LED-GFN decomposes the terminal state energy into learnable potential functions defined on state transitions which serve as local credit signals. The potentials are trained to approximate the terminal energy through summation and minimize variance along the action sequence. This provides dense and informative training signals. LED-GFN is evaluated on set generation, bag generation, molecule generation, RNA sequence generation, and maximum independent set problems. It outperforms GFlowNet baselines which do not use the  and achieves similar performance to methods using ideal intermediate energies.", "questions": "* Can you comment on the scalability of the approach and how much of the benefit does it provide for longer trajectories? \n* 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)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper proposes an interesting approach to tackle the problem of credit assignment and partial inference in GFlowNets. It builds upon the ability of forward-looking GFlowNets and addresses the limitation of having an intermediate potential function by learning it. \n* The learned potential function can also act as an important inductive bias during training: approximating energy through summation and minimizing variance over the trajectory can make the energy landscape easier to model for the sampler. \n* The method enjoys quite strong empirical performance over baselines on a diverse set of fairly complicated tasks.\n* The experiments in the ablations are quite thorough and well designed and provide interesting insights in the method's performance.  \n* Reproducibility: The authors provide code to reproduce the results for the molecule generation experiments and includes most details to reproduce the results.", "weaknesses": "* Learning the decomposition of the terminal potentials is interesting, but there are no theoretical guarantees that the learned decompositon provides meaningful local credits in all settings. \n* In terms of the empirical results, while the method performs quite well on a variety of tasks - there is an important caveat to note which is 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 assginment would be critical and thus demonstrate the efficacy of the approach. \n* 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. \n* Another important caveat of the empirical analysis is that it focuses on discrete problems and does not consider the continous case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yRmZ3PtrOU", "reviewer": "Reviewer_q5Nd", "summary": "This paper proposes to improve GFlowNet training by assigning rewards or credits to the intermediate steps. Previous work like forward-looking(FL) GFlowNet propose to reparametrize the state flow function and make use of the partial reward. However, this evaluation of the reward of intermediate state can be expensive. Therefore, the author proposes an interesting learnable energy function on state transitions and obtain good results in some tasks.", "questions": "The author(s) use two energy-based decomposition schemes as a base-line, namely, the model-based GFlowNet (Jain et al.) and LSTM-based decomposition (Arjona-Medina et al.). Why not test other base-lines? It seems like a small comparison, and I’m not entirely sure if the LSTM is comparable to a GFlowNet based model given its simplicity.  Another question would be the paper discussed both DB-based and subTB based objectives. What about other objectives?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper targets at an interesting and valuable question. Both empirical and theoretical results seem solid and convincing. By learning decomposition of the reward in the terminal stage using potential functions, it can improve both detail-balance based or sub-trajectory balance based GFlownet training.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o2J4LxnMlI", "reviewer": "Reviewer_RMo5", "summary": "A method is proposed to improve training of generative flow networks (GFlowNets) using a learned reward shaping scheme. Specifically, one trains an auxiliary model to predict an energy delta for every edge, with the objective that the sum of energy deltas along any trajectory should equal the negative log-reward of the state at which the trajectory terminates. The learned energy delta is then used as a correction to the log-flow difference that appears in the detailed balance training objective. This method is shown to accelerate convergence and improve mode discovery in several very different problems from past work where GFlowNets have been successfully applied.", "questions": "- On alternative potential learning:\n  - It is a little surprising to me that learning the $\\phi(s\\rightarrow s')$ works so much better than learning an extension of $\\cal E$ to nonterminal states, just replacing $\\phi(s\\rightarrow s')$ by ${\\cal E}(s')-{\\cal E}(s)$ in the current objective (section 4.5). Such an expression automatically guarantees a cycle consistency property for $\\phi$. \n  - Why do you need to learn a proxy model to predict terminal energies instead of using true rewards?\n  - Additionally, one could regress each $\\phi(s\\to s')$ to $\\frac1T{\\cal E}(s_T)$ (or, using energies, regress ${\\cal E}(s_t)$ to $\\frac tT{\\cal E}(s_T)$). This would also decrease variance and sparsity.\n- I am curious how well the energy decomposition performs with SubTB and how it could be used with TB, which does not require learning state flows (while the current system requires **three** estimators for every edge, in addition to one per state). For example, one could add the predicted energy difference to the logits of the forward policy. Do you have any ideas about this question?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- Good exposition; I have very few complaints on the presentation and math.\n- Natural and well-carried-out idea that seems to substantially improve GFlowNet training.\n- Experimental validation in very diverse settings from past work. Comparisons to alternative ways of learning the energy decomposition make the paper stronger.\n  - But see questions below.", "weaknesses": "- 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.\n  - \"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.\n  - The third paragraph of the introduction and the second paragraph of 2.2 attributes \"partial inference\" to [Pan et al., 2023a], but in fact that paper does not introduce this term and does not ever use it.\n  - The minimal fix would be to make the title \"partial inference ~of~ **in**/**for** GFlowNets\", which would more accurately describe what is being done.\n- Small errors in exposition on GFlowNets:\n  - Error on the top of p.3: $\\exp(-{\\cal E}(x))$ is the reward, not the energy.\n  - 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$).\n  - 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 (+ briefly discuss possible sources of the partial energies).\n- Limitations/costs/failure modes of the proposed algorithm are not discussed.\n- Missing details on the form of the model that predicts $\\phi$ (I could not find them in the appendix). Does it share 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.\n- Some theoretical characterization of the optimal energy decomposition would be helpful (even in the case of a tree-shaped state space).\n- In all plots with curves, please use markers or line styles, and not just colour, to distinguish the curves.\n\nI am very willing to raise the score to 6 or even 8 if the above weaknesses and below questions are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qtDQYoaAQD", "reviewer": "Reviewer_qpM8", "summary": "This paper investigates generative flow networks (GFlowNets), which sample a composite object via a sequence of constructive steps, with a probability proportional to a reward function. In contrast to prior training objectives for GFlowNets, which mainly focus on learning from complete trajectories, Looking-Forward GFlowNets (FL-GFlowNets, as introduced by Pan et al.) take advantage of the computability of intermediate rewards or energies and employ an additive energy factorization to facilitate learning from incomplete trajectories. However, the authors argue that there are two limitations: 1) intermediate rewards are too expensive to be evaluated, e.g., frequent evaluation over partially-constructed molecules at all intermediate steps can indeed be time-consuming, especially in the context of chemical synthesis or molecular design; 2) the significant variability in intermediate rewards along the entire trajectory, e.g., rewards are low or zero at the beginning, and experience a sudden surge at the end (as illustrated in Figure 1), thus leading to less informative intermediate signals. To this end, inspired by the reward decomposition method in Ren et al., the authors propose to learn a potential function that is additive over transitions. This is done by minimizing the least square loss between $R(s_{n})$ and a summation over potential functions that incorporates a dropout-based technique, i.e., $\\sum_{t=0}^{n-1} z_{t} \\phi_{\\theta} (s_{t} \\rightarrow s_{t+1})$, together with some scaling coefficients. The learned potential function can be directly incorpoated into FL-GFlowNets, thus leading to LED-GFlowNets. Experimental results on several datasets show the effectiveness of LED-GFlowNets.", "questions": "### Method:\n\n- 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?\n\n- In terms of Figure 2, \n   - Why do we want to have (b) for potentials --> approximately same energies to minimize variance (more smooth transitions)? In this case, we just simply set $\\frac{1}{3}$ if $\\mathcal{E} = 1$?\n   - 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? I conjecture this might be a common case in many tasks?\n\n- In GFlowNet settings, we hope to achieve a transition probability distribution: $P_{F}(s_{t+1} | s_{t}) \\propto F(s_{t} \\rightarrow s_{t+1})$. I am curious - since we have $- \\log R(s_{n}) = \\mathcal{E}(s_{n}) = \\sum_{t=0}^{n-1} \\phi_{\\theta}(s_{t} \\rightarrow s_{t+1})$, then we might be able to learn $P_{F}(s_{t+1} | s_{t}) \\propto \\phi_{\\theta}(s_{t} \\rightarrow s_{t+1})$? With such policy, can we achieve - sampling $x$ with probability proportional to $R(x)$? If not, with GFlowNet training objectives, the learned policy would be modified accordingly?\n\n### Experiment:\n\n- In terms of Figure 9(b), how to understand number of calls? Assume 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, 16 calls vs. $\\sum_{i=1}^{16} n_{i}$ calls? But, I think we should have more than 16 calls for LED-GFlowNets, as we need to train the potential function.\n\n- For experimental details, 'We set the dropout probability as 10% for tasks with a trajectory length less than 10 and 20% for others.' -- Did you try other proportions? How does $\\lambda$ affect the performance? Abalation studies on $\\lambda$ for Figure 9 (a) are missing.\n\n\n### Some writing issues, typos and inconsistencies:\n\n1) $\\log$ is missing for $P_B$ regarding the DB loss, as well as in the Appendix.\n\n2) As far as I understand, it should be --> SubTB($\\lambda$) [Madan et al.] is practically useful interpolation between TB and DB losses?\n\n3) 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.?\n\n4) Figure 4: SubTB-based objectives --> subTB-based objectives\n\n5) For Figure 4, it is clear that DB or subTB is considered. But for Figure 5 & 6, it's just LED-GFN or FL-GFN. Which training objective are you using? (though '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? Try to find a way to avoid such confusion.\n\n---\n\n**Update after rebuttal**\n\nI raised the score: 6 --> 8.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "### Motivation:\n\nA major limitation of GFlowNets is that they suffer from inefficient credit assignment when trajectories are long, as the learning signal only comes from the terminal reward (episodic setting). This is where partial inference comes in - the learning signal is from partial reward signals. To this end, FL-GFlowNets take advantage of intermediate rewards and an additive energy factorization over transitions to facilitate learning from incomplete trajectories. However, there might be two limitations: 1) intermediate rewards are too expensive to be evaluated; 2) the significant variability in intermediate rewards along the entire trajectory. These motivate LED-GFlowNets.\n\n\n### Originality:\n\nThe proposed method builds on two existing works - reward decomposition [Ren et al.] and FL-GFlowNets [Pan et al.]. The authors extend the reward decomposition method by introducing a dropout-based regularizer to reduce variance and promote sparsity. The learnable potential function $\\phi_{\\theta} (s_{t} \\rightarrow s_{t+1})$ can be directly used in FL-DB or FL-SubTB. The combination of existing ideas shows good performance on various tasks. The originality should be ok.\n\n\n### Clarity:\n\nThe paper is well-organized and easy to follow, making it accessible to readers.", "weaknesses": "Please see the following questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)