Trajectory balance: Improved credit assignment in GFlowNets
Authors: Nikolay Malkin
Moksh Jain
Emmanuel Bengio
Chen Sun
Yoshua Bengio
Venue: NeurIPS 2022 Accept (2022)
Keywords: GFlowNets
generative models
reinforcement learning
molecule design

Abstract: Generative flow networks (GFlowNets) are a method for learning a stochastic policy for generating compositional objects, such as graphs or strings, from a given unnormalized density by sequences of actions, where many possible action sequences may lead to the same object. We find previously proposed learning objectives for GFlowNets, flow matching and detailed balance, which are analogous to temporal difference learning, to be prone to inefficient credit propagation across long action sequences. We thus propose a new learning objective for GFlowNets, trajectory balance, as a more efficient alternative to previously used objectives. We prove that any global minimizer of the trajectory balance objective can define a policy that samples exactly from the target distribution. In experiments on four distinct domains, we empirically demonstrate the benefits of the trajectory balance objective for GFlowNet convergence, diversity of generated samples, and robustness to long action sequences and large action spaces.  

Trajectory balance:
Improved credit assignment in GFlowNets
Nikolay Malkin Mila, Université de Montréal Montréal, Québec, Canada
Moksh Jain Mila, Université de Montréal Montréal, Québec, Canada
Emmanuel Bengio
Chen Sun
Mila, McGill University, Recursion Montréal, Québec, Canada
Mila, Université de Montréal Montréal, Québec, Canada
Yoshua Bengio
Mila, Université de Montréal
Montréal, Québec, Canada
fnikolay mecolay..mallee malkin mww.h. jain, chen. sun, yoshua hew.ooooaa....iioeeeee........ emmanuel, b
Abstract
Generative flow networks (GFlowNets) are a method for learning a stochastic policy for generating compositional objects, such graphs strings, from given unnormalized density by sequences of actions, where many possible action sequences may lead to the same object. We find previously proposed learning objectives for GFlowNets, flow matching and detailed balance, which are analogous temporal difference learning, to be prone to inefficient credit propagation across long action sequences. We thus propose new learning objective for GFlow trajectory balance, as more efficient alternative previously used objectives. We prove that any global minimizer of the trajectory balance objective can define a policy that samples exactly from the target distribution. experiments on four distinct domains, we empirically demonstrate the benefits of the trajectory balance objective for GelowNet convergence, diversity of generated samples, and robustness long action sequences and large action spaces.
Introduction
Generative flow networks [GFlowNets: 3 are models that exploit generalizable structure in an energy function & to amortize sampling from the corresponding probability density function on space of compositional objects X, for example, graphs composed of nodes and edges. A GFlowNet learns stochastic policy that generates such structured objects by producing stochastic sequence of actions that incrementally modify partial object (state), e.g., by adding an edge or a node to a graph, starting from a universal initial state (like an empty graph). A special "exit" action signals that the construction
object æ E X' is completed, and the policy is trained so as to make the likelihood of generating x proportional to the given unnormalized probability or reward R(x) = -E())
Like other models deep reinforcement learning (RL. GFlowNets are trained with a parametric policy that can be given desired inductive biases (c.g.. particular deep net architecture) and allows generalization to states not seen in training. Natural domains for applying GFlowNets are those where exact sampling is intractable and local exploration (MCMC) methods perform poorly, but
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

diversity of samples is desired [3 32 14 8). For example, GFlowNets have been used 3 to generate graphical descriptions of molecules by incremental addition of simple building blocks, where the reward R(x)) is the estimated strength of binding the constructed molecule to protein target: the number of candidates grows rapidly with the molecule size and the reward has many separated modes. Like all RI. models that iteratively sample action sequences for training, GFlowNets pose the learning challenges of exploration/exploiatiio and credit assignment, i.c. propagation of a reward signal over an action sequence [27 22 17]. The efficiency of credit assignment and training in GFlowNets is the focus of the present paper.
The learning problem solved by GelowNets also has two fundamental differences with the standard doward-maxiiizuue paradigm of RL. First, a GFlowNet aims to make the likelihood of reaching a terminating state proportional to the reward, not to concentrate it at a maximal-rewe reward state. Thus, a GFlowNet must model the diversity in the target distribution, not only its dominant mode. Reward maximization in RL can be turned into sampling proportionally to the reward with appropriate entropy maximization regularization, if there is only one way to reach every state [3] The second difference with deward-maximization in RL is indeed that the GFlowNet training objectives still lead to correct sampling even when multiple action sequences lead to the same terminating state. Note that the likelihood of reaching state is the sum of likelihoods of all action sequences leading to it, and that
the number of such paths may be exponential in their length.
The set of all achievable sequences of actions and states can be conceptually organized in directed graph G (S,A) in which the vertices S are states (some of them designated as terminal states, in bijection with X)) and the edges u->r in A each correspond to applying an action while in a state E S and landing in state u. In B a GFlowNet is described by a nonnegative function on the edges, called the edge flow F A + Rzo. where F(u->e) is an unoormalized likelihood of taking the action that modifies state u to state v. The GFlowNet policy samples the transition u->v from state with probability F(u-x)/ En, F(u->e') where the denominator sums over the outgoing edges from u. By analogy with the classical notion of flows in networks (9. one can think of this flow like the amount of water flowing through an edge (like a pipe) or a state (like tee, where pipes meet). is shown that this GFlowNet policy samples I proportionally to R(x) if F satisfies a set of linear flow matching constraints (a conservation law: the sum of flows into state should equal the sum of flows out of it). These constraints are converted into a temporal difference- objective that can be optimized with respect to the parameters of a neural net that approximates F. An alternative objective based on detailed balance constraints was proposed in 4 These objectives, however, like temporal-difee -iffereeee learning, can suffer from slow credit assignment 27 17).
The main contribution of this work ($3) is a new parametrization and objective for GFlowNets. This objective, which we call trajectory bulance, is computed on sampled full action sequences (trajectories) from the initial state to terminal state, unlike the flow matching and detailed balance objectives. We prove that global minimization of trajectory balance implies that the learned action policy samples proportionally to R. We also empirically show that the trajectory balance objective accelerates training convergence relative to previously proposed objectives, improves the learned sampling policy with respect to metrics of diversity and divergence from the reward function, and allows learning GFlowNets that generate sequences far longer than was possible before. As secondary contribution, we perform the first empirical validation of the detailed balance training objective. Comparative evaluation of the three GFlowNet objectives and non-GFlowNet baselines is performed on four domains illustrating different features of the reward landscape:
Hypergrid (5..)) an illustrative synthetic environment with modes separated by wide troughs; Molecule synthesis (5.2), a practical graph generation problem, where the trajectory balance objective leads to significant computational speed-ups and more diverse generated candidates;
Sequence generation ($5.) where we show the robustness of trajectory balance to large action spaces and long action sequences on synthetic data and real AMP sequence data.
Since the initial appearance of this work on arXiv, several published papers and preprints have used trajectory balance and its generalizations successfully in various applications (32 14 g 18
2

---

2 Preliminaries
2.1 Markovian flows
We give some essential definitions, following $2 of [4]. Fix a directed acyclic graph G = (S, A) with state space and action space A. Let sg be the special initial (source) state, the only state with no incoming edges, and designate vertices with no outgoing edges as terminal (sinks) We call the vertices states, the edges actions, the states reachable through outgoing edges from state its children, and the sources of its incoming edges its parents.
complete trajectory is sequence of transitions - -s.. -8.) going from the initial state to a terminal state 8n with (8.-281++) 1 E A for all Let be the set of complete trajectories. frajectory flow is nonnegative function F T- With water analogy, could be the number of water molecules travelling along this path (the units don't I matter because the flow function can be scaled arbitrarily, since we normalize them to get probabilities). For any state s, define the state flow F(s) Eser F(T), and, for any edge s s-s',, the edge flow
F(s->8))
Σ
F(T).
(1)
T=
As a consequence this definition, the flow matching constraint (incoming flow = outgoing flow) is satisfied for all states that are not initial terminal:
F(s) = Σ FS)
Σ F(s-ts').
(8" -B)EA
(a-t
nontrivial (i.e., not identically zero) trajectory flow F determines a distribution P over trajectories,
1 P(r) 7 F(T) Z F(s))
F(T).. TET
The trajectory flow Markovian there exist action distributions (-(s) over the children of each nonterminal state s such that the distribution P has a factorization
P(T ( s8,)) = II Pr(s)s--1)
(4)
Equivalently (4. Prop. 3) there are distributions Pn(-Is) over the parents of each noninitial state such that for any terminal x,
P(T (so ... +8n)(sn =x) II
PB Ps(st-1sss)
(5)
t=1
If F is Markovian flow, then Pp and PB can be computed in terms of state and edge flows:
Pr(('ss F(s->5)) Pa(s|s')= = F(s')
F(s-ts)
(6)
F(s)
supposing denominators do not vanish. We call PF and Pn the forward policy and backward policy corresponding to F, respectively. These relations are summarized by the detailed balance constraint
F(s))F((((s))
Uniqueness properties. A Markovian flow is uniquely determined by an edge flow, i.e., nontrivial choice of nonnegative value every edge satisfying the flow matching constraint (2). By Corollary 1 of [4) a Markovian flow is also uniquely determined by either of
a constant Z F(so) and a distribution Pr(-ls) over children every nonterminal state; or a nontrivial choice of nonnegative state flows F(z) for every terminal state and choice of distribution (=(- s) over parents of every noninitial state.
3 allowed terminal states with outgoing edges. The difference is easily overcome by augmenting every such state by new terminal state with stop action x=>2"

---

2.2 GFlowNets
Suppose that nontrivial nonnegative reward function R: X' - R2z is given on the set of terminal [ aim to approximate a Markovian flow F on G such that
states. GFlowNets
F(x) R(x) Vo E X.
We adopt the broad definition that a GFlowNet any learning algorithm consisting of:
a model capable of providing the initial state flow = (50) as well as the forward action distributions Pr(--s) for any nonterminal state (and therefore, by the above, uniquely but possibly in an implicit way determining a Markovian flow F);
objective function, such that if the model is capable of expressing any action distribution and the objective function is globally minimized, then the constraint (8) is satisfied for the corresponding Markovian flow F.
The forward policy of GFlowNet be used to sample trajectories from the corresponding Markovian flow F by iteratively taking actions according to policy (-s). the objective function is globally minimized, then the likelihood of terminating
is proportional to R(x).
In general, an objective optimizing for (8) cannot be minimized directly because F F(x) sum over all trajectories leading to and computing it may not be practical. Therefore, two local objectives- flow matching and detailed balance - have previously been proposed.
Flow matching objective 33.. A model (s, s') F with learnable parameters 0 approximates the edge flows F(s->s). The corresponding forward policy is given by (s (s; 0) Fe s') (Eq. (6)) Denote the corresponding Markovian flow by Fe and distribution over trajectories by Pe. The parameters are trained to minimize the error the flow matching constraint (2) for all noninitial and nonterminal nodes 8:
(FM log Eise-seE Fo(s", 2
E(--A)JA Fs 8')
and a similar objective FTM pushing the inflow at x E X' to equal R(xx at terminal nodes ær This objective is optimized for nonterminal states and terminal states x from trajectories sampled from training policy Usually, To is chosen to be tempered (higher temperature) version of (- 0), which also helps exploration during training. The parameters are updated with stochastic gradient
Σ LFM(s:) + (s)) [=1
(10)
As per Proposition 10 of 44, if the training policy my has full support, and a global minimum of the expected loss (9) over states on trajectories sampled from mo is reached, then the GFlowNet samples from the target distribution (i.e., Fg satisfies (8))
Detailed balance objective [4. A neural network model with parameters has input and three kinds of outputs: an estimated state flow (s). an estimated distribution over children P=(-Is; 0). and an estimated distribution over parents PB( ( (s; 0). The policy Pr (-1-:0) and the initial state flow (so) uniquely determine a Markovian flow Fe. which is not necessarily compatible with the estimated backward policy ()) The error in the detailed balance constraint (7) is optimized
on actions (s-+') between nonterminal nodes seen along trajectories sampled from the training policy:
LDn(s, log Fo(s))P(s'ss:0) Fo(ss)Ppssss':0)
(11)
and a similar constraint Co (s) flow matching, the parameters are updated with stochastic gradient
x) is optimized at actions leading to terminal nodes. Similarly to
(DB(#-1, + LDB ( (8n-1, 8x)
(12)
along trajectories sampled from training policy TT. By Proposition 6 of [4], global minimum of the expected detailed balance loss under a Te with full support specifies a GFlowNer that samples from the target distribution, i.e., the flow Fg satisfies g
IIn practice. it is convenient and more economical to provide a representation of to the neural net, which simultaneously outputs the flows Fo s') for all s' that are reachable by action from

---

Algorithm Training a GFlowNet with trajectory balance
input Reward function R X' Ryo, model and optimizer hyperparameters
1: Initialize models PF, PB. Z with parameters 0
2: repeat
3: Sample trajectory 7 == (sg (s0-..... 8,) from policy Pr(-l-;0) or a tempered version of
4: 0 - 0 (TD((T (gradient update on (14))
5: until convergence monitoring on running (Tn(T)
Remarks. In some problems, such autoregressive sequence generation ($5.3) the directed graph G is a tree, so each state has only one parent. In this case, Pa trivial and the detailed balance objective reduces the flow matching objective, which in can be shown to be equivalent to Soft Q-Learning 12 5
with temperature a 1, uniform ga', and = 1.
3 Trajectory balance
Let F be Markovian flow and P the corresponding distribution over complete trajectories, defined by (3). and let Pr and P3 be forward and backward policies determined by F. A direct algebraic manipulation of Eqs. (145) gives the trajectory balance constraint for any complete trajectory
-81 +8p x):
2
II Pr (=())(111)) F(x) II Pa Pa((t-1(5)).
(13)
where we have used that P(s, = x) F(x) 22-
As explained in $A.2 the trajectory balance constraint 13 and the detailed balance constraint 7 special cases of one general constraint, which has been studied as a training objective in [8
Trajectory balance as an objective. We propose to convert [3) into an objective to be optimized along trajectories sampled from training policy. Suppose that a model with parameters é outputs estimated forward policy Pp[-|s; 9) and backward policy (=[-Is: 0) for states s (just as for detailed balance above), as well as global scalar Zo estimating ((p)). The scalar Zo and forward policy
Pr(-1-;0) 0) uniquely determine an implicit Markovian flow Fo.
For trajectory (So 381 .... Snn x)), define the trajectory loss
Zo I PF (p
(TB(T)= = [log R(x)) IT= PB (st 19:;)1));)))
(14)
1=
If Ty training policy usually that given by Pr(-l-;0) trajectory loss is updated along trajectories sampled from To. i.e., with stochastic gradient
0) or tempered version of it then the
E VOLTB(T).
(15)
The full algorithm, with batch size of 1, is presented as Algorithm and its correctness is guaranteed by the following.
Proposition 1. Let R be a positive reward function on X.
(a) If Pp(-1-:8), P=(-1-:0), and Zg are the forward and backward policies and normalizing constant Markovian flow F satisfying (8), then LTB(T) 0 for all complete trajectories (b) Conversely, suppose that LTB O for all complete trajectories т. Then the corresponding Markovian flow Fg satisfies and Pr(-1-:0) samples proportionally to the reward.
The proof is given in SA.1 In particular, if σe has full support and BTCr,CBB(7 is globally minimized over all forward and backward policies (Pr, PR)) and normalizing constants Z, then the corresponding Markovian flow Fg satisfies (8) and Pr(-l-;0) samples proportionally to the reward. (The positivity assumption on R is necessary to avoid division by in (14). but can be relaxed by introduction of smoothing constants, just was done for the losses proposed in 3 4.)
5

---

Remarks. (1) As discussed in 2 in the case of auto-regressive generation, G is directed tree, where each S has a single parent state. In this case Pn is trivially Pn vs E S. We get
(Ta(T)= = =(10g Ze II PF (=((1(sc-1:0)
(16)
R(x)
(2) We found it beneficial to parametrize Z in the logarithmic domain (log Z is the trainable parameter) and output logits for (-(-1s; 8; 0) and Pn(-|s; 0), that all products in (14) become logarithm. This is consistent with the log-domain parametrization of flows in [3] In addition, we
under the
found helpful
higher learning rate for Z than for the parameters of PF and PB
3.1 Canonical choice of reward-matching flow
The constraint [8) in general, does not have a unique solution: if the underlying undirected graph of G has cycles, there may be multiple Markovian flows whose corresponding action policies sample proportionally to the reward. However, by the uniqueness properties, for any choice of backward policy PA (-1-- there is a unique flow satisfying (8) and thus unique corresponding forward policy Pp(-1- for states with nonzero flow. (See Fig.
In some settings, may be beneficial fix the backward policy PB and train only the parameters giving Pr and Zg. For example, it may difficult to construct a model that outputs a distribution over the parents of given input state (e-g., for the molecule domain ($5.2), it is hard force invariance to molecule isomorphism). A natural choice is to set PB( to be uniform over all the parents of state s, i.e. Ps(-(s) #/# (s' (s' E A3.
Related work
Reinforcement learning. GFlowNets are trained to sample proportionally the reward rather than maximize it as usual in RL. However, tree-structured DAGs (autoregressive generation) equivalent to RL with appropriate entropy regularization or soft Q-learning and control as inference 12 1.. The experiments and discussion of [3] show how these methods can fail badly in the general DAG case well handled by GFlowNets. Signal propagation over sequences of several actions in trajectory balance is also related losses used in RL computed subtrajectories 22].
Local exploration vs. amortized generalization to unseen modes. GFlowNets are also related to MCMC methods for sampling from unnormalized densities. While there has been work on accelerating or partially amortizing sampling from unnormalized densities over discrete spaces when exact sampling is intractable 71] some of it domain- problem-specific 300, GFlowNets treat the compositional structure in data as a learning problem (enabling generalization to unseen modes), not bias build in to the sampler. Thus, the cost is amortized and borne by the learning of that structure through sampling, not through search at generation time.
Variational inference. GFlowNets are connected with variational methods for fitting hierarchical generative models. The squared log-ratio loss proposed in 200 as control variate in the optimization of an evidence lower bound can be seen as special case of trajectory balance. See AA.3 for further discussion, in which we prove that on-policy optimization of trajectory balance is equivalent to minimization of certain KL divergence. Two recent papers [3] [19] extend our observations.
5 Experiments
We evaluate the proposed trajectory balance objective against prior objectives for training GelowNets as well as standard methods for learning policies that approximately sample objects proportionally their rewards, like MCMC well against other RL techniques. Our experiments include the hypergrid and molecule synthesis tasks from and two new tasks in which G is directed tree.
'Because the loss (14 is quadratic in log Z, gradient updates on log 2 are equivalent to setting it to a weighted moving average of the discrepancy between log Pr and (R(())) Ps). Optimizers enhanced with momentum complicate things: we leave empirical investigation of these questions to future work.

---

5.1 Hypergrid environment
In this subsection, we study a synthetic hypergrid environment introduced by B.. This task is easier than others we study, but we include it for completeness, and because it allows us illustrate some interesting behaviours
In this environment, the nonterminal states 8 form D-dimensional hypergrid with side length H::
S = ((s p...,, s E (0, (0,1,..... H H -
and actions are operations of incrementing one coordinate in a state by without exiting the grid. The initial state is transitions corresponding terminal state ST (cf. footnote 1). The reward terminal state (s)..... (...,."))) s*) is given by
0). For every nonterminal state 8, there is also a termination action that
s H
sd
R(sT) Ro + 0.5
II d=1
0.5 (0.25, 0.5)
2
11-1 0.5 E d=1
III
(0.3, (043
where I is indicator function and Ro is constant controlling the difficulty of exploration. This reward has peaks of height 2.5 + Rg near the corners the hypergrid, surrounded by plateaux of height 0.5 + Rg. These plateaux are separated by wide troughs with reward Ro. An illustration with H = and D = 2 is shown in the left panel of Fig. 1 This enviromment evaluates the ability of a GFlowNet to generalize from visited states to infer the existence of yet-unvisited modes.
We train GFlowNets with the detailed balance (DB) and trajectory balance (TB) objectives with different H, D, and Ro, in addition to reproducing the flow matching (FM) experiments and non- GFlowNet baselines based upon (3)s published code. Our GFlowNet policy model is multilayer perceptron (MIP) that accepts as input one-hot encoding of state (with the goal of enabling generalization) and outputs logits of the forward and backward policies PF(-1-;0 and (-1(-1=;0) (as well as the estimated state flow (s) in the case of DB). The forward policy, backward policy, and state flow models share all but the last weight matrix of the MLP. This is consistent with (3's model, where identical architecture accepted as input and output estimated flows Fo s') for all children s' s. Details are given in BB.1
We consider 4 dimensional grid with H = and and 2-dimensional grid with H = 64. The two grids have the same number of terminal states, but the 2-dimensional grid has longer expected trajectory lengths. For both grid sizes, we consider Ro 0.1, 0.01, 0.001, with smaller Ro giving environments that are more difficult to explore due to the lower likelihood for models to cross the low-reward valley. For the models trained with DB and TB, we also explore the effect of fixing the backward policy to be uniform (.3.I).
Results. In Fig. 2 we plot the evolution over the course of training of the L error between the true reward distribution (the reward R(x) normalized over all possible terminal states x E X²) and the empirical distribution of the last 2- 10 visited states for all settings (which would have a probability of 0 on xts not visited). Although convergence the same stable minimum is achieved by all models and settings, DB and TB training tend to converge faster than FM, with a slight benefit of TB over detailed balance in the 4-D environment.
Effect of uniform PB- Note the difference in learning speed between models with fixed uniform backward policy PB and models with learned PB. As noted in $3.1 when Pg is fixed, there is a unique PF (=(--=:0) 0) that globally minimizes the objective, and it may be approached slowly. However, if PB and PF are permitted
evolve jointly, they may more quickly approach one of the many optimal solutions. This is confirmed by the much faster convergence of models with learned PB on the 64 x 64 grid. We have observed that, especially for large grid sizes, when Pg and Pp are both learned. the model has a bias towards first taking all actions in one coordinate direction, then proceeding the other direction until terminating (as the right panel of Fig. [ perhaps because a constant distribution over two actions ('continue to the right' and 'terminate") ) can be modeled with higher precision over a large portion of the grid than the complex position-dependenn distribution as shown in the centre panel of Fig.
7

---

Figure 1: Left: The reward function 8 8 grid environment ($5.1) with Ro 0.1. Centre and right: Two forward action policies - with fixed uniform PB and with learned non-uniform PB that sample from this reward. The lengths of arrows pointing up and right from each state are proportional to the likelihoods of the corresponding actions under PF, and the sizes of the red octagons are proportional the termination action likelihoods.
Bx8x8x6 Ru- 10
3xdx8888, 10-2
Sx
64x64, 10-1
64v64 10
64v64,
Figure 2: Empirical L' error between true and sampled state distributions on the grid environment with varying grid size and Rg. Mean and standard error over 5 seeds. The curves for PPO and MCMC baseline would lie outside the plot bounds.
5.2 Small drug molecule synthesis
24 introduced for GFlowNets in K We extend 's published code with implementation of the TB and DB objectives F
Next, we consider the molecule generation task 30 15 16 10
The goal is to generate molecules, in the form of graphs, with a low binding energy to the 4JNC inhibitor of the sEH (soluble epoxide hydrolase) protein. The graphs generated are junction trees [1]
of vocabulary of building blocks. The reward is defined as the normalized negative binding energy as predicted by a proxy model, itself trained to predict energies computed via docking simulations [26]. The maximum trajectory length is 8, with the number of actions varying between around 100 and 2000 (the larger molecule, the more possible additions exist), making |X| about 16
Results. We plot in Fig. (left and centre) the correlation log-reward and log- sampmimmiig probability (the likelihood that trajectory sampled from the learned policy terminates at 2) for GFlowNets trained using TB, FM and DB. This correlation is significantly higher for models trained with TB. The points shown are from fixed held-out set to which the GFlowNets do not have access in training. Note that perfect model would have correlation 1, as log R(x) and log po (x) would differ by constant (equal to log Z) that is independent of x.
In Fig. 3 (right) we plot the average pairwise Tanimoto similarity ① for the 1000 highest-reware samples generated over the course of training. We see that TB consistently generates more diverse molecules than FM. These results showcase the benefits of TB, not only for faster temporal credit assignment, but for generalization and diversity. In addition, TB has up to 5x runtime speedup over FM as the enumeration of parents not needed. See $B.2 for further discussion.
5.3 Autoregressive sequence generation
Finally, we evaluate the TB objective on the task of autoregressive sequence generation. In $5.3. we study the effect trajectory length and action space size on the learning dynamics GFlowNets. In 5.3.2 consider the more realistic task of generating peptides (short protein sequences with anti- anti-microbial properties) and evaluate GFlowNett against standard RL and MCMC basclines.
CCode: hutp://sthieeeeuuoooooooooooeerr ballnce

---

-20
Corr. log pabx) log R(x)
Top- 10000 Tanimolo sim.
FM DB
log Palx)
0.4
-60
022
-80
-0.2
08 0.8
PM 0.
-0.4
-120
0.5
1e5
log R())
objective
molecules seen
Figure 3: Left, Centre: Pearson correlations between rewards and sampling probability. log po (ar) is the log-likelihood that a trajectory sampled from the learned policy Pp(-1-:0) terminates at z. Left: Scatter plot
set of x's for the best hyperparameters of TB, FM, and DB. Centre: Violin plot of correlations for 16 hyperparameter settings and 3 seeds for each setting, showing TB being capable of fitting better. Right: Average pairwise Tanimoto similarity for the top 1000 samples generated by GFlowNets as training progresses. Lines are the average across runs, shaded regions the standard deviation. Models trained with TB have consistently lower similarity than those with FM, hence greater diversity. We hypothesize that the higher variance, in correlation and diversity, of TB relative DB is related to high variance of the stochastic gradient; see [8]
5.3.1
Bit sequences
Task. The goal is to generate bit sequences of a fixed length n = 120 (X' {0, 1)"), where the reward designed to have modes at given fixed set M X that is unknown the learner. The reward for sequence is defined as R(x)) = exp( mingeM d(I, y)), where d is the edit distance. We describe the procedure generate M in |B.3
For different integers is dividing n. we design action spaces for left-to-right generation of sequences in X, where complete trajectory has & actions and each action appends k-bit 'word' to the end of a partial sequence. A forward policy on this state space is the same an autoregressive sequence model over vocabulary of size 2* Varying while fixing X' and M allows us study the effect of the tradeoff between trajectory lengths ( and the action space sizes ([V) = 2*) without changing the underlying probabilistic modeling problem.
We compare GFlowNets trained with the TB objective against GFlowNets trained with the FM objective (equivalent to DB and Soft O-Learning in this case) and two Nen-GeloWNet baselines: A2C with Entropy Regularization
29 21. Soft Actor-Critic [13 and MARS [30]. We use a
Transformer-based architecture 28 across all the methods. See
B33 for details.
To evaluate the methods we use (1) Spearman correlation between the probability of generating the sequence p(z) = Fzx)/Z and its reward R(x) on test set sampled approximately uniformly over the possible values of the reward, (2) number modes captured (number of reference sequences from M for which candidate within distance å has been generated). In our experiments, = 120, |M 60, & E (1, 11,2,4 4, 6, 8, 10), and 8 = 20.
Results. Fig. 4 (left) presents the results for the Spearman correlation evaluation. We observe that GFlowNets trained with the TB objective learn policies that correlate best with the reward R(x) across all action spaces. In particular, we observe the effect of inefficient credit assignment in GFlowNets trained with FM, as the correlation improves with increasing i.e., shorter trajectories. On the other hand, large action spaces also hurt GFlowNets trained with the FM objective, while the TB objective is robust them. Additionally, we can observe in Fig. 4 (right) that for fixed GelowNets trained with TB discover more modes faster than other methods.
5.3.2 Anti-Microbial Peptide (AMP) generation
Table 1: Results on the AMP generation task.
Top 100 Reward
Top 100 Diversity
In this section, we consider the practical task gen- erating peptide sequences that have anti-microbial activity. The goal to generate a protein sequence (where the vocabulary consists of 20 amino acids and
GEN- Cru
0.85 + 0.03
18.35 +-166
GFN-Lrxs/Cns
0.78 0.05
12.61 4: 1.32
SAC
0.80 0.01
8.36 = 1.44
AAC-ER
0.79 +: 0.02
7.32 0.76
MCMC
0.75 - 0.02
12.56 + 1.45

---

GEN Belaxe
07-
Macling
Trnpy
Seet AcsinCont
Number of Modea
ACMM
044
0.1-
0.1
0.
Narber Notis
Todking
Figure 4: Left: Spearman correlation of the sampling probability under different learned policies and reward on a test set, plotted against the number of bits k in the symbols in V in the bit sequence generation task. GFlowNets trained with trajectory balance learn policies that have the highest correlation with the reward R(z) and are robust to length and vocabulary size. Right: Number of modes discovered over the course of training on the bit sequence generation task with k 1. GFlowNets trained with trajectory balance discover more modes faster.
special end-of-sequence action), with maximum length 60. We take 6438 known AMP sequences and 9522 non-AMP sequences from the DBAASP database [23]. We then train classifier on this dataset, using 20% of the data validation set. The probability output by this model for sequence to be classified as an AMP is used as the reward R(r) in our experiments.
The state and actions are designed just in the previous experiment, with each action appending a symbol to the right a state. We again compare TB and FM, as well as A2C with entropy regularization, SAC and MCMC baselines. We again use Transformers for all the experiments this task; see further details in Appendix [..4 We generate 2048 sequences from each method, and pick the top 100 sequences ranked by their reward R(z). As metrics, we use the mean reward for these 100 sequences and the average pairwise edit distance among them as a measure of diversity.
Results. We present the results in Table 1 where observe that GFlowNets trained with TB outperform all baselines on both performance and diversity metrics.
Discussion and conclusion
We introduced novel training loss for GFlowNets, trajectory balance (TB), which yields faster and better training than the previously proposed flow matching (FM) and detailed balance (DB) losses. We proved that this objective, when minimized, yields the desired GFlowNet property of sampling from the target distribution specified by an unnormalized reward function. This new loss was motivated by the observation that the FM and DB losses are local in the action sequence and may require many iterations for credit assignment to propagate to early actions: if gradient update introduces a flow inconsistency some state far from the initial state (such as when novel high-reward state is sampled), the parent of this state must be visited before the update is propagated closer to the root, akin to the slow propagation of reward signals in temporal difference learning.
We empirically found that TB discovered more modes of the energy function faster and was more robust than FM and DB to the exponential growth of the state space, due in part the lengths of sequences and in part to the size of the action space. A factor to consider when interpreting our experimental results is that because we use neural net rather than tabular representation of policies, the early states' transitions are informed by downstream credit assignment via parameter sharing. Early states also get many more visits because there more possible states near the ends of sequences than near the initial state. Finally, TB trades off the advantage of immediately providing credit to early states with the disadvantage of relying on sampling of long trajectories and thus potentially higher variance of the stochastic gradient. The high gradient variance is possible limitation of TB in difficult environments, and ways overcome it by interpolating between local and trajectory-level objectives have been studied in subsequent work (18).
All in all, found that trajectory balance is superior training objective in broad set of experiments, making it the default choice for future work on GelowNets.
10

---

Acknowledgments
This research was enabled in part by computational resources provided by Compute Canada. All authors are funded by their primary academic institution. We also acknowledge funding from CIFAR, Samsung. IBM, Microsoft, and the Banting Postdoctoral Fellowship.
The authors are grateful to all the members of the Mila GFlowNer group, in particular to Dinghuai Zhang, for many fruitful research discussions, well as Yiheng Zhu for feedback the published code. We also thank the anonymous reviewers for their comments.
References
[1] Bender, A. and Glen, R. C. Molecular similarity: a key technique in molecular informatics. Organic & biomolecular chemistry, 2(22):3204-3218, 2004.
[2] Bengio, E., Pineau, J., and Precup, D. Interference and generalization temporal difference learning. International Conference on Machine Learning (ICML), 2020.
[3] Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
[4] Bengio, Y., Lahlou, S., Deleu, T., Hu, E., Tiwari, M., and Bengio, GFlowNet foundations. arXiv preprint 2111.09266, 2021.
[5] Buesing, L., Heess, N., and Weber, T. Approximate inference in discrete distributions with Monte Carlo tree search and value functions. Artificial Intelligence and Statistics (AISTATS), 2019.
[6] Christodoulou, P. Soft actor-critic for discrete action settings, arXiv preprint arXiv:1910.07207. 2019.
[7] Dai. H., Singh, R., Dai, B.. Sutton, C., and Schuurmans, D. Learning discrete energy-based mod- els via variii local exploration. Neural Information Processing Systems (NeurIPS), 2020.
[8] Deleu, Góis, A., Emezue, C., Rankawat, M., Lacoste-Julien, S., Bauer, S., and Bengio, Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
[9] Ford, L. R. and Fulkerson, D. R. Maximal flow through a network. Canadian Journal of Mathematics, 8:243-240 1956.
[10] Gilmer, J., Schoenholz, S. S., Riley, F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chemistry. International Conference on Machine Learning (ICML), 2017.
[11] Grathwohl, W., Swersky, K., Hashemi, M., Duvenaud, D. K., and Maddison, C. J. Oops took a gradient: Scalable sampling for discrete distributions. International Conference on Machine Learning (ICML), 2021.
[12] Haarnoja, T., Tang, H., Abbeel, P., and Levine, S. Reinforcement learning with deep energy- based policies. International Conference on Machine Learning (ICML), 2017.
[13] Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critie: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. International Conference on Machine Learning (ICML), 2018.
[14] Jain, M., Bengio, E.,, Bermandez-Garcia A., Rector-Brooks, J., Dossou, B. F., Ekbote, C., Fu, J., Zhang, T., Kilgour, M., Zhang, D., Simine, L., Das, P., and Bengio, Y. Biological sequence design with GelowNets. International Conference on Machine Learning (ICML), 2022.
[15] Jin, W., Barzilay, R., and Jaakkola, T. Chapter 11. junction tree variational autoencoder for molecular graph generation. Drug Discovery, pp. 228-249. 2020. ISSN 2041-3211.

---

[16] Kumar, A., Voet, A., and Zhang, K. Fragment based drug design: from experimental to computational approaches. Current medicinal chemistry, 19(30))518-5544) 2012.
[17] Kumar, A., Agarwal, R., Ghosh, D., and Levine, Implicit ander-parameteriiationn inhibits data- efficient deep reinforcement learning. International Conference Learning Representations (ICLR), 2021.
[18] Madan, K., Rector-Brooks J., Korablyov, M., Bengio, E., Jain, M., Nica, A., Bosc, T., Bengio, Y., and Malkin, N. Learning GelowNets from partial episodes for improved convergence and stability. arXiv preprint 2209. 12782, 2022.
[19] Malkin, N., Lahlou, S., Deleu, T., Ji, X., Hu, E., Everett, K., Zhang, D., and Bengio, GFlowNets and variational inference. arXiv preprint 2210.00580, 2022.
[20] Mnih, A. and Gregor, K. Neural variational inference and learning in belief networks. Interna- tional Conference on Machine Learning (ICML), 2014.
[21] Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D., and Kavukcuoglu, K. Asynchronous methods for deep reinforcement learning. International Conference on Machine Learning (ICML), 2016.
[22] Nachum, O., Norouzi, M., Xu, K., and Schuurmans, D. Bridging the gap between value and policy based reinforcement learning. Neural Information Processing Systems (NeurIPS), 2017.
[23] Pirtskhalava, M., Amstrong, A., Grigolava, M., Chubinidze, M., Alimbarashvili, E.,, Vish- nepolsky, B., Gabrielian, A., Rosenthal, A., Hurt, D. E., and Tartakovsky, M. Dbaasp v3: Database of antimicrobiallcytotoiii activity and structure of peptides resource for develop- ment of new therapeutics. Nucleic Acids Research, 49(DI):D288 -018888-D29,.22 2021.
[24] Shi, Y., Huang, Z., Feng, S., Zhong, H., Wang, W., and Sun, Y. Masked label prediction: Unified message passing model for semi- comi-supervised classification. International Joint Conference on Artificial Intelligence (IJCAI), 2021.
[25] Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction. MIT Press, 2018.
[26] Trott, O. and Olson, A. J. AutoDock Vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, (1(2):455 461, 2010.
[27] van Hasselt, H., Doron, Y., Strub, F., Hessel, M., Sonnerat, N., and Modayil, J. Deep reinforce- ment learning and the deadly triad. arXiv preprint 18 02648, 2018.
[28] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. Neural Information Processing Systems (NeurIPS), 2017.
[29] Williams, R. J. and Peng, J. Function optimization using connectionist reinforcement learning algorithms. Connection Science, 3(3):241 -268, 1991.
[30] Xie, Y.. Shi, C., Zhou, H., Yang. Y.. Zhang, W., Yu. Y., and L. MARS: Markov molecular sampling for multi-objective drug discovery. International Conference on Learning Representa- tions (ICLR), 2021.
[31] Zhang, D., Chen, R. T. Q. Malkin, N., and Bengio, Y. Unifying generative models with GelowNets. arXiv preprint 2209.02606, 2022.
[32] Zhang, D., Malkin, N., Liu, Z., Volokhova, A., Courville, A., and Bengio. Y. Generative flow networks for discrete probabilistic modeling. International Conference on Muchine Learning (ICML), 2022.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] This is mainly a paper about theory and algorithms. Future work that applies these algorithms, in particular for domains where they can most immediately have an impact (e.g., molecule design for drug discovery), should consider the limitations and negative societal impacts of these applications.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] See previous.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See SA.1
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]] For the grid and molecule environments. [No] For the other environments.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See the Appendix.
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See the Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See the Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] See the relevant experiment sections.
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material as URL? [N/A] References to the molecule and AMP sequence data are provided.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A] No new data collection was done.
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A] Not relevant for the domains studied.
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "iYrGGzb4Z6", "reviewer": "Reviewer_F9mg", "summary": "The paper proposes a new objective for GFlowNets and shows improved performance in experiments.", "questions": "Can we add more explanations on how to interpret figure 3?", "limitations": "N.A.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Clear presentation and introduction of the background of GFlowNets.\n2. Theoretical backup on the proposed objective.\n3. Clear performance improvements.\n\nWeaknesses:\n1. Figures 2 and 3 are a bit hard to read. ", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "DazjlT8N2e", "reviewer": "Reviewer_9Af4", "summary": "This paper proposes a new objective function for Generative Flow Networks (GFlowNet), which aims to generate objects with probability proportional to a given energy function. The proposed objective function is termed trajectory balance. This objective comes from a constraint that the flow of each complete trajectory should be consistent by either forward computation or backward compution. This constraint is formulated by Eq. (13), and is further converted to an objective function given by Eq. (14).\n\nTrajectory balance suggests parameterizing the forward and backward probabilities (P_F and P_B) between connected states, as well as the flow volume Z of each state. The parameters can be trained by stochastic gradient descent. The authors further prove that, if Eq. (14) is minimized to zero, the GFlowNet can generate samples exactly from the target distribution. Experiments further show that trajectory balance outperforms the conventional two objective functions (flow matching and detailed balance) in terms of convergence, diversity of generated samples, and robustness to long action sequences and large action spaces.", "questions": "Please refer to the weakness.", "limitations": "Please refer to the weakness.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\nThis paper proposes an effective new objective function of training GFlowNets termed trajectory balance (TB). As a new learning method for generating complicated compositional objects, GFlowNet has recently gained an increasing research interest. Compared with conventional objectives, i.e., flow matching (FM) and detailed balance (DB), trajectory balance is empirically observed to overcome the difficulty of slow credit assignment, and gains a faster convergence speed and more diversified generated samples. This paper clearly sheds light on the research of generated flow methods, providing a new inspiring methodology for applying GFlowNet to various domains.\n\n**Weakness**:\n\nAlthough this paper demonstrates the success of TB through substantial experiments, it is not that clear why it performs better than FM and DB. It is argued that one benefit of TB is that the global minimizer can define a policy that samples exactly from the target distribution. However, this also holds for FM and DB. Another benefit claimed for TB is the improved credit assignment, but very few discussions are presented. Specifically, why does considering complete trajectories lead to better credit propagation in GFlowNet training?\n\nAnother concern is that the paper claims that the global minimizer of TB gives rise to exactly the target distributions. However, this is on the assumption that the parameterization of GFlowNet has sufficient representation power, i.e., Eq. (14) can be trained to exactly zero (Proposition 1). This may not always be the case if the task is too complicated or the training data is non-realizable. The authors may want to remind the readers of the above issue in the introduction and relevant discussions.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "7XkSLcThBNX", "reviewer": "Reviewer_Bg1m", "summary": "This paper proposes a new learning objective for GFlowNets, trajectory balance TB), as a more efficient alternative to previously used objectives. A GFlowNet is a trained stochastic policy or generative model, trained such that it samples objects $x$ through a sequence of constructive steps, with probability proportional to $R(x)$, where $R$ is some given non-negative integrable reward function.\n\nMore specifically, GFlowNets aim to approximate a Markovian flow F on G such that\n\n$F(x) = R(x)~~ \\forall{x} \\in X$\n\nGenerally speaking, a GFlowNet could be any learning algorithm consisting of:\n(1) a model capable of providing the initial state flow and the forward action\ndistributions for any nonterminal state s, determining a Markovian flow F.\n(2) an objective function, such that if the model is capable of expressing any action distribution and the\nobjective function is globally minimized, then the constraint is satisfied for the corresponding Markovian flow F .\n\nThis work focuses on finding a more appropriate objective function, distinct from those already tested and used in previous works. Namely,\nFlow matching (FM) objective and Detailed balance (DB) objective. Both are inefficient in credit propagation across\nlong action sequences.\n\nThe proposed alternative is called \"trajectory balance\", as the title indicates. Basically, this approach relies on the \"trajectory balance constraint\", which in fact balances the Markovian flow F (and its forward and backward policies) over complete trajectories. \n\nA set of experiments to test the performance of TB is provided in Section 5. More specifically: (1) The use of an Hypergrid environment to understand better the behavior with respect to MCMC and PPO, and more interestingly when using GFlowNets with FM and DB. (2) Experiment on Small drug molecule synthesis where TB, FM and DB are compared. (3) Autoregressive sequence generation in (3.1) bit sequences' problem and (3.2) Anti-Microbial Peptide (AMP) generation. \n", "questions": "1- One evident question would be: why publish this idea now? I mean, it results at least surprising that it has been already used in papers already published. I see all of them belong to 2022, so it may be for distinct reasons, but I feel curious about the chronological order in publication.\n\n2- Runtime gain (numerical speedup) is only explicitly commented (x5) in the drug molecule synthesis problem. I assume this is because it is more significant there. However, how does it affect the other problems? Is there any quantitative estimation of the general profit with respect to FM and DB? What aspects/factors this improvement will depend on?", "limitations": "NA", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Positive aspects:\n\n* Well presented and original idea (within GFlowNets' world)\n* Good quality and clarity, even though in the experimentation part there is little space for further analysis\n* Already successfully applied (as indicated in the last paragraph of Section 1)\n\nLess positive aspect:\n\n* It could be a subjective point, but the contribution is quite specific to be applied in a relatively new \"tool/model\" and solving a quite particular task on it.", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "4WlDqIlBbfA", "reviewer": "Reviewer_fSaD", "summary": "\nTrajectory balance: Improved credit assignment in GFlowNets\nIn this paper, the authors improve the original generative flow networks’ objective function by modeling it on top of the entire trajectories generated from the policy network. This new objective function improves the convergence, diversity, robustness etc. compared to existing methods.\n", "questions": "\n1) I am entirely sure why the proposed algorithm improves “credit assignment”. Maybe the authors will elaborate on this? I do see improved results in terms of the correlations between final reward and policy probability across different tested environments & tasks. \nBut why does it improve the credit assignment? Maybe it improved something else instead?\nIs there a metric or discussion to support this claim?\n\n2) Does the algorithm scale to high-dimensional problems? A lot of the robotics tasks also involve actions that incrementally modify the state of the robot to solve the tasks (in gym for example). How well does the proposed algorithms work with these problems?\nAlso does the algorithm work with problems that have continuous state space?\n\n3) The fonts of the figures are way too small (for example Figure 2, Figure 4), and it is hard to read without going back to the pdf version.\n\n", "limitations": "1) As far as I know, I don’t think it will have any potential negative societal impact directly.\n2) In terms of the limitations of the paper itself, I think it would be great for the authors to comment on how well the algorithm will scale for high-dimensional robotics problems. As I have no related knowledge, I am not entirely sure how difficult the tasks of drug molecule synthesis and AMP generation are and how well the baselines should perform.\n\n", "rating": 7, "confidence": 1, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nUnfortunately I don’t have any experience or knowledge working with generative flow networks, and therefore I can only provide my evaluation and feedback as an outsider.\n\n1) Originality: The paper discusses a relatively new technique of generative flow network and proposes a novel formulation of objective functions.\n\n2) Quality & Clarity: The paper is very well written. The preliminary section provides complete and valuable background knowledge for readers not familiar with the topic. \nThe paper is supported both with solid math and well designed experiments. The results of the experiments are also fully discussed.\n\nThe code is given, and the results can be easily reproduced.\nThe comments and remarks are interesting and inspiring. The related work is very complete and adequate, and the connections to some other algorithms or methods are discussed.\n\n3) Significance: In this paper, the generative flow networks are extended and potentially can be used in more complicated real-life problems. Although it is not clear how well the method can be scaled to high-dimensional continuous state problems such as robot arm target reaching tasks.\n\n", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept