Learning GFlowNets from partial episodes for improved convergence and stability
Authors: Kanika Madan
Jarrid Rector-Brooks
Maksym Korablyov
Emmanuel Bengio
Moksh Jain
Andrei Cristian Nica
Tom Bosc
Yoshua Bengio
Nikolay Malkin
Venue: Submitted to ICLR 2023 (2023)
Keywords: GFlowNets
probabilistic modeling
reinforcement learning

Abstract: Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks. Existing training objectives for GFlowNets are either local to states or transitions, or propagate a reward signal over an entire sampling trajectory. We argue that these alternatives represent opposite ends of a gradient bias-variance tradeoff and propose a way to exploit this tradeoff to mitigate its harmful effects. Inspired by the TD($\lambda$) algorithm in reinforcement learning, we introduce subtrajectory balance or SubTB($\lambda$), a GFlowNet training objective that can learn from partial action subsequences of varying lengths. We show that SubTB($\lambda$) accelerates sampler convergence in previously studied and new environments and enables training GFlowNets in environments with longer action sequences and sparser reward landscapes than what was possible before. We also perform a comparative analysis of stochastic gradient dynamics, shedding light on the bias-variance tradeoff in GFlowNet training and the advantages of subtrajectory balance.

LEARNING GFLOWNETS FROM PARTIAL EPISODES
FOR IMPROVED CONVERGENCE AND STABILITY
Anonymous authors Paper under double-blind review
ABSTRACT
Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks. Existing training objectives for GelowNets are either local to states or transitions, or prop- agate reward signal over entire sampling trajectory. We argue that these alternatives represent opposite ends of gradient bias-variance tradeoff and pro- pose a way to exploit this tradeoff to mitigate its harmful effects. Inspired by the TD(A) algorithm in reinforcement learning, we introduce subtrajectory bal- ance or Su) a GFlowNet training objective that can learn from partial ac- tion subsequences of varying lengths. We show that SubTB(2) accelerates sam- pler convergence in previously studied and new environments and enables train- ing GFlowNets in environments with longer action sequences and sparser reward landscapes than what was possible before. We also perform a comparative anal- ysis of stochastic gradient dynamics, shedding light on the bias-variance tradeoff in GelowNet training and the advantages of subtrajectory balance.
INTRODUCTION
Generative flow networks (GFlowNets: Bengio et al. 2021a) are generative models that construct objects lying in a target space X by taking sequences of actions sampled from learned policy. GFlowNets are trained so as to make the probability of sampling an object X proportional to given nonnegative reward R(x). GFlowNets' use of a parametric policy that can generalize to states not seen during training makes them competitive alternative methods based local exploration in various probabilistic modeling tasks (Bengio et al. 2021a Malkin et al. 2022 Zhang et al. 2022]
Jain et al. 2022 Deleu et al. 2022).
GFlowNets solve the variational inference problem of approximating a target distribution over X with the distribution induced by the sampling policy, and they are trained by algorithms reminiscent of reinforcement learning (although GFlowNets model the diversity present in the reward distribu-
tion, rather than maximizing reward by seeking its mode). In most past works (Bengio et al. 2021a Malkin et al.] 2022 Zhang et al. 2022, Jain et al., (022). GFlowNets are trained by exploratory sampling from the policy and receive their training signal from the reward of the sampled object.
The flow matching (FM) and detailed balance (DB) learning objectives for GFlowNets proposed in
Bengio et al. (2021ab) b) resemble temporal difference learning (Sutton & Barto 2018).
A third objective, trajectory balance (TB), was proposed in Malkin et al. (2022) to address the problem of slow temporal credit assignment with the FM and DB objectives. The TB objective propagates learning signals over entire episodes, while the temporal difference-like objectives (FM and DB) make updates local to states or actions. It has been hypothesized by Malkin et (2022) that the improved credit assignment with TB comes at the cost of higher gradient variance, analogous to the bias-variance tradeoff seen in temporal difference learning (TD(n) or TD(2)) with different eligibility trace schemes (Sutton & Barto 2018 Kearns & Singh 2000 van Hasselt et al. 2018 Bengio et al. 2020). This hypothesis one the starting points for the present paper.
In this paper. we propose a new learning objective for GFlowNets, called subtrajectory balance (SubTB, or (.)) when its real-valued hyperparameter A is specified). Building upon theoretical results of Bengio al. (202Tb); MMakinn et al. (2022), we show how the Su)TB(2) objective allows the flcxibility of learning from partial experiences of any length. Experiments on two synthetic and four real-world domains support the following empirical claims:

---

(1) SubTB()) improves convergence of GFlowNets in previously studied environments: models trained with SubTB(a) approach the target distribution in fewer training iterations and are less sensitive to hyperparameter choices.
SubTB()) enables training of GFlowNets environments where past approaches perform poorly due to sparsity of the reward function or length of action sequences.
(3) The benefits of SubTB(A) are explained by lower variance the stochastic gradient, with the parameter 1 allowing interpolation between the high-bias, low-variance DB objective and the low-bias, high-variance TB objective.
METHOD
2.1 PRELIMINARIES
In this section, we summarize the necessary preliminaries GFlowNets. We follow the notation of Malkin et al. (0022) to which the reader directed for a more thorough exposition written with a view towards motivating the trajectory and subtrajectory balance objectives. A deeper introduction given in Bengio et al. (20TTb).
Let G A) be a directed acyclic graph. The vertices s E are called states and the directed edges (u-vv E A are actions. If (u~v) is an edge, we say is a child of and is parent of v. There is unique initial state so S with no parents. States with no children are called terminal, and the set of terminal states is denoted by X.
A trajectory
action sequence is sequence of states 35mmSm++1 - ... Sn). where each (511);+41 is an action. The trajectory is complete if sm = so and sn is terminal. The set complete trajectories is denoted by T.
A (forward) policy is a collection of distributions (-(- over the children of every nonterminal state $ S. forward policy determines a distribution over T by
n-1
PF(T = (so- = [ Pr Ps(si))s)). 1=0
(I)
Any distribution over complete trajectories that arises from forward policy satisfies Markov prop- erty: the marginal choice of action out state s is independent of how s was reached. Conversely, any Markovian distribution over T arises from a forward policy (Bengio et al. 2021b).
forward policy can thus be used sample terminal states x E X by starting at so and iteratively sampling actions from Pr, or, equivalently, taking the terminating state of a complete trajectory (F))). The marginal likelihood of sampling X is the sum of likelihoods of all complete trajectories that terminate at x.
Suppose that a nontrivial (not identically 0) nonnegative reward function R X - Rg0 is given. The learning problem solved by Nets is estimate policy PF such that the likelihood of sampling x X is proportional to R(x). That is, there should exist constant such that
R(x) Z
Σ
PF(T)
Vx
(2)
T=(50 75.
If 2 is satisfied, then Z= Σxex R(x).
2.2
GEL TRAINING OBJECTIVES
Because the sum in P may be intractable to compute, is in general not possible directly convert this constraint into a training objective. To solve this problem, GFlowNet training objectives intro- duce auxiliary variables in the parametrization in various ways, but all have the property that (2) is satisfied at the global optimum. The key properties of these objectives are summarized in Table
Flow matching (FM; Bengio et al., 2021a). Motivating the 'flow network' terminology, Bengio al. (2021a) proved that (2) is satistied Pr arises from an edge flow function satisfying certain constraints. Namely, an assignment F A + Rg0 of a nonnegative number (flow) to each action defines a policy via
(r(()) = Er:ss-r) F(s-r\)
A sufficient condition for the terminating distribution of Pp to be proportional to the reward R(x) is that family of fow-matching (flow in flow out) conditions satisfied at all interior states and
2

---

Table 1: Summary of GFlowNet training objectives.
Objective
Parametrization
Locality
Flow matching
odge flow F(s-t;
state
Detailed balance
state flow F(s; 0), policies Pr(-l-:0), (p(-)))pppp330000
action s->t
Trajectory balance
initial state flow Zo. policies Pp( P-(01-00))
complete trajectory r
Subtrajectory balance
state flow F(s:0) policies (-((--:0),
(partial) trajectory
family of reward- reward-matching matching conditions is satisfied terminal states:
Σ s: Prr
F(s-r)=
Σ
F(I-u)
V WES (so)),
Σ
F(s-x) R(x)
(4)
The flow F(s-t) is then proportional to the marginal likelihood that complete trajectory sampled from Pp includes the action 5->1.
In Bengie et al. (2211), , a GFlowNet described by a parametric estimate of the edge flow function, F(u-v;;0) 8) (a neural net with parameters 0). These conditions can be converted into objectives that minimized when (4) is satisfied. For example, the flow-matching objective at nonterminal state s is defined by
(FM(s) = [tog Ex:(--EE F(s-1; +
(5)
Em:rr-aie> F(--u; 0) +
where e a smoothing constant that can safely be set to the flows are constrained be strictly positive, and a similar objective (or constraint by construction) is defined to force the flow F(s-x) into terminal states match R(x). If these objectives
globally minimized for all states s, then
the policy (r(-1-:0) 0) defined by F(-;0) via (3) satisfies (2. with Z -m::s0-))e F(s-1;0)
ΣxeX Rxx). The question of how to sample states s for training is discussed below.
Detailed balance (DB; Bengio et al. 2021b; Malkin et al. 2022). In the DB parametrization, forward policy model Pr P- 0) is learned directly, jointly with two additional objects: backward policy model P -l) which can predict a distribution over the parents any noninitial state, and state flow function F(s; 0) (typically parametrized in the log domain). The detailed balance conditions state that
F) 0) F(r; P(;:0))((s)::0
(6)
for all actions (s-r) and F(x; 0) R(x) for terminal. Satisfaction of these conditions for all actions (s-x) and x E X implies that Py samples proportionally to the reward (i.c., satisfies (2). with F(so)). The DB condition
(6) can be converted into a squared log-ratio objective LD (s-) in
the same way that (4) yields (5) ) and (DB(s-1)) can be optimized over sampled actions (s (s-t).
Trajectory balance (TB; Malkin et al. 2022). The parametrization required for the TB objective includes forward and backward policy models P-(-1-:0) and P=(-1=;0), as well as an estimate Ze of the constant of proportionality in (2). Satisfaction of the following condition for all complete trajectories (50 sp) implies that (2) is satisfied:
ZoP (T; 0) P
where we have used the conventions
8-1
PF(T:0)
[ Pa (=(s)1s;;)) i=0
Pb(TIs,;
[ P Pp(si(smi:)) 1=0
The condition 7 can again be made into squared log-ratio objective LTB(T) and optimized for complete trajectories т taken from some training policy. In Malkin et al. (2022), the TB objective was empirically demonstrated to have better convergence properties than FM and DB on various problem domains.
Training policy and exploration. Global minimization of the FM, DB, and TB objectives for all values of their respective arguments (states, actions, or complete trajectories) implies satisfaction of 2). Therefore, given sufficiently expressive model and convergence of the optimization procedure, a GFlowNet policy that samples x with likelihood proportional to R(x) can be trained by minimizing any of these losses distribution with full support, enabling offline training of GFlowNets. As
3

---

in other RL algorithms, the distribution over sampled states, actions, episodes can be fixed and off-policy, or can vary over the course of training and use available information about terminal states in interesting ways (Zhang et al., 2022 Deleu et al. 2022). The simplest approach, which is also taken in this paper, on-policy learning or very similar off-policy variant that flattens the current policy to ensure exploration. Complete trajectories T = (sp-.... (so=....-s,) are sampled from the forward policy Px(-1-:0) (tempered or mixed with uniform policy with small weight so as to ensure full support and exploration). One then takes gradient descent steps on LTB(T), on (DB(S)-S(A1) over all actions in r, or on (sM(s;) for all intermediate states in T.
The GelowNets in this paper trained on-policy, or off-policy with training policy that is mixture of PF with a uniform policy: T = (50 -5p) is sampled with 5/41 ~ (1 ( [s;; 0) + *) Here is the random exploration weight.
2.3 SUBTRAJECTORY BALANCE: LEARNING FROM PARTIAL EPISODES
Recall the GFlowNet parametrization used in the DB objective above, with a state flow estimator
F(-I-; 9) and a pair of policies Pr(-l-;0) 9), (n(-I=;
It is shown in $A.2 of Malkin et al. (2022)
that the detailed balance conditions (6) are satisfied for all actions if and only if the following subtra-
jectory balance constraint holds for all (not necessarily complete) trajectories T = (sm-.... sn):
[ Pr (serls;;;) = ((sn:)) 0)
I PB (8)
F(sm: 0)
(8)
where we again enforce that F(x; = R(x) if x is terminal. Observe that the DB condition (6) is a
special case of (8) when the trajectory consists of one action, and the TB condition 7 is precisely the case when r is complete, with the identification Zo = F)s;
The above constraint yields the subtrajectory balance objective
LSuBTB (T) = log F(Sm' 0) I Pr(siulss;;0)
(9)
F(sni 0) 1- Pa(si
If this objective is made equal to O for all partial trajectories т, where Rss,) is substituted for F(xni 0) if sn is terminal, then the policy PF satisfies the desired condition ( (Proof: When LSubTB(() T) = 0, (8) is satisfied, implying satisfaction of both (7) and (6). Either of these conditions is a sufficient condition for 2 as shown by Bengio et al. (2021b); Malkin et al. (2022
Extracting subtrajectories for training. Suppose that episode (complete trajectory) (5o bo..)_...., is sampled for training. There are (*=²)) 0(n?) nontrivial subtrajectories:
Trjj := (51 ss)). 0si<ja
(10)
Having sampled a complete trajectory 7 for training, we make gradient steps on a convex combina- tion of the subtrajectory balance losses LS (T:;): 0 0 VeL, where
Σngiejan 1 LSuNTB(TE.))
(11)
Eosiejsn 2i11
Here A O is a hyperparameter controlling the weights assigned to subtrajectories of different lengths, and when is set leads to uniform weighting scheme. Notice that the - 0* limit leads precisely to the average detailed balance loss LDS (S; +1+1 over all transitions in r, while the
- 400 limit gives the trajectory balance objective LTm (T)
Other schemes for weighting subtrajectories are possible and should be explored in future work.
Computational considerations. It may appear that the optimization of (11) induces a computation cost that is quadratic in the trajectory length. However, a closer inspection of the gradient of (TT) with respect to the state flows log F F(;;; 0) and the forward and backward policy logits shows that gradient computation requires only one forward and one backward pass through the neural networks giving log Fs;; 0), log (=(-1a;; (). and log PB (- 0). The quadratic computation cost is incurred only in performing linear operations on these log- flows and policy logits, not in the evaluation of the deep networks. Thus the SubTB loss has little computation overhead over DB or TB.
I When batch of trajectories is used for training, the convex combination weights may either be normalized over all subtrajectories of all trajectories in the batch, or normalized independently over the subtrajectories of each trajectory. For consistency, we choose the first option for the experiments in this paper.
4

---

is regressed to Hypothesized benefits. We hypothesize that SubTB(1) brings two benefits to GFlowNet training:
VARIANCE REDUCTION. The TB loss terms LTB (T) for trajectories T that take a given sequence of actions until state s, then diverge. share the terms log Z and the policy logits for all transitions preceding s inside the square. However, the 'tail' of the TB loss, involving the forward and back- ward policy logits for transitions that appear after s in T, can be seen as a stochastic least-squares regression target. That is, if s = in trajectory (50-5))
sn), then
Pr (s(s))))))
log Z
(12)
P (s(s
i=0
log
R(sn)
(13)
Pf(sia
i=
Similarly, for trajectories that share the transitions following s but may differ in their initial actions, is stochastic regression target for (13)
The subtrajectory balance loss terms LSubTB (myy) for partial trajectories beginning at s regress the log-state flow log F(s) to (parts of) expressions like (13). while loss terms LSabTB ( (Tem) regress (parts of) expressions like (12) to the log-state flow log F(s). The learned log F(s) thus learned estimate of a stochastic piece of the TB loss for trajectories that contain s. Replacing a stochastic term in the TB loss by a learned estimate of its expectation is guaranteed to introduce bias into the gradient (with respect to the gradient of the TB loss), but expected to reduce variance. This is akin to the variance-reducing effect of actor-critic methods in RL.
This hypothesis is studied empirically in our experiments and in particular $4.11 where we provide evidence that SubTB(2) is a practically useful interpolation between TB (high variance) and DB (low variance, high bias relative to the true TB gradient) losses.
FASTER LEARNING DUE TO GENERALIZATION OF STATE FLOWS. Another benefit of subtrajec- tory balance for convergence speed may come from the ability of estimated state flow functions log F(s; f) to be modeled with high precision and generalize between states faster than the often high-dimensional policy logits log (-|[=s; 0), log (= 8). Such generalization is important in problems where the state graph becomes 'wide' far from the initial state, making the learning signal sparse at states that are near termination. Indeed, in all of our experiment domains except the hy- pergrids in $4.1 and for the largest hypergrids the number of terminal states is many orders of magnitude larger than the total number of states seen in training.
RELATED WORK
Eligibility traces. SubTB(2) draws inspiration from the TD(1) algorithm in RL (Sutonn 1988
Sutton & Barto 2018), which forms an estimate of the expected return via convex combination of n-step returns, each weighed by (1- (1-1) d) The parameter A E [0, 1] enables a bias-variance tradeoff (Kearns& & Singh 2000). Intuitively, larger à leads to lower bias and higher variance, since the estimate of the expected return approaches the single-point Monte Carlo estimate as A - 1. We take inspiration from this idea to mix together different (possibly all) subtrajectories, akin to how n-step returns are mixed together. We hypothesize that the right mixing may reduce variance, compared to TB, with the additional benefits of inducing consistency between the flows of interme- diate states, and thus of helping propagate credit faster and enable faster convergence. In addition, GFlowNet training objectives are reminiscent of residual gradient RL methods (Baird) 1995 Zhang et al. 2020) since the "endpoint" (e.g. F F(s,)) in (9) is also considered in the gradient.
maximum entropy, methods (Ziebart 2010 Mnih et al. 2016 Haarnoja et al.) 2017 / Nachum et al. 2017]] Schulman et al. [017 HHarnoja ct al. 2018), which are close equivalent to the GFlowNe Net framework in certain settings (in par- ticular when the MDP has a tree structure (Bengio et al.) (021a)). Also related are methods that maximize entropy not on the policy, but rather on the state visitation distribution (Hazan et al. 2019 [slam et al.) 2019 Zhang et al.) 2021) some proxy of it (Eysenbach et al.) 2018). which achieve a similar objective to GFlowNer Net models by flattening the state visitation distribution. the state graph of the environment is directed tree, the loss LSubTB on individual subtrajectories equivalent to that of path consistency learning (Nachum et al. (017).. However, attempts to use path consistency learning in settings without intermediate rewards have only computed the loss on subtrajectories that have length 1 or include terminal state (Guo al. 2021)
MaxEnt RL. RL has rich literature on energy-based,
5

---

16x16
16x1Ex16316 20
15
2.5
2.0
0.0
10f
107
10'10"
107
10'10
101
107
10
107
10"
10'101
10"
10
28
28
32 32
40
50
60 60
23 32
2.3
3.6
10"
10'00
10
10 1042" 10s
110
8 8 (hard)
12 (hard)
16 (hard)
20x 20 (hard)
24 x4 (hard)
28 thard)
32 (hard)
1.5
3 1.0
10"
10"
16a
10
10
30410
10"
30t
10
10
10"
10
10t
121
10900
10
36 10t
10.00
10"
Figure 3: L distance between empirical and target distributions the hypergrid environment. SubTB(I = 0.9) consistently gives faster convergence than TB, the strongest objective from past work, on all grid sizes. The difference is especially visible for the harder variant of the reward function (last row). The x-axis the cumulative number of training trajectories (episodes).
the course of training
EXPERIMENTS
4.1
HYPERGRID: ROBUSTNESS TO SPARSE REWARDS
We study the synthetic hypergrid environment introduced in Bengio et al. (202la). The set of interior states is d-dimensional hypergrid of H H x ... x H with a multimodal reward function concentrated near each of the 2d corners of the hypergrid (sce Bengio et al. (202 Malkin al. (2022) and Fig. [ The initial state is (0,0, ... O), and each action step that increments one of the d coordinates by 1 without leaving the grid. special termination action also allowed from each state. This environment is designed to challenge a learning agent to infer and discover new modes from those that have been already been visited.
Figure 1: 16 x 16 hyper- grid reward function.
We study various sizes of 2-dimensional and 4-dimensional hypergrids, using the hardest variant of the reward function from past work (the minimal reward, away from the corners of the grid, is set to 10-3). We train GFlowNets to sample from the target reward functions and plot the evolution the L' distance between the target distribution and the empirical distribution of the last 2 10 states in training In all cases, tune the learning rates for the TB and SubTB(1 0.9) objectives. (See $A for details.)
The results (mean and standard deviation over three random runs) shown in the first rows of Fig. Models trained with SubTB(d) converge faster, and with less variance between random seeds, to the true distribution than with TB for all hypergrid sizes.
25-
We also study an even sparser variant of the environment, in which the background reward is set to 10 In this case, SubTB(.) con- tinues to perform strongly (last row of Fig. 3}, while models trained with TB do not even discover all modes of the target distribution for grids larger than 8 8 (Fig. 2).
Figure 2: Distribution of x
10 samples from GelowNets trained on the harder variant of the 32 32 grid with TB and SubTB(A) objectives.
Additional results are given in I In particular, SubTB(2) con- tinues to perform strongly when only subtrajectories of less than certain length are used for training, which can be beneficial in realistic settings where only partial episodes are given. We also show the effect of l on the convergence rate (Fig. A.2 and of more exploratory training policies (Fig. A.3)
synthetic environment because the exact target distribution function be tractably computed. Note that the metric shown in Fig. 3 differs from what is called 'L! distance" in past work, do not divide by the total number of states.
2Scch an evaluation possible

---

Figure 4: Mean cosine similarity between small-batch (2 and large-batch (1024) gradients at se- lected training iterations. Left: Small-batch vs. large-batch gradients of DB, SubTB(.), and TB objectives. Right: Small-batch DB, SubTB(a), and TB gradients vs. large-batch TB gradient.
4.1.1 A CLOSER LOOK AT GRADIENT VARIANCE
We take a closer look at gradient bias and variance to understand the benefits of training GFlowNets with SubTB(2). The method- ology of these experiments is inspired by [lyas et al. (2020).
Codlee Sinkllo
Losa Cemo
0.7
9.6
SubTD SubTo
TDl ve
0.5
We train GFlowNets on the 8 x 8 grid environment using SubTB(d 0.8) and monitor various gradient metrics during training. To remove the effect of parameter sharing between policies at different states and to isolate the effect of the objec- tive, we use tabular representation of the GFlowNet, i.e., all flows and policy logits are optimized as independent parameters.
1500
Simbrity Gridlen, Learned Rows
Codle Santar
0.5
Gradient variance. To measure gradient variance, use the following procedure for each training objective (DB, TB, or SubTB())). A large batch of 20 1024 trajectories is sam- pled, and the gradient / of the objective with respect to the policy logits at all states is computed for each trajectory Tÿ in (0) the batch. Then, for each k [0, I........ 9}, the gradients 8; are combined into 210- sub-batches, each of size 24, The sub- batch gradient (*a) for the i-th sub-batch is
Lasa Camp Narms
0.1
Traning Step
Figure 5: Mean cosine simi- larity between small-batch (64) and large-batch (1024) gradients on the 8 x 8 grid environment. Above: Self-similarity of the DB, SubTB(I), and TB gradi- ents, showing DB < SubTB(a) TB in gradient variance. Be- low: Similarity of small-batch DB, SubTB(a), and TB gradi- ents to the bate-bbacch TB gradi- ent, showing that the small-bated SubTB(A) gradient is good es- timator of large-batch TB.
the average (0) contained within the sub-batch and 210-2). We then report the average
of trajectory gradients 8
computed for E cosine similarity between the sub-batch and full-batch gradients:
210-6 (A)
(10)
210-
8
If this quantity positive, then gradient steps of infinitesimalyy small norm along the stochastic bub-batch gradient decrease the full-batch objective in expectation. Fig. 4 (left) shows the dependence of this metric on k at various iterations. A steeper curve, such as those of DB and SubTB(I), indicates lower gradient variance.
Fig. 5 (top) shows the metric at k 6 (corresponding to the batch size of 64 used for training) over the course of training. We see that the DB gradient has the highest self-consistency at all iterations, TB has the lowest, and SubTB(t = 0.8) is in between.
Gradient bias. We next compare the small-batch stochastic gradients with large- batch stochastic gradients, using different objectives for the small and full batches. Specifically, we compare the small-batch DB, SubTB(,)) and TB gradients with the full-batch TB gradient. (The full-batch TB gradient can be seen
'canonical" gradient against which bias can be measured, its expectation equals the gradient of the KL divergence between the distribution over trajectories defined by PF and that defined by the reward R and PB; $A.3 of Malkin et al. (2022).)
Fig. 5 (bottom) shows the cosine similarity at the batch size used for training. Notably, at interme- diate iterations, the similarity of SubTB(2) with TB is higher than that of TB with TB: despite its bias, the small-batch SubTB(I) gradient estimates the full-batch TB gradient better than the small-batch TB gradient does. Fig. 4 (right) shows the dependence of the similarity on k at se- lected iterations and suggests that this effect may be even larger for smaller batch sizes. Moreover, at = 10, the similarity of SubTB(d) vs. TB always lies between DB vs. TB and TB vs. TB, indicating that SubTB(t) interpolates between TB's 's unbiased and DB's 's biased estimates of the TB gradient.
The effect of learned state flows. For additional experiments, see A.2

---

beggy to legitti corretatiio
02
10 leaming rate
32-
0.5
reward exponent
SubTB parameter
Figure 6: Correlation between marginal sampling log-likelihood and log-reward on the molecule task. For each hyperparameter setting on the x-axis, we plot the best result over choices of the other hyperparameter(s)) - a in the left plot, & in the centre plot, and both a and / in the right plot - with solid line. The mean result over values of other hyperparameter(s) is plotted with dashed line.
4.2 SMALL MOLECULE SYNTHESIS
We use SubTB(1) to train models on the molecule generation task of [Bengio et al. (2021a) The task is generate binders of the sEH (soluble epoxide hydrolase) protein, based a docking prediction TTrot & Olson 2010). To be precise, molecules are generated by sequentially joining 'blocks' from a fixed library to the partial molecular graph (Jin et al.) 2020 Kumar et al.) 2012). resulting in a state space of estimated size 10 The reward function R is given by pretrained proxy model made available by [Bengio et al. (2021a). To adjust the greediness of the agent, an inverse temperature hyperparameter B is used, i.c., the reward used for training R(x) R(x)² where R(x) is the proxy's 's prediction.
We train models with the DB, TB, and SubTB(A) objectives, with four values each of a, B, and learning rate, averaging the results over 3 random runs for each setting. We measure how well the trained models match the target distribution by the correlation of log R(x) and log (x the log-probability assigned to by the GFlowNet, computed
held-out set of terminal states
The results are shown in Fig. 6 SubTB(2), in particular with 2 1, performs better than both DB and TB when the optimal hyperparameters a, f are used (solid lines) and is far more robust to the choice of hyperparameters (dashed lines). Additional details can be found
4.3 SEQUENCE GENERATION
We consider three sequence generation tasks which sequences are generated left to right, with each action appending one symbol from a vocabulary to partial sequence: a synthetic task with varying sequence lengths and vocabulary sizes ($4.3.1). D. a practical biological sequence design task (44.3.2. and a new protein design task with longer sequences (4.3.3). 3. For all three tasks, we consider the baselines Soft Actor-Crriit regularization (Williams Peng [991] et al. 2016) and MARS-like MCMC (Xie et al. 2021 and compare them with three GFlowNet training objectives: TB, FM, and SubTB(.).
(Haarnoja al.) 2018 Christodoulou 2019), A2C with Entropy
In $ we also study a non-autoregressivve sequence generation problem (inverse protein folding).
4.3.1 BIT SEQUENCES
We consider the synthetic sequence generation setting from Malkin et al.) (022 where the goal is to generate sequences of bits of fixed length = 120. The reward is specified by a set of modes M C X {0. 1}" that is unknown to the learning agent. The reward of a generated sequence is defined in terms of Hamming distance from the modes: R(x) exp(- minyeM y(xxy)).
The vocabulary size can be varied: for any integer k dividing 120, we take a vocabulary consisting of words length (so that the vocabulary size is 2 and the full sequence is generated in actions). By varying the value of & and keeping п and M constant, we study the behavior of learning agents with varying action space sizes and trajectory lengths without changing the underlying modeling problem. Most experiment settings are taken from Malkin er al. (2022); see $C
Comparing the exact sampling and target distributions. like in $4.1 not possible here, since we cannot enumerate all terminal states. However, the marginal likelihood that trained GFlowNet generates given tractable to compute by dynamic programming. For model that samples perfectly from the target distribution, log( (x)) and log po would differ constant log Z independent of and thus be perfectly correlated.
8

---

04-
Q55 SelThahggppaa
Eagday bduze
CBN The Maciag
Spumes Cundaka
ABC Ernnyy
Sea Arve Orli
MONC
01-
1000
2000 20000
20000
Figure 7: Left: For the number of bits & E 21, 4, 10] in each vocabulary token, we plot the Spearman correlation between the sampling probability and reward on test set for each method. Training with SubTB(A) leads to policies that have the highest correlation with the reward across all lengths and vocabulary sizes. Right: For k = I, the number modes discovered by each method over the course of training is plotted. SubTB()) discovers more modes faster.
Models are evaluated by computing the Spearman correlation, on a test set of sequences x, between the probability of generating x and the reward R(x). We also track the number of modes discovered during the training process for all the methods, see Fig. We find that models trained with the SubTB(d) objective have a higher Spearman correlation the end of training and discover modes faster compared to the other GFlowNet objectives and non- Non-GFlowNeet baselines.
4.3.2 ANTIMICROBIAL PEPTIDE GENERATION
Table 2: Results on the AMP generation task (mean and standard error over 3 runs).
Next, we consider the task of generating peptides with antimicrobial properties (AMPs). These quences have maximum length 60 and use a vocab- ulary of 20 amino acids (and an end-of-sequence to-
Algorithm
Top- 100 Reward Top-100 Diversity
GEN- (SN-LSOBBA)
0.96 = 0.02
42.23 = 3,4
GFN-LIN
0,90 0.03
31.42
GEN- L
0.78 = 0.05
12.61 # 1.32
ken), resulting in a state space of size 21 The re- ward function is pretrained proxy neural network that estimates the antimicrobial activity. (See Jain et al. (2022) for details on this task.)
SAC
0.80 0.01
8.36 1.44
AAC-ER
0.79 = 0.02
7.32 0.76
MCMC
12.56 1.45
We train GFlowNets with the SubTB(J). TB, and FM losses and compare them with baselines. To evaluate the trained models, we sample 2048 sequences from the policy, then compute the mean re- ward and mean pairwise edit distance of the top- 100 reward sequences. The metrics and model archi- tecture are taken from Malkin et al. (2022); see $D The results are presented in Table 2 SubTB(2) provides significant improvements over all the baselines (including TB, FM, and DB GFlowNets) in both reward and diversity.
4.3.3 FLUORESCENT PROTEIN GENERATION
Table 3: Results on the GFP generation task
(mean and standard error over 3 runs).
We consider the task of generating protein sequences with fluorescence properties (Trabucco et al. 2022)
Algorithm
Top- 100 Reward Top-100 Diversity
to evaluate Su) in settings with longer trajec- tories. In this task, sequences have a fixed length of 237, and the size of the state space is 20237 The proxy reward function R(x) is trained on a dataset of proteins with their fluorescence scores from Sark- isyan et (2016). The metrics and models are the same as in $4.3.2; see $ for details.
GEN LSUSTTDD
1.18 0.10
201.44 0..5
GFN- LTH
0.76 = 0.19
204.31 = 0.44
GFN- Lna
0.30 0.08
190.21 6.78
SAC
0.23 0.03
120.32 # 15.57
AAC-ER
0.22 0.02
113.65 21.31
MCMC
0.28 0.01
169.17 + 12.44
The GFlowNet objectives outperform all other methods in both metrics, finding more diverse and higher-reward sequences (Table 3.. SubTB(I) significantly outperforms TB, while achieving sim- ilar diversity. We note that the advantage of SubTB(2) is greater than that in the AMP task (Table @ and speculate that the benefits of SubTBB() become more prominent for longer action sequences.
DISCUSSION AND CONCLUSION
We have given evidence of a bias-variance tradeoff in GFlowNet training algorithms. The high- variance stochastic regression objective of TB and the low-variance local consistency objective of DB lie at opposite ends of this range. We showed that SubTB()) can harness the variance-reducing effects local objectives while retaining the fast credit assignment properties of trajectory- objectives. We see leurnable strategies for selecting and weighting (sub)trajectories for training e.g., a dynamic choice of 1 and an active-learming approach to sampling trajectories as the most teresting questions for future work. The ability of subtrajectory objectives to learn from incomplete episodes also makes their application in RI. environments an appealing research direction.

---

REPRODUCIBILITY STATEMENT
We provide extensive experiment details, such as learning rates, batch sizes, number of training steps, choices of A, description of attempted hyperparameters, and additional clarifying experiments the Appendices. Code for experiments on the hypergrid domain ($4.1) and on the molecule is also provided with the submission.
domain 4.2
REFERENCES
Leemon Baird. Residual algorithms: Reinforcement learning with function approximation. Inter- national Conference on Machine Learning (ICML), 1995.
Emmanuel Bengio, Joelle Pineau, and Doina Precup. Interference and generalization temporal difference learning. International Conference Machine Learning (ICML), 2020.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow net- work based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS). 2021a.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. arXiv preprint 2111.09266, 2021b.
David Brookes, Hahnbeom Park, and Jennifer Listgarten. Conditioning by adaptive sampling for robust design. International Conference on Machine Learning (ICML), 2019.
Sidhartha Chaudhury, Sergey Lyskov, and Jeffrey J Gray. PyRosetta: a script-based interface for im- plementing molecular modeling algorithms using Rosetta. Bioinformatics, 26(5):689-691, 2010.
Petros Christodoulou. Soft actor-critic for discrete action settings. arXiv preprint 1910.07207, 2019.
Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julics Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity all you need: Learning skills without reward function. International Conference on Learning Representations (ICLR), 2018
Han Guo, Bowen Tan, Zhengzhong Liu, Eric P. Xing, and Zhiting Hu. Text generation with efficient (soft) Q-learning. arXiv preprint 7770007704 2021.
Tuomas Haarnoja. Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. International Conference on Machine Learning (ICML), 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. International Conference on Machine Learning (ICMI), 2018
Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. International Conference on Machine Learning (ICML), 2019.
Andrew Ilyas, Logan Engstrom, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. A closer look at deep policy gradients. International Con- ference on Learning Representations (ICLR), 2020.
Riashat Islam, Zafarali Ahmed, and Doina Precup. Marginalized state distribution entropy regular- ization in policy optimization. arXiv preprint 1912.05/28, 2019.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Gareia. Jarrid Rector-Brooks, Bonaventure F.P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological sequence design with GFlowNets. Inter- national Conference on Machine Learning (ICML), 2022.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Chapter 11. junction tree variational autoen- coder for molecular graph generation. Drug Discovery, pp. 228 228-249, 2020. ISSN 2041-3211.
10

---

Michael J Kearns and Satinder P Singh. Bias-variance error bounds for temporal difference updates. Conference on Learning Theory (COLT), 2000.
Ashutosh Kumar, Amout Voet, and Kam Y.J. Zhang. Fragment based drug design: from experimen- tal to computational approaches. Current medicinal chemistry, (100)):128-5147,, 2012.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems (NeurIPS), 2022.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. Neural Information Processing Systems (NIPS), 2016.
Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap between value and policy based reinforcement learning. Neural Information Processing Systems (NIPS), 2017.
Malak Pirtskhalava, Anthony A Amstrong, Maia Grigolava, Mindia Chubinidze, Evgenia Alim- barashvili, Boris Vishnepolsky, Andrei Gabrielian, Alex Rosenthal, Darrell E Hurt, and Michael Tartakovsky. Dbaasp v3: Database of amtimicrobial/yyotoxii activity and structure of peptides as resource for development of new therapeutics. Nucleic Acids Research, 40(D1):D288-D297 2021.
Carol A Rohl, Charlic EM Strauss, Kira MS Misura, and David Baker. Protein structure prediction using Rosetta. In Methods in enzymology, volume 383, pp. 66-93, Elsevier, 2004.
Karen S Sarkisyan, Dmitry A Bolotin, Margarita V Meer, Dinara R Usmanova, Alexander S Mishin, George V Sharonov, Dmitry N Ivankov, Nina Bozhanova, Mikhail Baranov, Onuralp Soyle- mez, et al. Local fitness landscape of the green fluorescent protein. Nature, 533(7603):397-440, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1707.06347, 2017.
Sam Sinai, Richard Wang, Alexander Whatley, Stewart Slocum, Elina Locane, and Eric Kelsic. AdaLead: simple and robust adaptive greedy search algorithm for sequence design. arXiv preprint 2010.02141, 2020.
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3 (1):9-44, 1988.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT Press, 2018.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench: Benchmarks for data-driven offline model-based optimization. International Conference on Machine Learning (ICML), 2022.
Oleg Trott and Arthur Olson. AutoDock Vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, 3122):55-461, 2010.
Hado van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Mo- dayil. Deep reinforcement learning and the deadly triad. arXiv preprint 1812.02648, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Neural Information Processing Systems (NIPS), 2017.
Ronald Williams and Jing Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science, 3:241-, 09 1991. doi: 11111800095540999999999969
Yutong Xie. Chence Shi, Hao Zhou, Yuwei Yang, Weinan Zhang, Yong Yu, and Lei Li. MARS: Markov molecular sampling for multi-objective drug discovery. International Conference on Learning Representations (ICI.), 2021.
11

---

Chuheng Zhang, Yuanying Cai, Longbo Huang, and Jian Li. Exploration by maximizing Rényi entropy for reward-free RL framework. Association for the Advancement of A'Aifficial Intelligence (AAAI), 2021.
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. International Conference on Machine Learning (ICML), 2022.
Shangtong Zhang, Wendelin Bochmer, and Shimon Whiteson. Deep residual reinforcement learn- ing. Autonomous Agents and Multi-Agent Systems (AAMAS), 2020.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mellon University, 2010.
12

---

1x0x0x3
12-122022
20 20
20
20x:23:<200 20 32 32 x3
30 60.00-4
8 8
12 12
16 x 16
20 20
24 24
150
400
600
60
200
300
400
200
50
100
100
0
10a
10a
10a
10
10"
10a
109
10a
10a
10a
x
32 32
50
60 6060
1.000
1500
3000
600
400
500
2000
200
500
1000
10P
10 10a
10P
10a
10a
The evolution of the L between empirical sampling and target distributions the harder variants of 4-dimensional grids, in the same format as Fig. 3 Below: The number cumulative distinct terminal states visited function of training time on the standard 2-dimensional grid. Models trained with SubTB(.)) discover more states faster.
Figure A.1: Additional results for hypergrid experiments. Above:
A EXPERIMENT DETAILS: HYPERGRID
The environment is identical to that in Malkin al. 2022 with reward function parameters (Ro. R1, R₂) (10-3, 0.5, 2) for the standard variant of the grid and (10-² 1.0, 3.0) for the harder variant. The models giving logits of (=((ls) and PB (-)),, well log F(s), MIPs of the same architecture as in Bengio et al. (022a)), taking a one-hot representation of the coordinates of s as input and sharing all layers except the last. The initial state flow log Z log F(so) is an independent parameter whose learning rate
10x the learning rate other parameters.
All models are trained with the Adam optimizer and batch size of 16 for total of 10 trajectories (62500 batches). The optimal learning rate for each experiment is chosen from (0.0005, 0.00075, 0.001, 0.003, 0.005, 0.0075, 0.01}. and a = 0.9 is chosen as the optimal value from the set {0.8, 0.9.0.99).
Gradient bias and variance experiments are conducted in the harder variant of the 8 8 grid. The tabular GelowNet is trained using Adam with a learning rate 0.007 and the SubTB(2 0.8) objec- tive.
A.1 ADDITIONAL EXPERIMENTS
Fig. A.I shows additional results on more difficult grid environments.
We perform another experiment in which only short (up to length 4) subtrajectories are used for training with the SubTB(I) objective (i.e., the sum in (II) truncated to exclude pairs (i. with j-i> 4). The results, shown in Fig. A.4 show that SubTB(2) continues to perform strongly in this restricted setting.
13

---

interpolation (hard
=
2.00
2.00
2.00
1.75
1.25
1.25
1.25
1.25
1.00
1.00
1.00
1000
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
10
10'
10>
10
101
10*
8.
8 0.9
8 A 1.3
1.50
1.50
1.50
1.50
1.2
1.25
1.25
1.25
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
10'
10'
30'
10'
10'
10'
10'
10'
8 8x..AA3.00
8 8, m 5.0
8. - 10.0
8 A 15.00
1.75
1.75
1.75
1.75
150
1.50
1.50
1.50
0.75
0.75
0.75
0.75
0.50
0.50
0.50
0.50
10
10'
10'
10'
8. 20.0
8 A 30.0
x A 50.0
8
100.0
2.00
2.00
2.00
2:00
1.75
1.75
1.75
1.75
1.50
1.00
1.00
1.00
1.00
0.75
0.75
0.75
0.75
0.50
SuLTBIA
0.00
0.00
0.00
0.00
10'
Figure A.2: Empirical L curves on the 8 x 8 grid for varying values of d.
Fig. A.2 shows the effect of the SubTB parameter 1 on the training curves, showing gradual interpolation between DB and TB and fastest convergence at values slightly less than 1.
Fig. A.3 contains visualizations of the exploration behavior of different training algorithms. It shows that TB can perform better with off-policy training and can benefit from higher temperature of the policy logits, but still does not learn as fast as SubTB(1), nor does it find all the modes in the maximum number of training iterations.
A.2 MORE ON BIAS AND VARIANCE: THE EFFECT OF LEARNED STATE FLOWS
To better understand the variance-reducing properties of SubTB(d), perform the gradient bias experiments with modified computation of gradients that removes the factor of learning the state flows.
Recall from $2.1 that a forward policy PF uniquely determines distribution over trajectories. If the initial state flow Z and forward policy Py. are fixed, there is unique state flow function FF and backward policy Pn that satisfy the detailed balance conditions (6). This 'true forward' flow function, written FF (s) 7. Erwer PF(T), is determined by an initial state flow fixed to the true
14

---

graph Tuaerr wit buiring pally
Thanirg untamreetoon actin siobebriless
2AR
$6.00
$.401
21227 21.23 21.41
Figure A.3: Training GFlowNets the harder variants of 2-dimensional grids using tempered training policy (left), and a training policy that takes uniformly random action with probability at each sampling step (right).
12 x-22
40 x.40
SLET2g-21ahat
1.3
23
124
33
10"
23
31
23
30
Figure A.4: Training GFlowNets using only short subtrajectories in different grid environments using the SubTB(1 = 0.9) objective.
partition function Exxx R(x) and the learned forward policy Pv. Similarly, the 'true backward' flow function, written ()) ETwer PB(T) where x, is the terminal state of т, is uniquely determined by the reward function R and the learned backward policy Pg. In particular, B (sp)) Σxex R(x).
We repeat the experiments on gradient bias, but by replacing the learned state flows F in the losses by either the true forward or the true backward state flows (FF or FB respectively) computed exactly using the current values of the learned Pf and Pg. (These modifications are not applied in training, but used only compute the gradient similarities. The small size of the environment makes computation of the true state flows tractable; this is not possible general.)
The gradient similarity over the course of training is shown in Fig. A.5 (cf. Fig. 5 in the main text). The similar behavior of SubTB()) with learned and true forward state flows suggests that the learned state flows remain close enough to their optimal values and that the variance-reducing benefits of SubTB()) with true state flows are retained.
B EXPERIMENT DETAILS: MOLECULES
All experiments with SubTB()) are based upon the published code of Malkin et al. (2022). which extends that of Bengio et al. (2021a). The proxy model giving the reward, the held-out set of molecules used compute neural network are identical to those in Bengio et al. (2021a). and the off off-policy exploration rate and early stopping likelihood are the same as those tuned for the training with the TB objective in
correlation metric, and the GFlowNet model architecture
15

---

Similarity TB Gradient, Forward Flows
Simllarity with TB Gradient. Bacioward
0.8
Cosine Simiarity
Coeine Simiority
0.4
Loss Comp Name D6 vs TB SubTB TB
0.2
Loss Comp Name D6 vs TB SubTB
500 1030
15:00 2000 2500 3000 2500 4000 Training Stop
500
1930
1500
2030 2500 3000 3500 4000 Training Stop
Figure A.5: Gradient similarity with state flows analytically computed in two ways (see AA.2) (Compare with Fig. 5
Malkin et al. 2022). All models are trained for maximum of 50000 batches of 4 trajectories each. (Some training runs terminated early because of numerical overflows in the gradients, in which case we report the metric of the last stable model whose cumulative number of batches a multiple of 5000.)
C EXPERIMENT DETAILS: BIT SEQUENCES
The modes M as well as the test sequences are selected as described Malkin et al. (2022). The policy for all methods is parameterized by a Transformer (Vaswani et al. (017)) with 3 layers, dimension 64, and 8 attention heads. All methods are trained for 50,000 iterations with minibatch size of 16 using Adam optimizer. For GFlowNets with FM objective as well as the baselines, we use the exact same implementation and hyperparameters reported in Malkin et al. (2022). For TB and SubTB(2), we pick the best learning rate from (0.0075, 0.001, 0.00I. 0.003, 0,005) for forward logits, and for Z, use a learning rate of 10x the learning rate for forward logits. For SubTB(2), we found the best value of 1.9 from the values {0.8, 0.9, 1.1, 1.3, 1.5, 1.7, 1.9, 2.03.
D EXPERIMENT DETAILS: ANTIMICROBIAL PEPTIDE GENERATION
Following Malkin et al. (202)) we use the following amino acids: [ A'' 'C' 'D', 'E', 'F', , 'G' 'H'
K' 'I', 'M', 'N', 'p', 'Q' , 'R' '5' 'T' 'V' W' 'Y' We take 6438 known AMP sequences and 9522 non-AMP sequences from the DBAASP database Pirtskhalava et al. (202I). l). The classifier that serves as the proxy reward function is trained on this dataset, using 20% of the data as the validation set. The reward model is a Transformer, with hidden layers, hidden dimension 64, and 8 attention heads. We train it with a minibatch size 256, with learning rate 10 and with early stopping on the validation set. We use Transformer with 3 hidden layers with hidden dimension 64 with 8 attention heads as the architec- ture of the policy for all methods. All methods are trained for 20, 000 iterations, with a minibatch size of 16, using the reported hyperparameters for all the baselines from (Malkin et al. 2022). For pick the best learning rates from (0.005, 0.007, 0.01, 0.03, 0.05, 0.07) for forward logits and from (0.007, 0.01, 0.03, 0.05) for log Z. For SubTB(l), the best performing value of 1.9 chosen from (0.9, 0.99, 1.1, 1.2, 1.3, 1 ..4, 1.6, 1.7, 1.8, 1.9, 2.0) is used.
TB and SubTB(d),
E EXPERIMENT DETAILS: FL PROTEIN GENERATION
We consider a variant of the GFP task from Trabucco et al. 2022) The vocabulary of 'H' 'I' K'' 'L' 'M' 'N' 'p', '', 'R', '5' 'T' 'V' W 'Y' Following Tra- bucco et al. 2022), we consider the dataset of 56,086 proteins from Sarkisyan et al. (2016) processed based on Brookes et al. (2019). Each protein is accompanied by score quantitying its fluorescence. As with the AMP data, we keep 20% of the data as validation set used for early-stopping. The regressor trained with the dataset is a Transformer, with 4 hidden layers, hidden dimension 64, and attention heads. We train it with a minibatch of size 256, with learning rate 10-4, with early stop- ping on the validation set. The architecture of the policy for all methods is Transformer with 3 hidden layers with hidden dimension 64 with 8 attention heads. All methods are trained for 20, 000 iterations, with minibatch size of 16. We use the same implementation for all methods as the ones used in D
amino acids is the same as D ('A,, 'C' 'D' 'E', 'F' G
16

---

0.4
0.2
-0.1
1.0
1.5
2.0
167
Figure F.I: The Spearman correlation between the sampling probability and reward on test is plotted over the course of training for each value of A.
To define an exploratory training policy, we set the the random action probability to 0.01 selected from (0.0001, 0.0005, 0.001, 0.01) and the reward exponent & (having the same meaning as in $4.2) learning rate of x 10 selected from
to 3 selected from {2. 4). For trajectory balance {10-5, -10-$110~$,00 xx 55 10 10-3, 10 for the flow parameters and x 10 for log Z. For SubTB(.)), we choose the best from (0.7, 0.8, 0.9, 0.99), and found 0.99 to perform the best. For TB and SubTB(d), we tune for the best learning rates from {0.0001, 0.0003, 0.0005, 0.00075, 0.001) for the forward logits. For log Z, we
learning rate of 10x the learning rate for the forward logits.
For FM we use learning rate of 10-3 selected from (10-5, 10 5 x 10 10-3, X 10-3) with leaf loss coefficient d7 30. For A2C with entropy regularization we share parameters between the actor and critic networks, and use learning rate of 5 x 10-3 selected from [10-5, 10-4.5 x 10-4. 10-3 50-33 with entropy regularization coefficient 5x 10 selected from {10-4, 10-3,00 5x 10-3, 10-2.xxx 10-2). For SAC we use the formulation in Christodoulou (2019) with a learning rate of 10-3 selected from {10-5 -,0-5_00 10-4, 4, 5 10-4 10 3, 5 10-"), target network update frequency of 400 and initial random steps of 200. For the MARS baseline,
set the learning rate to x 10 selected from (10-5, 10-4,5 5 10-4, 10-3,5 x 10-3). We run the experiments on 3 seeds and report the mean and standard error over the three runs in Table
F INVERSE PROTEIN FOLDING: NON-AUTOREGRESSIVE SEQUENCE GENERATION
We consider the inverse protein folding problem suggested in Sinai al. 2020): A target protein
3D backbone conformation is given, and the task is to sample amino acid sequences of a fixed length
I. = 40 from the Boltzmann distribution corresponding to their energy in the target conformation.
The energy is provided by a physics model (Rohl al. 2004 Chaudhury et al. 2010). The policy model is a 3-layer convolutional architecture that closely follows previous work (Sinar et al. 2020). Specifically. for the policy function, the convolution size was to 7 with 32 hidden features and ReLU activation each layer. The policy network has one additional convolutional layer of size 20 (number of amino acids), and without the activation function. The flow network has an additional two linear layers of sizes [1280.64], and [64, 1] with ReLU activation in between. We report mean result over three runs.
For this task, rather than generating sequences from left to right, consider an action space in which actions modify one letter at a time at arbitrary positions. The first action uniformly randomly samples an amino acid sequence. On each subsequent action, the agent selects position in the sequence and replaces the letter in this position with another letter in the vocabulary. Generation terminates after exactly N 40 replacement steps. The forward policy is conditioned on the number of steps taken far in the trajectory; the backward policy is fixed to be uniform over the N L actions.
As metric of how well the learned model matches the target distribution, we measure the correlation between log Rxx) and the marginal sampling likelihood log px( held-out set of terminal
17

---

states. The results are presented in Fig. F.1 We observe that intermediate values of lambda lead to the best to the target distribution.
18


Reviews:
[{"review_id": "pQvk9CHp4Kz", "reviewer": "Reviewer_VBco", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper propose a novel GFlowNet training objective called subtrajectory balance $SubTB(\\lambda)$ that can learn from partial The paper proposes a novel GFlowNet training objective called subtrajectory balance $SubTB(\\lambda)$ that can learn from partial action subsequences of varying lengths. Empirical results demonstrate that $SubTB(\\lambda)$ can improve the convergence of GFlowNets in some simple environments.", "summary_of_the_review": "I have some concerns and thus do not recommend acceptance. ", "strength_and_weaknesses": "The proposed method is novel, and the paper is well-written. However, I have some concerns and thus do not recommend acceptance. \n\nQ1. My biggest concern is that the experiments are quite simple. How well does $SubTB(\\lambda)$ perform over complex tasks in high-dimensional environments?\n\nQ2. The motivation of $SubTB(\\lambda)$ is not very celar. More theoretical analysis is required to explain why $SubTB(\\lambda)$ improves the convergence of GFlowNets.\n\nQ3. The current version lacks discussion about the limitations of $SubTB(\\lambda)$.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "oSuOLtyhpJy", "reviewer": "Reviewer_aGGf", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper discusses an improved training procedure to solve GFlowNet problems. The paper begins by outlining GFlowNets, a DAG with exactly one root node where \"actions\" are decisions about which edge to follow in the DAG. When the agent reaches a terminal node (a node with no edges leaving that node), the \"episode\" ends. The goal is to find a policy such that the probability of reaching a terminal state is proportional (up to a constant) to a reward function R.\n\nGFlowNets have strong assumptions not in MDPs (absence of cycles) and a different objective (probability of reaching terminal states), the authors review previous work which shows one can solve the GFlowNets problem with a solution involving an edge flow function which satisfies certain constraints. \n\nFrom there the authors outline previous methods used to solve the edge flow problem, such as Detailed balance and Trajectory balance all of which involve minimizing an objective, which once minimized is guaranteed to result in a policy matching the desired distribution.\n\nThe authors note that in  Malkin et al. (2022), it is shown that a flow solves the balance equation from the DB method if it holds for sub-sequences of nodes in trajectores. Using this observation, the authors present a new objective subtrajectory balance objective, resulting in their method.\n\nThey then hypothesize on why their new method brings a benefit.\n\nThe authors then repeat the experiments from Malkin et al. (2022), showing better results and an experiment on gradient variance and gradient bias.\n\n\n", "summary_of_the_review": "I do not recommend this paper for publication, since it's only a small modification of previous work that doesn't advance the reader's understanding of science. Plus, there's an array of open questions about why the author's method delivers the observed performance gains.", "strength_and_weaknesses": "Strengths:\n\n- well written, structured and easy to follow\n- idea is plausible and understandable\n- experiments improve on state of the art\n\nWeaknesses (in no particular order)\n\nNovelty+Relevance: the paper uses a slight modification on Trajectory balance objective, replacing whole trajectories with sub-trajectories. In the paper, the authors show how it provides better experimental results but don't argue how it advances theoretical understanding of GFlowNets. This issue is compounded since it's work on GFlowNets, which have such strong assumptions (must not contain cycles) that it's not applicable to most problems (chess, go, atari...)\n\nSoundness: The authors claim the improvements they show in experimental results come from their novel SubTB method, but there seems to be another key difference between their work and Malkin et al. (2022). In Malkin, they claim they facilitate exploration by choosing a policy based on a \"tempered (higher temperature) version\" of P_F, while the authors use a \"training policy that is a mixture of P_F with a uniform policy.\" It becomes unclear if gains are due to a different exploration policy or SubTB.\n\nAnother point that is unclear is the following: by using the criteria from eq. 11 with the O(n^2) different subsequences, the gradient's variance is reduced (as the authors claim), allowing for a larger learning rate (0.007 vs. 0.001 in Malkin et al. (2022)). The question is the following: if something is done to reduce the variance of TB, (for example using batches from a replay buffer) so that TB can be trained with a larger learning rate, will the author's method still be better? So is it the sampling of sub-sequences that makes the difference, or the variance reduction that comes from averaging the loss from lots of sequences.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "_M3fN3Z08p_", "reviewer": "Reviewer_dWW3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper considers the recently developed GFlowNets, a family of models using a policy network to sequentially sample from an unnormalised target distribution defined on a discrete domain. Three main objectives are typically used to train such models, each enforcing a flow matching \n(ensuring same incoming and outgoing flow at each node), a detailed balance (at each transition) and a trajectory balance condition (at global level). The work proposes a new objective unifying the last two conditions. The objective (i) extends the notion of balance to portions of trajectories, (ii) computes the balance contribution for each possible sub-trajectory and (iii) finally averages all these contributions. Importantly, the average used in the objective is governed by a single parameter, which is experimentally shown to control the bias-variance tradeoff. Indeed, a large parameter value offers similar behaviours to the trajectory balance objective (low bias, high variance), while a small parameter value shows similar effects to the ones obtained by the detailed balance objective (high bias, low variance). This bias-variance trade-off enables to reduce the time required to train GFLowNets and to generate better and more diverse solutions compared to the ones from previous objectives, as demonstrated by experiments on a hyper grid toy case, on small molecule synthesis and on three sequence generation tasks.", "summary_of_the_review": "The paper proposes a new objective to train GFlowNets. The objective allows to tradeoff between existing local and global criteria (using the detailed and the trajectory balance conditions, respectively). Extensive experimental analysis provides support on the benefits of this unified view. The quality of the paper can be improved by including additional details about the training, including missing baselines and smoothen some strong statements in the explanations of the experimental analysis. Overall, the paper contributes to the advancement of GFlowNets from an empirical perspective.\n\n---- POST REBUTTAL ------\n\nFirst of all I would like to thank the authors for the clarifications to my questions. Also, I appreciated the inclusion of the baseline (based on detailed balance) in Figure 3, which allows for a better understanding.\n\nIndeed, experiments in Figure 3 are the most important in the paper, as providing an empirical analysis of the bias-variance trade-off and giving insights about the behaviour of the strategy. From these figures I currently see three main potential issues:\n1. While it is clear that SubTB($\\lambda$) has the \"lowest\" curve, its convergence occurs at the same number of training iterations of the other approaches. Therefore, why is it possible to claim that the convergence is accelerated (as mentioned for example in the abstract)?\n2. It is unclear why the baselines have strange behaviour. Indeed detailed balance is achieving both lower bias and lower variance than trajectory balance. This suggests that there is some additional factor influencing the results, probably the hyperparameters?\n3. Hyperparameters seem to play a major role also for SubTB($\\lambda$). Indeed it is unclear what is the effect of initialization (for instance note that in 16 x 16 hard the proposed strategy seems to start from a different condition from the other approaches) as well as how $\\lambda$ is chosen.\n\nIn summary, the proposed objective unifies existing local and global objectives for GFlowNets and represents an interesting and novel idea. However, in absence of a theoretical analysis, the paper should provide a solid sensitivity analysis on the hyperparameters and utlimately a methodology for choosing them.\n\nBased on these considerations and also based on others' reviews, I reduce my score.\n\n", "strength_and_weaknesses": "**Strenghts**\n- The paper is well-written. However additional details can be included in order to make it more self-contained (see comments later)\n- The idea is novel, original and the technical content is correct.\n- Several tasks are used for experimental evaluation. \n- Sequence generation tasks provide convincing evidence.\n- Code is available. However, I haven’t run it to check if experiments are reproducible.\n\n**Weaknesses**\n- The sampling strategy doesn’t come with statistical guarantees on the quality of the obtained samples (like MCMC does). However, this is an intrinsic limitation of GFlowNets rather than of the proposed strategy.\n- No supporting theory is developed to analyze the bias-variance tradeoff.\n- Some experiments (hypergrid, bit sequences) miss baselines. Additionally, some experiments miss explanations of observed phenomena. Please refer to subsequent comments. \n\n**Clarifying questions to improve the quality of the paper**\n\nAbout training details:\n- Can you provide an algorithmic table to describe the training of GFlowNets and the computation of the proposed objective in order to make the paper self-contained? Additionally, is the partial trajectory objective computed at the end of or during each episode in a TD-like fashion?\n\nAbout hypergrid experiments:\n- The experimental setting allows only for actions that increase coordinates at each time step. This is quite a strong condition and I wonder what are the effects of such condition. Specifically, It seems that the proposed approach is biased and it doesn’t satisfy general conditions, like detailed balance and ergodicity used in Markov chains (Note also that the proposed strategy is in effect a Markov chain). Consequently, the proposed strategy doesn’t provide any guarantee of correct sampling from the target. For instance, this is visible in Figure 3 (e.g. 2-D grid for 16x16, 60x60 and 16x16-hard and almost all cases with 4-D grid), where even the unbiased approach using the trajectory balance objective doesn’t converge to 0. Note that the phenomenon becomes even more pronounced with the increasing of dimensions. Can you comment on these aspects? Additionally, how does the proposed strategy perform by also allowing actions that decrease the coordinates?\n- The baseline using the detailed balance objective is missing in Figure 3. Can you please include it in order to have a more complete picture of what is going on?\n- Regarding the analysis of gradient bias. Why in Figure 5 (bottom) small-batch SubTB provides a better gradient estimate than small-batch TD? The explanation in the text seems overly stated (indeed the superiority is visible only at 500 training steps and in between 1500-2000 steps), can you be more precise? Additionally, in Figure 4 (right), SubTB is extremely biased compared to TB even for full batch size (K=10). Can you elaborate more on that?\n\nAbout bit sequences experiments:\n- Can you add the baseline using the detailed balance objective in Figure 7?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject