A Variational Perspective on Generative Flow Networks
Authors: Heiko Zimmermann
Fredrik Lindsten
Jan-Willem van de Meent
Christian A Naesseth
Venue: Accepted by TMLR ()

Abstract: Generative flow networks (GFNs) are a class of probabilistic models for sequential sampling of composite objects, proportional to a target distribution that is defined in terms of an energy function or a reward. GFNs are typically trained using a flow matching or trajectory balance objective, which matches forward and backward transition models over trajectories. In this work we introduce a variational objective for training GFNs, which is a convex combination of the reverse- and forward KL divergences, and compare it to the trajectory balance objective when sampling from the forward- and backward model, respectively. We show that, in certain settings, variational inference for GFNs is equivalent to minimizing the trajectory balance objective, in the sense that both methods compute the same score-function gradient. This insight suggests that in these settings, control variates, which are commonly used to reduce the variance of score-function gradient estimates, can also be used with the trajectory balance objective. We evaluate our findings and the performance of the proposed variational objective numerically by comparing it to the trajectory balance objective on two synthetic tasks.

Published in Transactions on Machine Learning Research (04/2023)
A Variational Perspective on Generative Flow Networks
Heiko Zimmermann
husinommerrannnnaaa...
Amsterdam Machine Learning Lab, University of Amsterdam
Fredrik Lindsten
fredrik. dinddsst @liu.se
Division of Statistics and Machine Learning, Linköping University
Jan Willem van de Meent Amsterdam Machine Learning Lab, Universily of Amsterdam
j. jo ww..aandemeenn@@wv..tt
Christian A. Nacsscth Amsterdam Machine Learning Lab, University of Amsterdam
c. L...nassschhuuuu..t
openreview. Reviewed OpenReview: https: https:///peereeviau net/ forum? id= AZ4GobeSLy
Abstract
Generative flow networks (GFNs) are class of probabilistic models for sequential sampling of composite objects, proportional to a target distribution that is defined in terms of an energy function or a reward. GFNs are typically trained using a flow matching or trajectory balance objective, which matches forward and backward transition models over trajectories. In this work we introduce variational objective for training GFNs, which is convex combination of the reverse- and forward KT. divergences, and compare it to the trajectory balance objective when sampling from the forward- and backward model, respectively. We show that, in certain settings, variational inference for GFNs is equivalent to minimizing the trajectory balance objective, in the sense that both methods compute the same ecore-function gradient. This insight suggests that in these settings, control variates, which are commonly used to reduce the variance score-function gradient estimates, can also used with the trajectory balance objective. We evaluate our findings and the performance of the proposed variational objective numerically by comparing it the trajectory balance objective on two synthetic tasks.
1 Introduction
Generative flow networks (GFNs) [Bnggio et al. 2021a b) have recently been proposed as a computationally efficient method for sampling composite objects such as molecule strings (Bengio et al.] (2021a). DNA sequences (Jain et al., 2022) or graphs (Deleu et al. 2022). To generate such objects, GFNs sample a trajectory along directed acyclic graph (DAG) in which edges correspond to actions that modify the object. A trajectory sequentially constructs to terminating node (final composite object) which is scored according to reward signal. While sampling sequences of actions has been well studied in the reinforcement learning lilerature (Sutton and Barto 2018), the objective typically to find a policy which maximizes the expected reward of the trajectory. By contrast, problem Toussaint et al., 20066 by
object by transitioning from a root node (initial object, null state)
GFNs are trained learn a policy that solves a planning-as-infent learning a distribution over trajectories ending in a terminating state with probability proportional to the reward assigned to it. This is done by optimizing objectives which aim to satisfy flow matching detailed

---

Published in Transactions on Machine Learning Research (04/2023)
balance condition (Bengio et al. 2021a). Malkin et al.
(022b) has since found that these objectives are prone to ineffective credit propagation across trajectories and proposes alternative objective based on trajectory balance (TB) condition to alleviate these problems. Most recently, Madan et al. (2022) proposed an objective
that can be optimized on partial trajectorics and [Do et al. 2022) proposed an optimal-transport-baaed
objective to further improve generalization and exploration.
A positive reward function can be interpreted as an unnormalized distribution, which one wishes to generate samples from. In this view, such that the marginal distribution of the final state is approximately equal to the corresponding normalized distribution. Generating approximate samples from an unnormalized target distribution is a well-studied task in probabilistic inference, for which many methods have been developed. Examples include methods
interested in sequentially sampling form factorized joint distibution,
based on MCMC [Hoffman and Gelman 2014 Salimans et al. 2015 Li et al., |017]] Hoffman
2017 Naesseth
et al. 2020 Zhang et al. 2022c), importance sampling Neal (2001
| Del Moral ct al. (2006); Nacssoth et al.
(2019 and variational inference (Blei et al.] 2017 Nacssoth al. 2018 Maddison et al. 2017 Le et 2018
Zimmermann al., 2021]. Recent work on GFNs by Zhang al. (2022b) takes similar view by treating the reward function
energy-based model, which can be trained to maximize the data likelihood following
contrastive divergence- based approach (Hinton 2002), while the forward- and backward transition models are trained by optimizing the TB objective.
In this work we show that, in certain settings, optimizing the TB objective is equivalent to optimizing a forward- or reverse Kullback Leibler divergence. To this end we compare the TB objective when optimized with samples generated from the forward transition model, backward transition model, or a mixture of both, to training with a corresponding variational objective, which takes the form of convex combination of the forward- and reverse Kullback-Lebberr divergence. We identify cases in which the TB objective is equivalent to the corresponding variational objective and leverage this insight to employ variance reduction techniques from variational inference. Finally, we run experiments, to evaluate our theoretical findings and the empirical performance of the trajectory balance and the corresponding variational objective.
Related Work Recent work by Zhang et al. (022aa) identifies equivalences between GFNs and certain classes of generative models. The authors observe that hicrarchical variational auto-encoders are equivalent to a special class of GFNs, and that training hierarchical latent variable models with the forward KL divergence between the full backward- and forward transition model of the GFN is equivalent to training a hierarchical VAE by maximizing its ELBO.
In concurrent and independent work, Malkin al. (2022b) derive the same equivalences between optimizing the TB objective and forward- and reverse KT. divergence that establish in this work. The difference with our work is that we establish these equivalences based on composite objective which is a convex combination of the reverse and forward Kullback-Leible divergences. Using the composite objective we also explore setting for which training is not equivalent but conceptually similar in the sense that the same proportion of samples is taken from the forward and backward model. Furthermore, we discuss and study this objective in context of learning energy-based models. Finally, we also study the differences between variational inference and trajectory balance optimization when the forward and backward trajectory distributions share parameters.
2 Background
In the following we give brief introduction to Generative Flow Networks, variational inference, and variance reduction techniques for discrete variable models. It is important to note that we introduce the concept. of flows, which are used to define the a forward- and reverse transition model based on the flow along the odges of a DAG, for completeness only. In the remainder we only assume the existence of a proper forward- and backward transition model for a given DAG, which does not need to be parameterized by a flow. For a comprehensive study of flows and generative flow networks we refer to Bengio et. al. 2021a b)
2

---

Published in Transactions on Machine Learning Research (04/2023)
2.1 Generative Flow Networks
Generative Flow Networks (Bengio et al. 2021a aim to generale trajectories (so. .1,...;88 87; 88) along the edges of directed acyelic graph G= (S, E). Each trajectory starts in the root, and terminates in a terminating state, ST, before transitioning to a special final state, 8,, which is the single Icaf node G. A non-negative reward signal R(sr) is assigned to each terminating state 8T- The task to learn sampling procedure, or flow, for simulating trajectories, such that the marginal distribution of reaching the terminating state ST is proportional to H((T). We adopt the convention that sf = ST+1. The structure of the DAG imposes partial order, <, states s, E S such that < s' if is ancestor of s'. Hence, any trajectory satisfies 8j < 8g for 0 j < k <
S + +1 and consequently does not contain loops. In Figure [ we illustrate possible DAG structure and corresponding forward- and backward transition models over the domain of (extended) binary vectors.
2.1.1 Trajectory Flows
A trajectory flow is non-negative function FG T R+ on complete trajectories T, i.e. trajectories starting in initial state so and ending in the final state 8y associated with DAG G. Below, we drop the graph subscript for notational convenience. A trajectory flow defines probability measure P over complete trajectories, such that for any event AST
F(A) Z
P(A)
F(A)
[F()
Z Σ F(T), TET
TEA
where Z can be interpreted the total amount of flow. The flow F(s) through state and the flow F(s ss) along an edge s') are denoted by
F(s) = E(f
F(s -
ET (
ETT).
The probability of a trajectory containing the state s, and the forward- and backward transition probabilities are denoted by
F(s) Pr(8118)= P(s + +8 8)= = F(s + 8') F(s)
Pn(s ([(=1))== = P(s - (8)) = F(s - F(s-7)) F(s')
P(s) = 7
A flow is referred to as a Markovian flow if its corresponding probability measure satisfies P(s - s [T) P(s + s' s) for any consccutive states s, and partial trajectory (su)...,8) ending in For Markovian flow and complete trajectory T E T we have,
T
T II t=
P(T) II P (1) t=
Pn(s: 8(41).
GFNs parameterize a Markovian flow on DAG by modeling the forward transition probabilitics Pp(s' の) together with normalizing constant Zep, which can be interpreted an approximation to the total amount. of flow. The trajectory flow is given by
T
I F(st - 8441:0,)) II F(s(4100
T
F(T: Z PF(S41 ((
II t=
Zu II PB(s: s(+1:)).
1=
2.1.2 Training Generative Flow Networks
For reward function RG we want to find transition probabilities Pr and PR, i.e. parameters &, such that Pe(sT $y: (Zsr)/Z= =: (T()). In some scenarios we want to fix the backward transition model, e.g. a uniform distribution model can be advantagcous for exploration, or parameterize with distinct set of parameters In this case, the forward and backward transition probabilities do not correspond to the same flow and, under slight overload of notation, refer to (a 8/: () the backward transition probabilities.
3

---

Published in Transactions on Machine Learning Research (04/2023)
$2
1
Sn
80
8
Q(T) (s 80)
(s 81) PB(s1 (2)
Pr(83 82) P= 88n-1)
Pr(s; (8))
P(r) (s 81)
Ps(32 8a) Ps3 (sn-1 8n) Pa(8n 188))
Figure 1: Graph representing the forward- and backward transition model GFN over states s (0,0, 1 (0 white, 0- light gray, 1 dark gray); the graph DAG in the forward- and backward direction respectively. Each forward transition sets one of the H-bits to cither 0 1 (example in blue), whereas each backward transition deletes set bit (example in orange).
Bengio et al. 2021a) originally proposed objectives to train GFNs based on the flow matching conditions and a detailed balance condition. Malkin et al. (2022a) observe that optimizing these may lead to inefficient credit propagation to early transitions, especially for long trajectories. To alleviate this, Malkin et al. 2002a propose an alternative TB objective for complete trajectories
Zy I PF (= R(ST) I PB
log ZeQ(T:)) ZP(T;
ATB(T,
log
(1)
where A = 00,0,10 1b and we define
P(r:0) = R(ST)
T-1
II Pr(as se+1;0),
((+ = II Pr( ():
Trajectories are sampled from proposal distribution y with full support over the space of trajectories
T. The TB objective is optimized using stochastic gradient descent. The gradient w.r.t. all parameters À = (Ø, 0, vt) is computed as the average over a batch of S ii.d. samples. Solutions correspond to fixed points of the (negative) expected gradient
d
Erng da CTB(T:)) O.
We can compute an unbiased estimate of this gradient using samples from the proposal distribution,
8 Σ
(TB(A) = 5 dA fIB(Ts,)), A),
Ts q.
In section 3 we show how optimizing GFNs using the TB objective corresponds variational inference on complete trajectorics. Going forward, we refer to the probability mass functions Q(T: and P(r: 8) over complete trajectories as forward and backward model, respectively.
2.2 Variational Inference
The problem of finding corresponding forward and backward transition probabilities can alternatively be phrased as a variational inference problem. The goal is to find parameters ø and 0 such that the difference between the forward and backward transition probabilities, measured by a suitable divergence, is minimized. Two commonly used divergence measures are the forward Kullback-Leible divergence (FKL) and reverse Luibbeck-Leibler divergence (RKT.),
log Q(T: Φ)
CAKL(O, 0) KL(Q(: :6) (())))) ENNQ
E |- log w]
(2)
P(r; 0)
log P(r:0)
LFKI(0, (K.() :0) ((:+))) ETar
(T((P::
E (log
(3)
Q(T:Ø)
4

---

Published in Transactions on Machine Learning Research (04/2023)
with the importance weights ((T: The divergences can be optimized using stochastic gradient descent with gradients estimated from samples from the forward model Q and backward model P, respectively, In most setting, samples from P are not readily available and one has to resort other techniques to generate approximate samples, c.g. using importance sampling or MCMC (Naesseth et al. 2020
Computing the derivative of LRKL. w.r.t. parameters 0 of the backward transition model is straightforward, the dependence only appears in the log-weights. We can approximate the resulting expected gradient using S samples from the forward model,
d d0
d (2) = =Ern [-ao log P(T:0) = MiKI (0. 0) = 5 d0
Σ d d0
log P(T.:0),
Ts ~ Q(16).
Similarly, the derivative of LEKL w.r.t. parameters & of the forward transition model and corresponding gradient estimator 9HKL are
d (0,0( =Erop/ = dø
[
[-as d do
log Q(T:0) (0,0())))
Σ d
log Q(Ts:)).
Ta ~ P(:0).
dø
Computing derivative of LRKL w.rtt. ø and derivative of CFKL w.r.t. on the other hand involves computing so-called score-function gradient,
Σ log @)::0) d P 0) dø
Q(T:)) + (2 log Q(T:0) P(r;0)
CRKI (0, do
Q(T:
FET
Q(T: Ø)
d
d log Q(T:))
Σ log
(( Q(;; Φ) de
log ((T;)) + Q(T:))
dø
TET
d (- (-log + log Q(::0) En~q |
[ log log Q(T:0)
d dø
Ernq
dø
Importantly, we can cancel-out the additional scoro-function term (last equality of above equation) Eynql la d de log Q(:: 0)] for any constant a. The corresponding score- functionn gradient estimator is thus
8
W. = P(T,; Q(Ta;))'
kkkL (0, 0) = s
Σ
log w. do
log Q(Ta;)),
Tg ~ Q(:
Analogously, we can compute a score function gradient of LPKL w.r.t. 0 and corresponding estimator
e (f, (2, ETAP log d log P(T:0)
1
d d0
log w. log P(Ts:0), Ts P(:0))
da
S
Score-functe forre-fnntion gradient estimators can exhibit high variance [Ranganath et al, 2013), which can be problematic for learning variational approximations via stochastic gradient descent, and hence is often essential employ variance reduction techniques.
2.2.1 Variance reduction techniques for ecore-function estimators
A commonly used technique to reduce the variance of score-function estimators is to control variate h Roos 1997) to replace the gradient estimator g with the modified estimator g' g + c(h -EA)) ) where is scaling parameter. Control variates leave the expected value of the gradient estimator unchanged, Elg] = E[g'], but has the potential to reduce the variance. Indeed, for a given control variate h we can minimize the variance of g'
Var[g'] Var[g] + e Var[h] 2cCovlg, h]
(4)
5

---

Published in Transactions on Machine Learning Research (04/2023)
with respect to the scaling The scaling that minimizes the variance is
c" = arggmmn Varlg'] = Cov (ovig.1 hl Var
c
When I, and hence g', is a vector, we are typically interested in the scalar- or vector-valued scaling c that minimizes the trace of the covariance matrix of g'. Note that the trace obtained with the optimal vector-valued scaling lower bound on the Irace obtained with the optimal scalar-valued scaling as
D Σ d=1
min
Varly'l
min Varlgal
min
Σ d=1
Var (arlo')-
The optimal scalar- and vector- vector-valued scalings are
1=1 Cov [9d, ha] d==1 Var Var[ha]
Coviga hal Var[ha]
Cvee,t
The score function zø log Ø) |Ranganath et al. 2013) is a useful and easy to compute control variate when optimizing the reverse KL divergence, which we will use as our running example. Using the score function as a control variate simplifies the expression of the optimal baseline and the resulting gradient estimator such that the scaling e can simply be added to the (negative) log-importance weight,
s
log Ø) +e( d de
d de
9 log w.
log (TT:;)))= - log ((T;0))
dø
g
d log (T):i)).
log w. + c)
de
Monte Carlo Estimation. We can estimate the optimal scaling with the same S i.i.d. samples Ta Q(T<i 0) used to estimate However, in order for the gradient estimator to remain unbiased, we have to employ a leave-one-out (L.0)) estimator e MMihh and Rezende 2016), which only makes use of samples (fx s * s), such that
1 s
- log + C.) d log Q(Ts; Φ))
=E [ log log Q(Tai b0
E
-log
S
do
dø
The leave-on-ord out estimate of the optimal scaling for the d-th dimension of c* is
Covelga, hal
Var, [ha]
where Cov.l:.], and Var.[-] are empirical LOO covariance and variance estimates, respectively. Note that estimating the optimal scaling requires access to per-sample gradients and hence requires S forward-backward passes on the computations graph in many reverse-mode automatic differentiation frameworks]] Two popular non-optimal scaling choices that are easily computed and do not. require access to gradient. information are llog [E[log w| and clog Z log B[w] with corresponding LOO estimators
log
Σ
log W
log = log
Σ ws.
S
8 8'=1.'' #
5'=1,5' #
The expressions for the optimal baseline simplifies due fact. that the expectation of the score function zero. 2W used jax (Bradbury al. 2018) efficiently compute per-sample gradients estimate the optimal control variates.

---

Published in Transactions on Machine Learning Research (04/2023)
Interestingly, for elog one show that it is sufficient to only compute the fixed scaling e and instead correct by a factor S-1 to obtain an unbiased estimate of g g
8 Σ (- log w. +
s
(- log wis + elo +
log ((=ai)) 5-1 dø
log w3)
log QTTx 6).
5
del
Glog
In Section 3 we show how can leverage these variance reduction techniques for training GFNs by identifying scenarios in which training GFNs with the TB objective is equivalent to performing variational inference with score-function gradient estimator.
3 Variational Inference for Generative Flow Networks
The trajectory balance objective and variational objectives, introduced in Section 2.2.1 all try to find a forward model Q and backward model such that
T-1 P(T;0) ((T()))=T(TT) II Pa(s: sc+1;0) se II Pr(s1+1 88; (); = Q(T:)),
t=
and hence terminating states 8T which are approximately distributed according to Tr, which is proportional to the reward R. While the TB objective can be optimized with samples from any proposal distribution that has full support on T, it is commonly optimized with samples from either the forward model TF ~ Q or the backward model тв ~ P. Similarly, variational inference commonly optimizes the RKL divergence or FKL divergence, which can be estimated by sampling from the forward model and backward model, respectively.
Zhang et al. (022bb) propose special case of the trajectory balance objectíve using a proposal that first
samples Bernoulli random variable u B(a). This variable then determines whether the trajectory samples are drawn from the forward model or the backward model. The corresponding expected gradient is
d - CTB(T.))
d CTB(T,1)) dt
Ba)B(a)
[u 0(:3~P(:00)
[u Q(:
(5)
dd
=al -aBBn=P((:0)
CTT 2) dd
d
+ (I)-B((()(())))
d at
CTn(T.A))
(6)
We approximate the expected gradient by approximating the expectation w.r.t. the forward and backward model for any backward ratio [0, 1], which is equivalent to optimizing weighted sum of TB objectives,
COTB(TF, TB. A) aCTB(TB, (1 CTB(TrA
where ~ Q(: and тB P): We can similarly define a convex combination of the two KL divergences, which penalizes the FKL objective and RKL objective with a and (1 - a), respoctively,
(0,000 =aLrKI 0). (1 (x,
Like RKL and FKL, this divergence which is non-negative and zero and only if P Q. For 0.5 the objective recovers the Jeffreys divergence.
We are now equipped to compare the various objectives for different settings of and different parameteriza- tions of the forward and backward model. Specifically, we differentiate between two settings: (1) Pp and PB (and hence Q and P) have distinct parameters & and respectively, and (2) Pp and Pn share parameters = o 0. The expected gradient of CaTB can be computed as the convex combination of the expected gradient of LTB w.r.t. samples from the forward model and backward model (see Equation 5 Similarly, LoKL can be computed convex combination of CRKL and CFKL. Thus, in the following we study the cases 0 and separately and results for 0 < follow accordingly.
7

---

Published in Transactions on Machine Learning Research (04/2023)
3.1 Forward model and backward model with shared parameters
If the forward and backward model share parameters 7) (d. 0), e.g. when they are parameterized by the same GFN, the expected gradient of the TB objective (Equation takes the form
[ d dd CTa(T. = -22Era((:o))
((oogu=
z [/1 Z
log Zo + d log Q(T:)) - d log P(r:))] dv
d
Br-a(-w)
+ log
dn
dn
where the proposal q is either the forward model Q d) (a 0) backward model P(T;0) (a 1). The corresponding gradients of the RKL and FKI divergences are
d (n) = -ErnQ(:)) (log w + c)
log Q(::))
+ log P(T:n)
dn d (F) = Br-r(rm) (logw+ + +c) dn
dn
dp
d
d
log P(T;n)
log Q(::m)
dn
dn
where c is scaling parameter as discussed in Section 2.2.1
3.2 Forward model and backward model with distinct parameters
Sampling from the forward model (a 0). In the case where using samples from the forward model 5 ~ @) only, the expected TB gradients reduce to
(1og log
Z
d
E, ((:0()
CTB 2) = - 2E, -(.00((00)
log Q(T;
dø
Zpp
dø
2
d
d0 (IB(T,A)
log + llog
log P(:: 0)
Z
d0
d
2
d
E [ dej
(TB(T,
2T, Q((1))
log w + log
log 7op
Z
dup
Interestingly, the expected gradient w.r.t. & docs not depend log Z and is proportional to the gradient of the standard score-function gradient for the reverse KI-divergence
d dø
(0, 0) = -By-Q(10) (l0g + +c) d dø log Q(::0)) 2 E, B.-0((10)) do (Tb(T
d
Hence, solutions of the corresponding optimization problem correspond to fixed points of the (negative) expected gradient. Moreover, the term log Z/Zo can be interpreted as learned scaling parameter cy for
0,
T Q(::))
==1,
P(:0)
26 (o,
E
log de
log ((::
E
log Q(T:
de
E d (p(60,,,))
2E
log w de
log (T);))
(log w + log z 2E
log Q(T: 0)
do
Z
de
1 CKL
d log P(T:0)]
log w d0 log P(r; 0)]
do
E B CTn( (+)) 2E (1og w + log Z de log P(T:
2E [log w d log P( 0)]
Gradient expressions w.r.t. paramcters of the forward model (Ø) and backward model (8) for the aTB and aKL objective. For a = 0, the gradient w.r.t. o of the aKL objective is proportional the expected gradient of the TBB objective (highlight in first column). Analogously, for 1, the gradient w.r.t. 0 of the aKL objective is proportional to the expected gradient of the aTB objective (highlight in second column).
Table 1:
8

---

Published in Transactions on Machine Learning Research (04/2023)
variance reduction similar to the control variates discussed in section 2.2.1 Optimizing the TB objective w.r.t. parameters of the forward model is equivalent to optimizing RKT. divergence using score-function estimator with a learned scaling parameter Co, updated according to the gradient described above. This insight also suggests that the control variate described in Section 2.2.1 can be used as an alternative to the learned baseline reduce the variance of the expected gradient estimates of the trajectory balance objective.
The expression of the gradient of the RKL w.r.t. parameters of the backward model differs from the expected gradient of the corresponding TB objective
d CRKL(()00) -
d
log P(T:0)
dB
do
The integrand differs by a multiplicative factor log w + cp.
Intuitively, the likelihood of a sample is higher under the backward transition model P than under the forward transition model by more than predicted by -cs objective trics to increase the likelihood of the sample under P and vico versa. In contrast, the gradient of the RKL objective tries to always maximize the likelihood of samples under the backward transition model, which achieves its global maximum for P = Q. Due to the fact that P(;: 1, increasing the probability of P(T;0) 0) for some 7 decreases the probability of other trajectories indirectly. Hence, while both objectives have the same global minima for flexible enough Q and P, their optimization dynamics may differ.
then log w + Cg. and the TB
Sampling from the backward model (a = 1). When samples are taken from the backward model T P(:0) the expected TB gradients reduce to
Z
d
E P(:00
CTB (T, 2)
P(x00)
log w log
log
Q(T:
de
Z
dø
7
d
E, P
d8 (IB(T,
log w + log
log P(T:
Z
d0
E, P(:0( dub
d
(TB(T; 2)
(log w + log
2
d
log Zoo
Z
dup
Here, a similar observation holds. The expected gradient, w.r.t. 0, of the TB objective is proportional to the corresponding gradient of the forward KT--divrgeecc w.r.t. parameters 0
d d CoKL(O, = Er-r(vo)) log w d log P(T:0) d0
1 1. = B BBB
P(:0) d0 CTB(0,0,T)
2
Again, solutions of the corresponding optimization problem correspond to fixed points of the (negative) expected gradient. Morcover, analogously to the provious case, optimizing the TB objective w.r.t. is equivalent to optimizing a FKL divergence w.r.t. using a score-function estimator with a learned scaling parameter Cy-
The expression of the gradient of the FKL w.r.t. parameters of the forward model & analogously differs from the expected gradient of the corresponding TB objective by factor log + co in the integrand,
d LrKI(O,0) = -Br-p(:0) do dō
d log Q(T:0))
Observing the expected gradients of the TB objective and corresponding gradients of the RKL and FKL shows that in certain cases optimizing the TB objective is equivalent to variational inference using reverse or forward KL divergences. This obscrvation also suggests that we can leverage the various variance reduction techniques for scoro-function estimators developed the variational inference literature. We summarize the gradients w.r.t. the parameters of the forward- and backward model in Table
9

---

Published in Transactions on Machine Learning Research (04/2023)
4 Experiments
We have shown that for certain settings, optimizing the aTBB objective is equivalent to optimizing the aKL objective, in the sense that the fixed points are the same and the expected gradient of the aTB objective is proportional to the gradient of the aKL objective. In these settings we can use the variance reduction techniques for score-function gradient estimators reduce the variance of the expected gradients of the TB objective. In settings where optimizing the oTB objective and aKL objective is not equivalent, it is not immediately clear if optimizing the oKL objective is advantageous over optimizing the aTB objective, or vice versa. In the following we compare the performance of the aTB and oKL objective with different control variates and different backward ratios a.
Generating samples from the target density. Optimizing the oTB and aKT, objective with a > 0 require samples from the ground truth distribution 8T = x TGT. In these setting we require either (1) access online samples x TGT (e.g. in the synthetic density experiment in 4.1 or (2) access to a data set X' (= containing samples Ti aer which can be uniformly subsampled during optimization.
Evaluation metrics. If samples from the ground-truth target distribution are available we can sample trajectories from the backward model conditioned on æ. Let
T-1 II Pa( 8441:0) and Pp(s1r-1 1180: Φ) = II 1=0
T-2
PB(So-r-11 87;0)
Pr(s441566))
1=0
Then, we can estimate the marginal likelihood of the data under the the forward model using importance sampling,
s ~ PB(SO.T-11 l2;0) 0)
(7)
N
PB(S6or-11 x;
If no data is available we will report the expected log-weight B [log w| $ log Z.
Structure and representation of the state space. Following Zhang et al. (002bb) target discrete 1) by consecutively sampling values in {0, 1) for each step. To this end we define the state space S = {0, 0, 1) {syh, where 0 indicates that no bit value has been sampled for the corresponding position yet. We further define odges
distribution over lerminating states on ST 1)
E
(8y) A E (')) U ((s,8)):8EE3
s'(s) (s E (s)) |s|= |s'1 1},
where |s| denotes the number of set bits in 8. With these definitions in place we define a DAG G(S, E) that specifies the structure of the state space (see Figure D For mathematical convenience, we map the states s to numeric representations & in which 0, 0 and 1 are replaced by 0, -1 and 1 respectively. This allows us compute the number of set bits |s| = Ea [Sal, and the location and type of the bit added by transition + s' as the signed one-hot vector z' - 3. We can also compute the state -8'(s,'')
(3' B) that results from flipping the newly added bit in s'. These operations are useful for defining the transition model.
Transition model. We consider a fixed backward transition model PB(s, | se+1) which uniformly at random select a set bit and replaces with The forward transition model Pe($+41 sti Φ) uniformly at random selects a (-bit and and replaces with a bit value sampled from a Bernoulli distribution, whose parameters (logits) are the output of function fo S S R.. The corresponding probability mass functions of the forward- and backward transition model are
fo(F))
PB(se] 8141) = |st41
Pr(se+1 Φ) = D D |sal fo()) + fo(-8' (st,)
In practice fo RDD + Dx2 is vector valued function paramcterized by a Multilayer Perceptron (MLP) with weights d. Given a state st, produces D pairs of logits associated with positions in the state vector. The uniformly drawn position d of the added bit is then used to select the corresponding logits fo()) E R2
10

---

Published in Transactions on Machine Learning Research (04/2023)
4.1 Synthetic densities
To model discrete ground-truth distribution тGT over terminating states we follow Dai et al. (2020) Zhang et al. (202bb) and discretize continuous distribution "GT :R² - R+ into 2 equally sized grid cells along each dimension. The cells are remapped to Gray code such that neighbouring grid cells differ in exactly one bit and the resulting pair of 16-bit vectors concatenated to obtain a single 32-bit vector.
We are interested in two settings: (1) Learning a forward model Q(T: Φ) such that its marginal distribution Qr(sr: () approximates fixed distribution (s())) <Gr(sr) over terminating states, and (2) learning forward model and energy function Ç: {0, 1 - IR jointly such that тт (sx; 9) x exp(-E(sr; f)) = (Gr(sT) We optimize the energy function by minimizing the negalive log-likelihood via stochastic gradient descent, interleaving gradient updates to the forward model and energy function.
We approximate the gradient of the log- marginll likelihood
d log (T(:T(T:)0))) = d0 (E(ST: + log Zo) do ((sT:)) - -0 Esy (=0) (5(ST: 0)]
d0
which replaces the expectation w.r.t. TT with an expectation w.r.t. the marginal distribution of K-step Metstnolis-Hasttingss (MH) chain m(x' xx) initialized at data x,
using contrastive divergence- based approach (Hinton 2002
d To ((x:0) 0) - (e [E(x':0]]
Ex)((x)
The MH updates uses the GFN construct proposals (Zhang et al. 2022b). For + oo this gradient update recovers the expected gradient of the log-marginal likelihood.
We evaluate the aTB objective and aKL objective with learned control variate for different backward ratios a. For each backward ratio we consider two settings: (1) jointly learning the energy function & and parameters of the GFN, and (2) using a previously learned fixed energy and learning parameters only. We report the negative log- lieliieli in Table 2 All mumbers are averages over 10 independently trained models.
We find that, unsurrissnglly, for a = 0, in which case optimizing the aTB objective is equivalent to optimizing the aKL objective with a learned control variate both objectives perform comparably, i.e. their average negative log-likelihoods
withing one standard deviation of each other (sce Table 2. Similarly, for 0 < 1, both objectives perform similarly, with aTB having slight edge over aKL in terms of negative log-likelihood. For 1, i.e. when sampling from backward model only, the performance of aTB drops significantly while the performance of the aKL objective remains stable.
24
0-0.0
2spirals learned
8gausslans learned g
2spirals fixod
8gaussians
0m3.25
method
fixed g
22-
cTB, a=0.0
20.163++00013
20.006 015
20.201+:0 032
20.080 20.038:+0..0 014
21
a=1.0
oTBB a=0225
00.1333:000010
10.0000--..011
20.149::0..0 015
20.002:1:000 015
wTB, a05.5
20.11910.000
19.997 10.010
20.134.10.014
19.99910.011
70
oTB, 007.5
20.118::0.00 008
20.002-+0.008
20.126-10.00
19.998-10.011
74
<TB, 1=1.0
20.99440.000 037
20.08840.000 008
21.207:40.00 051
20.076+0.012
6TB
oKI, a=00.0
20.1711:..01
20.021 10.0018
20.196:10.
20.045:10.013
22
aKL, a=0.25
20.142:1.002
19.999110.009
20.147:10.00
20.003:10.000
oKI, a=0.5
20. 100.0001
20.003-10.014
00.14441..008
20.006:10.010
oKL, a=0.75
20.160 0088
20.019 009
20. 152::00000000
20.011#0.011
20
20k
406
30k
202
200k
aKL, a=1.0
20.174+:0.000
20.019#0.010
20.171#0..10
20.016+0.000
Figure Negative likelikel over training on 2spirals for fixed energy and different backward ratios
Table 2: Negative log- liceliiol of test data under learned GFN policy with learned control variate for learned and fixed energy functions and different backward ratios a.
11

---

Published in Transactions on Machine Learning Research (04/2023)
0 =
1.0
0
0 -0.6
[= -0.4
A= 0.2
200
200
Ellog
200
nome
100
140
-50
110
150
-100
-202
-200
0.2
0.4
0
0.8
1.0
100
400
Ellogw]
100
200
-100
0-
200
200
14
24
36
50
54
14
26
nterations
(terations
Figure 3: Expected log importance-weight over training. Using the estimated baselines leads to faster convergence during early stages of training.
4.2 Ising model
We are modeling a discrete distribution over terminating states ST E {-1,1) corresponding to the grid cells of an Ising model,
1
(T(ST) x exp(-BH
H(ST) = ANST
(8)
where A the adjacency matrix of a N x (D= = N2) grid with periodic boundary conditions, and / the interaction strength.
As we do not have access to ground truth samples from the Ising model, are training the GFN with a 0. In this setting, optimizing the aTB objective and aKT. objective is equivalent and hence focus the effect of replacing the learned control variate, that is used in the original aTB objective, with estimated control variate typically used to reduce the variance in score-function estimators.
We report the expected leigweights for different control variates and values B averaged over 10 independently trained GFNs Table 3 and show samples from a GFN and samples generated by rumning MH chain for qualitative comparison Figure | For computational efficiency we compute a slightly biased global estimate for the optimal control variates instead of an unbiased LOO estimate, which have found to not noticeable influence the result. While using the estimated control variates helps the GFN to converge faster, specifically in early stages of training (see Figue find that it has significant influence on the final performance.
method
control variate
8 == -1.0
-0.8
= -0.6
3 = -0.4
-0.2
aKI, a=0.0
none
279.54 +: 34.78
230.50 +: 17.23
161.12 + 20.86
-127.92 4 4.82
-135.94
: 1.71
aKL, a=0.0
learned
369.30 + 0.90
283.92
185.28
+ 0.92
40.38
3.08
-140.76
aKL, a=0.0
log Z
369.57 = 0.60
283.80
+ 1.15
185.70
+ 1.16
40.05
= 3.00
-138.13
0.22
aKL, a=0.0
opt. scalar
369.16 + 0.92
283.84
4 0.79
186.14
4 1.48
39.48
+ 2.87
-138.16
0.18
aKL, a=0.0
opt.
0.78
283.88
1.02
41.50 +
-138.01
0.27
method
control
8
aKI, a=0.0
none learned
1135.94 :: 2.84
-126.93
4.39
143.81
1 22.36
235.06
1. 17.64
274.11
1 31.61
aKL, a=0.0
-140.69
0.40
87.20
+ 21.32
214.36
+ 44.48
326.74
34.11
413.30
+ 35.51
aKL, a=0.0
log Z
-138.27
+ 0.31
18.88-
+ 35.28
211.71
+ 37.90
305.98
+ 44.27
111.81
+ 36.92
aKL, a=0.0
opt.
aKL, a=0.0
opt. vector
-138.05
0.36
74.51
33.47
211.8
1 37.44
326.55
1 34.62
405.99
1 34.91
Table Expected log-weights of oKL with different control variates for ten 15 x 15 Ising models with different. interaction strengths B. Note that aKL with learned control variate is equivalent to aTB for 0.
12

---

Published in Transactions on Machine Learning Research (04/2023)
B=-1.0
8=-0.8
8=-0.6
B=-0.4
6=-0.2
B=0.2
3=0.4
0=0.6
0=
B=1.0
Mi
GFN
Figure Approximate samples from a. 15x Ising model generated using MH chains (top row) and the forward model of a trained GFNs (bottom row).
5 Conclusion
In this paper, we draw connections between the recent literature on generative flow networks and the literature variational inference methods. We observe that GFNs be trained using variational objectives that minimize a divergence between a forward and a backward distribution over trajectories. When minimizing the reverse Kullback-Leibler divergence, the objective is analogous to that used in standard variational inference methods that maximize a lower bound on the log- marginal likelihood (Blei et al. forward Kullback-Leibler divergence, we obtain a variant of the objective that is commonly used in wake-sleop methods and related approaches (Hinton et al. 1995 Bornschein and Bengio 2015 Naesseth et al. 2020 It also possible to optimize convex combination of the two. These objectives are closely related to the tra brajectory-balacee objective that is typically used when training GFNs. Specifically, the gradient of the RKL is proportional to computing the expected gradient of the TB objective with respect trajectorics that are sampled from the forward distribution. Evaluations synthetic densities and an Ising model demonstrate that variational objectives for GFNs achieve a comparable performance in terms of the expected log weight relative to variants of the trajectory balanco objective. This observation opens up opportunitics to explore new variational objectives for GFNs that incorporate credit assignment methods Schulman el al. (2015) as well as importance sampling methods for GFNs based on e.g. variational sequential Monte Carlo Naesseth let al. 2018 or nested variational inference (Zimmermann al. 2021).
2017). When minimizing the
References
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. In Advances in Newral Information Processing Systems, volume 34, pages 27381-27304. Curran Associates, Inc., 2021a.
Yoshua Bengio, Tristan Deleu, Edward J. Hu, Salem Lahlou, Mo Tiwari, and Emmanuel Bengio. GFlowNet Foundations. ar arXiv:2111.09266 [cs, statj, November 2021b. arXiv: 2111.09266.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):859- 877, 2017.
Jörg Bornschein and Yoshua Bengio. Reweighted wake-sleep. In International Conference on Learning Representations, 2015.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnee, and Qiao Zhang, JAX: composable transformations of Python-NumPy programs, 2018. URL http::/githut www com/google/jar
Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, and Dale Schuurmans. Learning Discrete Energy-based Models via Auxiliayyvvriibbl Local Exploration. In H. Larochelle, M. Ranzato, R. Hadsoll, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 10443-10455
Curran Associates, Inc., 2020.
13

---

Published in Transactions on Machine Learning Research (04/2023)
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Sociely: Series B (Slatistical Methodology), 68(3))411 1346, 2006. ISSN 1467-9868. doi: (AL.L/..44679888.200000.66
Tristan Deleu, António Góis, Chris Chinenye Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian Structure Learning with Generative Flow Networks. In The 38th Conference on Uncerlainly in Artificial Intelligence, June 2022.
Anh Do, Duy Dinh, Tan Nguyen, Khuong Nguyen, Stanley Osher, and Nhat Ho. Improving Generative Flow Networks with Path Regularization, September 2022. arXiv:2209.10992 [cs, stat].
Geoffrey E. Hinton. Training products of experts by minimizing contrastive divergence. Newral Computation, 14(8):1771-1800, August 2002. ISSN 0899-7667. doi: 001112//899766220110001,
Geoffrey E Hinton, Peter Dayan, Brendan J Frey, and Radford M Ncal. The" wako-sleep' algorithm for unsupervised neural nelworks. Science, 268(5214::1158 1161, 1995.
Matthew D. Hoffman. Learning Doep Latent Gaussian Models with Markov Chain Monte Carlo. In Proceedings of the 94th International Conference Machine Learning, pages 1510-1519. PMLR, July 2017. ISSN: 2640-3498.
Matthew D Hoffman and Andrew Gelman. The No-U-Ture Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15, 2014.
Moksh Jain, Emmanuel Bengio, Alex-Hernandez Garcia, 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, March 2022. arXiv:2203.0413 04115 [cs, q-bio].
Tuan Anh Le, Maximilian Tgl, Tom Rainforth, Tom Jin, and Frank Wood. Auto-encoding sequential monte carlo. In International Conference on Learning Representations, 2018.
Yingzhen Li, Richard E. Turner, and Qiang Liu. Approximate Inference with Amortised MCMC, May 2017. arXiv:1702.0834 [cs, stat].
Kanika Madan, Jarrid Rector-Brooks Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability, September 2022. arXiv:2209.22782 [cs, stat).
Chris J Maddison, John Lawson George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Teh. Filtering variational objectives. In Advances in Neural Information Processing Systems, pages 6583, 2017.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory Balance: Improved Credit Assignment in GFlowNets. 132559 fes, stat), Jamuary 2022a. arXiv: 2201.13259.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference, October 2022b. arXiv:2210.00500 [es, stat] version: 1.
Andriy Mnih and Danilo Rezende. Variational inference for monte carlo objectives. In International Conference Machine Learning, pages 188-2196. PMLR, 2016.
C. A. Naesseth, S. W. Linderman, R. Ranganath, and D. M. Blei. Variational sequential Monte Carlo. In Proceedings of the 21st International Conference on Artificial Intelligence and Statistics (AISTATS), Lanzarote, Spain, Apr 2018.
C. A. Nacssoth, Lindsten, and T. B. Schön. Elements of sequential Monte Carlo. Foundations and Trends@ in Machine Learning, 12(3):307 392, November 2019. Now Publishers, Inc.
Christian Nacsseth, Fredrik Lindsten, and David Blei. Markovian score climbing: Variational inference with LL(pl))) In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 15499 15510. Curran Associates, Inc., 2020.
14

---

Published in Transactions on Machine Learning Research (04/2023)
Radford M. Neal. Annealed importance sampling. Statistics and Computing, 11(2):125-13) April 2001. ISSN 1573-1375. doi: 10. 10.1023/ A:10009200000000000 5028.
Rajesh Ranganath, Sean Gerrish, and David M. Blei. Black Box Variational Inference. arXie:1401.0118 (es, stat/, December 2013. arXiv: 1401.0118.
Sheldon M Ross.
Simulation. academic press, 1997.
Tim Salimans, Diederik Kingma, and Max Welling. Markov Chain Monte Carlo and Variational Inference: Bridging the Gap. In Proceedings of the 32nd International Conference on Machine Learning, pages 1218 1226. PMLR, June 2015. ISSN: 1938-7228.
John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. Advances in Neural Information Processing Systems, 28, 2015.
Richard Sutton and Andrew G. Barto. Reinforcement Learning, second edition: An Introduction. MIT Press, November 2018. ISBN 978-0-262-35270-3.
Marc Toussaint, Stefan Harmeling, and Amos Storkey. Probabilistic inference for solving (PO)MDPs. Neural Compulation, (D(eecember:3357-3737) -37,, 2006.
Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying Generative Models with GFlowNets, September 2022a, arXiv:2209.06000 [cs, stat].
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative Flow Networks for Discrete Probabilistic Modeling. In Proceedings of the 39th International Conference Machine Learning, pages 26412-26428. PMLR, June 2022b. ISSN: 2640-3498.
Liyi Zhang, David M. Blci, and Christian A. Nacsseth. Transport score climbing: Variational inforence using forward KL and adaptive neural transport. arXiv:2202.01841, 2022c.
Heiko Zimmermann, Hao Wu, Babak Esmacili, and Jan Willem van de Meent. Nested Variational Inference. In Advances in Neural Information Processing Systems, volume 34, pages 14423-20435. Curran Associates, Inc., 2021.
15

---

Published in Transactions on Machine Learning Research (04/2023)
layer number
layer type
activation function
input shape
output shape
number of parameters
1
Dense
Swish
B D
B 256
d 256
2
Dense
Swish
B x 256
B x 256
65536
3
Dense
Swish
B x 256
B x 256
65536
Dense 4
None
B x 256
B x 2D
d 512
Table 4: Neural network architecture of forward model which outputs 2D logits, where B is the batch size and is the size of the state space.
A Implementation details
The forward model first selects one of the 0-bit positions in the input vector uniformly at random. The bit position is then used to select the corresponding logit values from the output of the neural network (architecture described in 4) which are consecutively used to construct Bernoulli distribution to draw the new bit value for the selected bit position. In all our experiments we use fixed backward model, which uniformly at random selects a bit position in the vector representation of the state and deletes it, i.c. sets to 0. Hence, the overall mumber of parameters the number of parameters of the forward-model (Nparams 768D + 131072) and one additional parameter if the control variate is Icarnod. For all experiments trained the forward model using the Adam optimizer with learning rate le-3 and batch size B 256.
16


Reviews:
[{"review_id": "ywobDedn5E", "reviewer": "Reviewer_QqZW", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is well written and felt easy to understand. The contribution is fairly clear: there are links between GFNs and variational methods. The paper feels a bit weak in that it doesn't go further than that, it establishes equivalences and differences on a theoretical level, but the empirical follow up is minimal. There is no characterization of the settings where these differences matter, or of where trade offs between methods appear.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper highlights links between two frameworks, the recent GFlowNet framework and variational methods. In particular it shows that under certain conditions (notably the choice of objective and data distribution) the two frameworks induce the same gradients, namely the on-policy GFlowNet Trajectory Balance objective corresponds to the reverse KL divergence objective when $P_B R/Z$ is taken as the prior and $P_F$ is taken as the posterior.\n\nThe paper then suggests that the variance reduction techniques used in variational methods can therefore be translated to GFNs. This is shown both in theory and in a series of experiments that suggest control variates may have an impact on training GFNs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "In section 4 the authors suggest that the objective of Zhang et al., $L_{\\alpha TB}$ corresponds to sampling from either $P$ or $Q$ with probability $\\alpha$. This doesn't seem correct, at least, as far as I can tell this is not what Zhang et al. propose, since the backward trajectories that they use are obtained by first sampling $x$ _uniformly_ from a dataset, and only then sample a backward trajectory from $P_B$. This is because they assume an identical reward/energy for all instances in the data (which they are jointly learning), but it is not exactly sampling from $P_B$. In fact, in the general case for GFNs, as far as I know there is no tractable way to sample from $P_B$.  \n\nNote that this remark doesn't invalidate the theoretical results shown here, the equivalences drawn still hold -- but this makes some of the empirical work a bit incomplete to me. If we assume that $\\pi_T$ is given, then there's not really any point in training a model that essentially recovers $\\pi_T$. What I understand to be one of the strengths of GFN is that it is off-policy; it might have been interesting to compare using different distributions.\n\nThe Figure 1/Table 1 results are suspicious, not in bad sense, but in that something feels like it should be investigated more. The gap between KL and TB seems to be problem dependent, it only appears for the two spirals. This might be saying something about the optimization dynamics (as the authors point out) or the combo of dynamics + function approximation, which it would be nice to investigate.\n\nFinally, I understand that this is concurrent work with Malkin et al., but it might be worth demarcating the differences more clearly. Also note that the $\\alpha KL$ objective is not exactly novel, although its treatment here matters, see _On a generalization of the Jensen-Shannon divergence_ by Frank Nielsen.\n", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HqyTrU5bqw", "reviewer": "Reviewer_LtBg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Stengths: The idea is ok.\nWeakness: The proposal in this paper need more justification on  real  data set or in theory. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": " In this work the authors introduce a variational objective for training GFNs, which is a convex combination of the reverse- and forward KL divergences, and compare it to the trajectory balance objective when sampling from the forward- and backward model, respectively. \n In certain settings, variational inference for GFNs is equivalent to minimizing the trajectory balance objective, in the sense that both methods compute the same score-function gradient. The proposed variational objective is evaluated numerically  on two synthetic  data. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The authors should either justifying  the method on    real  data set or giving  theoritical study. ", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "npQZnQIxh5", "reviewer": "Reviewer_5oNV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I would like to emphasize the I have not worked a lot with the graphical models before and, in particular, have not published papers in this area.  I am not very familiar with recent discrete models in his field. Therefore, I am not very familiar with the related work and can not provide a confident assessment of the significance of the current work. In my review, I will try to focus on readability and presentation of the results and hope that the other reviewers will assess the significance and novelty.\n\n- The presentation in the paper assumes that the reader is already familiar with the GFNs and only briefly recalls their main concepts. In general, I think this not true, especially because the GFN model (if I correctly understand) has been introduced only about a year ago (2021). The authors provided all the necessary mathematical details of the in Section 2, which is good. However, I would expect the authors to provide a more detailed introduction of GFNs, possibly with some introductory toy example of a GFN model/visual material, e.g., particular cases of direct acyclic graphs for certain problems. This would notably simplify the understanding of the model to the reader.\n- It is not very clear which parts are the authors contribution and are novel and which parts contain known content. For example, I assume that section 3.1 is well known, although it is a little bit unclear from the current text.\n- I did not find some implementation details in the paper (and I do not see “supplementary material” on openreview as well), e.g., learning rates, etc. How much computational time do we need to reproduce the experiment? This is not good for reproducibility.\n- In Section 4, the authors study the TB objective from the variational perspective. They derive a lot of formulas and, if I understand correctly, show that they is some cases match to the gradients of the losses (2), (3). However, the referencing in this paper is extremely inconvenient. For example , it would be nice if the authors in phrases such as “proportional to the [...]” would also add reference to some line from Section 3. Otherwise, it is hard to compare those sections.\n- While the paper focuses on the GFN model, I wonder what are the other models in the field for the tasks in view. The discussion of this aspect is limited. I think it should be added. Otherwise, a reasonable question arises: why do we need to further develop GFNs? I think the authors should provide a clear list of advantages w.r.t. the other (non-GFN) models. The same question is about the experiments: how do GFNs compare with the other models in the field (which can be used as the baselines)? As far as I understand, the experiments mostly focus on testing GFN with various parameter alpha.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors study the recently introduced Generative Flow Networks (GFN) model by (Bengio et. al., 2021). Specifically, the authors focus on the training of such networks with the TB objective (Malkin et. al., 2022), equation 1 in this paper. This objective optimizes two models simultaneously: forward and backward distribution models. The training of GFN requires sampling the trajectories from some proposal distribution over trajectories q, but in practice it seems like everyone samples trajectories from the forward model, backward model or their mixtures. THe authors show that in certain cases such training can be viewed as the “mixed” variational inference, i.e., optimizing the weighted sum of forward and reverse KL divergences between the forward and backward model. This point of view opens possibilities to employ the variance reduction techniques from the variational inference. The authors run several experiments to analyse their findings and the effect of the mixing weight.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- A more deeper discussion of GFNs with toy examples and use-cases.\n- I recommend the authors to add a table summarizing the gradient formulas so that it would be easier for the reader to compare.\n- Add more implementation details, at least to appendix. \n\nSome questions:\n- The authors consider the case when P,Q have the same parameters. How does this work in practice? More precisely, this boils down to the general question of parameterization & practical implementation.\n- I do not understand how large are the final models.\n- Do we build the trajectory flow (dyrect acyclic graph) on our own for  for a particular modeling task, or it is already given?\n- Why do we need both forward and backward models in GFNs? Are there approaches to optimize only one of them?\n- What is the interpretation of the TB objective? Do we want to match the forward model’s and reverse model’s trajectory probabilities?\n", "broader_impact_concerns": "-", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: There is a consensus among two reviewers that the paper is well-written and presents interesting ideas. The third review is too short to take it into account. The summary of the reviews:
- The paper is well-written.
- The links between GFNs and VI are clear.

The authors did a great job in the rebuttal. They followed (most of) the points raised by the reviewers, thus, it seems that the claims are well supported with evidence now.

In my assessment, the paper fulfills all requirements to be accepted as is. 

Decision: Accept (poster)