GFlowNets and variational inference
Authors: Nikolay Malkin
Salem Lahlou
Tristan Deleu
Xu Ji
Edward J Hu
Katie E Everett
Dinghuai Zhang
Yoshua Bengio
Venue: ICLR 2023 poster (2023)
Keywords: variational inference
GFlowNets
probabilistic modeling
weighted importance sampling

Abstract: This paper builds bridges between two families of probabilistic algorithms: (hierarchical) variational inference (VI), which is typically used to model distributions over continuous spaces, and generative flow networks (GFlowNets), which have been used for distributions over discrete structures such as graphs. We demonstrate that, in certain cases, VI algorithms are equivalent to special cases of GFlowNets in the sense of equality of expected gradients of their learning objectives. We then point out the differences between the two families and show how these differences emerge experimentally. Notably, GFlowNets, which borrow ideas from reinforcement learning, are more amenable than VI to off-policy training without the cost of high gradient variance induced by importance sampling. We argue that this property of GFlowNets can provide advantages for capturing diversity in multimodal target distributions. Code: https://github.com/GFNOrg/GFN_vs_HVI.

Published as a conference paper at ICLR 2023
GFLOWNETS AND VARIATIONAL INFERENCE
Nikolay Malkin', Salem Lahlou', Tristan Deleu' Mila, Université de Montréal
Xu Ji, Edward Hu
Mila, Université de Montréal
Katie Everett Google Research
Dinghuai Zhang Mila, Université de Montréal
Yoshua Bengio
Mila, Université de Montréal, CIFAR
ABSTRACT
This paper builds bridges between two families of probabilistic algorithms: (hi- erarchical) variational inference (VI), which is typically used to model distribu- tions over continuous spaces, and generative flow networks (GFlowNets), which have been used for distributions over discrete structures such graphs. We demonstrate that, in certain cases, VI algorithms are equivalent to special cases of GFlowNets in the sense of equality of expected gradients of their learning ob- jectives. We then point out the differences between the two families and show how these differences emerge experimentally, Notably, GFlowNets, which borrow amenable than VI to off-policy train-
ideas from reinforcement learning, ing without the cost of high gradient variance induced by importance sampling. We argue that this property of GFlowNets can provide advantages for capturing diversity in multimodal target distributions.
Code: https: ////tttttne tuuu..uub...... com/ GFNOrg/ VS HVI
INTRODUCTION
Many probabilistic generative models produce sample through a sequence of stochastic choices. Non-neural latent variable models (e.g., [Bei et al.) 2003), autoregressive models, hierarchical varia- tional autoencoders (Sønderby et al. 2016). and diffusion models (Ho et al. 2020) can be said to rely upon shared principle: richer distributions can be modeled by chaining together sequence of sim- ple actions, whose conditional distributions are easy to describe, than by performing generation in single sampling step. When many intermediate sampled variables could generate the same object, making exact likelihood computation intractable, hierarchical models are trained with variational objectives that involve the posterior over the sampling sequence (Ranganath et al. 2016b).
This work connects variational inference (VI) methods for hierarchical models (i.c., sampling through a sequence of choices conditioned the previous ones) with the emerging area of research on generative flow networks (GFlowNets; [Bengio et al., 2021a). GFlowNets have been formulated as reinforcement learning (RL) algorithm- with states, actions, and rewards that constructs an object by a sequence of actions so as to make the marginal likelihood of producing an object pro- portional to its reward. While hierarchical VI is typically used for distributions over real-valued objects, GFlowNets have been successful at approximating distributions over discrete structures for which exact sampling is intractable, such as for molecule discovery (Bengio et al., 2021a), for Bayesian posteriors over causal graphs (Deleu et al. 2022), or for approximate mixximum--ikellhooe training of energy- energy-based based models (Zhang et al., 2022b). Al- though GFlowNets appear to have different foundations (Bengio et al. (2021b) and applications than hierarchical VI algorithms, we show here that the two closely connected.
amortized learned sampler
As main theoretical contribution, GFlowNets coincide in their expected gradients. In particular, hierarchical VI (Ranganath et al. 2016b) and nested VI (Zimmermann et al.] 202 1) are related to the trajectory balance and detailed balance objectives for GFlowNets (Malkin al. 2022 Bengio al., (0216). We also point out the differences between VI and GelowNets: notably, that GFlooNetss automatically perform gradient variance reduction by estimating marginal quantity (the partition function) that acts baseline and allow off- off-policy learning without the need for reweighted importance sampling.
show that special cases of variational algorithms and
Our theoretical results are accompanied by experiments that examine what similarities and differ- ences emerge when one applies hierarchical VI algorithms to discrete problems where GFlowNets
Equul contribution. Contact: n kol ay. mm k nimi 1a quebec.

---

Published as a conference paper at ICLR 2023
have been used before. These experiments serve two purposes. First, they supply missing hi- erarchical VI baseline for problems where GFlowNets have been used in past work. The relative performance of this baseline illustrates the aforementioned similarities and differences between VI and Gets Second, the experiments demonstrate the ability of GFlowNets, not shared by hi- erarchical VI, to learn from policy distributions without introducing high gradient variance. We show that this ability to learn with exploratory off-policy sampling is beneficial in discrete proba- bilistic modeling tasks, especially in cases where the target distribution has many modes.
THEORETICAL RESULTS
2.1 GELOWNETS: NOTATION AND BACKGROUND
We consider the setting of Bengio al. (202 We are given a pointed directed acyclic graph (DAG) G = (S, A), where S is a finite set of vertices (states), and A C S x S is a set of directed edges (actions). If s->' is action, one state that has no incoming edge, called the initial state so E S. States that have no outgoing edges are called terminating. We denote by X the set of terminating states. A complete trajectory a sequence r = (50- ,, such that each 5,-1 is an action and Sn E X. We denote by T the set of complete trajectories and by x, the last state of a complete trajectory T.
say is a parent of s' and s' is child of s. There is exactly
GFlowNets are a class of models that amortize the cost of sampling from an intractable target dis- tribution over X by learning functional approximation of the target distribution using its umnor- malized density or reward function, R : X - R*. While there exist different parametrizations and loss functions for GFlowNet they all define forward transition probability function, or forward policy, (- s), which is a distribution over the children of every state s S. The forward policy is typically parametrized by a neural network that takes a representation of s as input and produces the logits of a distribution over its children. Any forward policy Py induces distribution over com- plete trajectories r e T (denoted by PF as well), which in turn defines a marginal distribution over terminating states x E X (denoted by PT)::
PF(T= (50 s,,
PF ss)
VTET,
(1)
1=0
()(()= =
Σ PF(T) TETxx=xx
Vx X
(2)
Given a forward policy PF. terminating states x E X can be sampled from PF by sampling trajecto- ries T from (T(T) and taking their final states Xy.
GFlowNets aim find a forward policy PF for which (x))) R(x). Because the sum in 22) is typically intractable to compute exactly, training objectives for GFlowNets introduce auxiliary objects into the optimization. For example, the trajectory balance objective (TB; Malkin et al. 2 introduces an auxiliary backward policy PB, which is learned distribution PB(- s) over the parents of every state E S, and an estimated partition function Z, typically parametrized as exp(log Z) where log Z is the learned parameter. The TB objective for complete trajectory is defined as
LTB(T:PF, PB. = log z. PF))
(3)
R(x)) xx)
where PB(T 1 xx) [[o PB(S) s'). If LTB made equal for every complete trajectory t, then (x) or R(x) forallx X and Z is the inverse constant of proportionality: Z Evex R(x).
The objective is minimized by sampling trajectories 7 from some distribution and making gradi- ent steps on 3 with respect to the parameters of Pr. PB. and log Z. The distribution from which r is sampled amounts to choice of scalarization weights for the multi-objective problem of minimiz- ing (3) over all T. If T is sampled from (T) note that this is a nonstationary scalarization we say the algorithm runs on-policy. If is sampled from another distribution, the algorithm runs off-policy; typical choices to sample from tempered version of PF encourage exploration (Bengio et al. 2021a Deleu et al. 2022) or to sample from the backward policy P starting from given terminating states (Zhang et al. (022b). By analogy with the RL nomenclature, we call the behavior policy the one that samples 7 for the purpose of obtaining a stochastic gradient, e.g, the gradient of the objective LTB in (3) for the sampled r.
A pointed DAG is one with designated initial state.
2

---

Published as a conference paper at ICLR 2023
Other objectives have been studied and successfully used in past works, including detailed balance (DB; proposed by Bengio et al. (2021b) and evaluated by Malkin et al. (2022)) and subtrajectory balance (SubTB: Madan et al. 2022). In the next sections, we will show how the TB objective relates to hierarchical variational objectives. In s(C we generalize this result to the SubTB loss, of which both TB and DB are special cases.
2.2 HIERARCHICAL VARIATIONAL MODELS AND GFLOWNETS
Variational methods provide way of sampling from distributions by means of learning ap- proximate probability density, Hierarchical variational models (HVMs: Ranganath et al. 2016bb Sobolev & Vetrov 2019, Vahdat & Kautz, 2020; Zimmermann et al. (02Ig) typically assume that the sample space is a set of sequences (Zi...... Zn of fixed length, with an assumption of condi- tional independence between 2 and 2,41 conditioned on 2,. i.e., the likelihood has a factorization g(z....., 2n) (z)1))())2z)) ... (z(zlZn=)) The marginal likelihood of 2n in hierarchical model involves a possibly intractable sum,
(zn))= = Σ ((a))((aal))..... G
The goal of VI algorithms to find the conditional distributions q that minimize some divergence between the marginal ((zn) and a target distribution. The target is often given as a distribution with intractable normalization constant: typical setting is a Bayesian posterior (used in VAEs, variational EM, and other applications), for which we desire q(zn) oc (
The GFlowNet corresponding to a HVM: Sampling sequences (EI...., Zn) from a hierarchical model is equivalent to sampling complete trajectories in certain pointed DAG G. The states of G distance of i from the initial state are in bijection with possible values of the variable 2ı and the action distribution is given by q. Sampling from the HVM is equivalent to sampling trajectories from the policy (zi++1 (2,) q(<==1|2z) (and (z(z))5))
q(1)), and the marginal distribution
(2n) is the terminating distribution p...
The HVM corresponding to a GFlowNet: Conversely, suppose G = (S. is a graded pointed 2 and that forward policy PF on is given. Sampling trajectories
(s0-s))>... ... -s,) -s)) in G is equivalent to sampling from a HVM in which the random variable 2; is the identity of the (i (i+D) 1)-th state s; in and the conditional distributions ((z=11Zz) are given by the forward policy Pr(St-]ss)) Specifying an approximation of the target distribution in a hierarchical model with n layers is thus equivalent to specifying forward policy PF in graded DAG.
D
The correspondence can be extended to non-graded DAGs. Every pointed DAG G = (S, A) can be canonically transformed into a graded pointed DAG by the insertion of dummy states that have one child and one parent. To be precise, every edge s-s' E A is replaced with a sequence of l - {(s) edges, where ((s) is the length of the longest trajectory from to s, l' = ((s') if s' & X, and We thus restrict our analysis in this = ea ((s") otherwise. This process is illustrated in SA section, without loss of generality, to graded DAGs.
Typically, the target distribution is over the objects X of the last layer of graded DAG, rather than over complete sequences or trajectories. Any backward policy PB on the DAG turns an unnormalized target distribution R over X into an unnormalized distribution over complete trajectories T:
The meaning of the backward policy:
VT T PR(T) P(x+) [x+), with unknown partition function 2 = R(x)) (4)
XEX
The marginal distribution of P3 over terminating states is equal to R(x)22 by construction. There- fore, if Pr is forward policy that equals Pn as a distribution over trajectories, then (x)
R(x)22 R(x).
VI training objectives: In its most general form, the hierarchical variational objective ("HVI ob- jective' in the remainder of the paper) minimizes a statistical divergence D, between the learned and the target distributions over trajectories:
p 1-
(5)
We recall some facts about partially ordered sets. A pointed graded DAG is pointed DAG in which all complete trajectories have the same length. Pointed graded DAGs are also characterized by the following equivalent property: the state space S can be partitioned into disjoint sets S = -Ifin S₁. with Sy- = (xp). called loyers, such that all edges s-ss are hetween states of adjacent layers (s Si,5' S for some /).
3

---

Published as a conference paper at ICLR 2023
Two common objectives are the forward and reverse Kullback-Leibler (KL) divergences (Mnih Gregor 2014), corresponding to f 1 I logr for DKI. (PalPP)) and f / ト - log/ for (P PB), respectively. Other f f-divergences have been used, as discussed in Zhang et al. 2019b); ; Wan et al. 2020). Note that, similar to GFlowNets, ( can be minimized with respect to both the forward and backward policies, or can be minimized using a fixed backward policy.
Divergences between two distributions over tra- jectories and divergences between their two marginal distributions over terminating states distributions are linked via the data process- ing inequality, assuming f is convex (see e.g- Zhang et al. (2019b)), making the former a sen- sible surrogate objective for the latter:
A comparison of algorithms for approxi- mating a target distribution in hierarchical vari- ational model or a GFlowNet. The gradients used to update the parameters of the sampling distribu- tion and of the auxiliary backward policy approxi- mate the gradients of various divergences between distributions over trajectories.
Table 1:
D (R(22111PT) ()(PP))PF)
(6)
Surrogate loss
When both P3 and Py are learned, the diver- gences with respect to which they are optimized need not be the same, as long as both objec- tives are O if and only if PF Pp. For ex- ample, wake-sleep algorithms (Hinton 1993) optimize the generative model Pr using DKL (P)l and the posterior PB using (KL(Pr A summary of common combinations is shown in Table
Algorithm
Py (sampler)
Pe (posterior)
REVERSE KL
DRL (P=l(p))
FORWARD KL WAKE-SLEEP (WS)
Dr. ( (PallPP
DxilPallPPr
DRa (PPllPp)
DxL(P= (P))
al.]
REVERSE WAKE-SLERP
DRL Ra.(Ppr (Pxl)PR)
(xa(Pal1P
On-policy TB
( (PrllPs)
see 2.3
We remark that tractable unbiased gradient estimators for objectives such as 3 may not always exist, as we cannot exactly sample from or compute the density of PB(T) when its normalization constant 2 is unknown. For example, while the REINFORCE estimator gives unbiased estimates of the gradient with respect to PF when the objective is REVERSE KL (see 223)), other objectives, such as FORWARD KI, require imiiotance-weighted estimators. Such estimators approximate sampling from PB by sampling a batch of trajectories (T,) from another distribution π (which may equal Pp) Pa(TT)
and weighting a loss computed for each Ti by scalar proportional to x(T)) Such reweighted importance sampling is helpful in various variational algorithms, despite its bias when the number
of samples is finite (e.g.. [Brnsschein & Bengio 2015 Burda et al. 2016), but it may also introduce variance that increases with the discrepancy between PB and πT
2.3 ANALYSIS OF GRADIENTS
The following proposition summarizes our main theoretical claim, relating the GFN objective of (3) and the variational objective of (5). In $ we extend this result by showing an equivalence between the subtrajectory balance objective (introduced in Malkin al. (2022) and empirically evaluated in [Mdan et al. (2022)) and a natural extension of the nested variational objective (Zimmermann) et 202T) to subtrajectories. A special case of this equivalence is between the Detailed Balance objective (Bengio et al. 2021b) and the nested VI objective (Zimmermann et al. 2021).
Proposition 1 Given a graded DAG G. and denoting by 0. the parameters of the forward and backward policies Pr, Py respectively, the gradients of the TB objective satisfy:
VODKI(PPllPP = Br- pal VV 1 (TL((T)).
(PF I (oLL((()
The proof of the extended result appears in C An alternative proof is provided in $
While (8) is the on-policy TB gradient with respect to the parameters of PF, 77) not the on-policy
TB gradient with respect to the parameters of PB. as the expectation is taken over PR. not Pr. The on-policy TB gradient can however be expressed through a surrogate loss
Er-PF (TD() V
P (PBlIPF) + 2(log log k)
(9)
where 2 = -.rex Rxx), the unknown true partition function. Here D Dog is the pseudo-f-divergene defined by fxx = 1 which not convex for large x. (Proof in B
The loss in 7 is not possible to optimize directly unless using importance weighting (cf. the end of $22)). but optimization of PB using 7 and PF using (8) would yield the gradients of REVERSE
WAKE-SI.EEP in expectation.
4

---

Published as a conference paper at ICLR 2023
Score function estimator and variance reduction: Optimizing the reverse KL loss D Pp) with respect to 6, the parameters of Pr. requires a likelihood ratio (also known as
REINFORCE) estimator of the gradient (Williams 1992), using a trajectory T (or a batch of trajec- tories), which takes the form:
PT(T)
A(T) = Vo log (r(T) where c(T) = log
(10)
R(x,) PT(T xx+)
(Note that the term Voc(T) that is typically present in the REINFORCE estimator is 0 in expectation,
since By-p, IVo log Pr(r)] = Σ. PT(T) high variance norm, thus slowing down learning. A common workaround is to subtract a baseline b from c(T), which does not bias the estimator. The value of the baseline b (also called control
(y voPF(T) 0.) The estimator of (10) known to exhibit
variate) that most reduces the trace of the covariance matrix of the gradient estimator
Br-Pr [c(T)]]V log (T:
E LII log (T;
commonly approximated with E,- [c(r), (see, e.g., Weaver & Tao (001 ) Wu et al. (2018). This approximation is itself often approximated with batch-dependent local baselme, from a batch of trajectories (T,)/=
B Iocal Σ
c
(11)
better approximation of the expectation E,-P, [c(T)] can be obtained by maintaining a running average of the values c(T), leading to global baseline. After observing each batch of trajectories, the running average is updated with step size r7:
bglobul
- b 1
(12)
This coincides with the update rule of log Z in the minimization of (PB.Pr.PB: Z) with learning rate 2 for the parameter log Z (with respect to which the TB objective is quadratic). Consequently, (8) of Prop. shows that the update rule for the parameters of PF. when optimized using the RE- VERSE KL objective, with (22) the same as the update rule obtained by optimizing the TB objective using on on-policy trajectories.
control variate for the score function estimator of its gradient,
While learning backward policy PB can speed up convergence (Malkin et al. 2022), the TB ob- jective can also be used with fixed backward policy, in which case the REVERSE KL objective and the TB objective differ only in how they reduce the variance of the estimated gradients, if the trajectories are sampled on- on-policy. In $4] we experimentally explore the differences between the two learning paradigms that arise when PB is learned, or when the algorithms run off-policy.
3 RELATED WORK
(Hierarchical) VI: Variational inference (Zhang et al. 2019a) techniques originate from graph- ical models (Saul al. 1996 Jordan et al. 2004). which typically include an inference machine and a generative machine to model the relationship between latent variables and observed data. The line of work on black-box VI (Ranganath et al., 2014) focuses on learning the inference machine given a data generating process, i.e., inferring the posterior over latent variables. Hierarchical mod- eling exhibits appealing properties under such settings as discussed in Ranganath et al.) (2016b); Yin & Zhou (2018) Sobolev & Vetrow (2019). On the other hand, works on variational auto-encoders (VAEs) (Kingma & Welling, 2014; Rezende et al., 2014) focus on generative modeling, where the inference machine the estimated variational posterior is a tool to assist optimization of the gener- ative machine or decoder. Hierarchical construction of multiple latent variables has also been shown to be beneficial (Sønderby et al. 2016 Maaløe et al. 2019 Child 2021).
While earlier works simplify the variational family with mean-fiee field approximations (Bishop
2006
modern inference methods rely on amortized stochastic optimization (Hoffman et al.] 2013) One
of the oldest and most commonly used ideas is REINFORCE (Williams 1992] Paisley et al.
2012)
which gives unbiased gradient estimation. Follow-up work (Titsias & Lazaro Gredilla 2014
Gre-
gor et al.) 2014, Mnih & Gregor 2014 Mnih & Rezende 2016) proposes advanced estimators to
reduce the high variance of REINFORCE. The log-variance loss proposed by Richter et al. (2020) is
equivalent in expected gradient of PF to the on-policy TB loss for a GFlowNet with a batch-optima value of log Z. On the other hand, path-wise gradient estimators (Kingma & Welling, 2014) have much lower variance, but have limited applicability. Later works combine these two approaches for particular distribution families (Tucker al. 2017 Grathwohl et al. 2018).
5

---

Published as a conference paper at ICLR 2023
Beyond the evidence lower bound (ELBO) objective used in most variational inference methods, more complex objectives have been studied. Tighter evidence bounds have proved beneficial to the learning of generative machines (Burda et al., [016)] Domke & Sheldon 2018; Rainforth et al., 2018) Masrani al. 2019). As KL divergence optimization sutters from issues such as mean-seeking behavior and posterior variance underestimation (Minka 2005). other divergences are adopted as in
expectation propagation (Minka, 2001; Li et al., 2013), more general f-divergences (Dieng et al., 2017 Wang et al. 2018 Wan et al., 2020), their special case ar-divergences (Hernández-Lobarl et al. 2016). and Stein discrepancy (Liu & Wang [0166] Ranganath et al. (016a). GFlowNetts could be seen as providing novel pseudo-divergence criterion, namely TB, discussed in this work.
Wake-sleep algorithms: Another branch of work, starting with Hinton al. (1995), proposes to avoid issues from stochastic optimization (such as REINFORCE) by alternatively optimizing the generative and inference (posterior) models. Modern versions extending this framework in- clude reweighted wake-sleep [Bornschein & Bengio (2015): Le et al. (2019) and memoised wake-
sleep (Hewitt et al. 2020; Le et al.) 2022). It was shown in Le ct al. (2019) that wake-sleep algo-
rithms behave well for tasks involving stochastic branching.
GFlowNets: GFlowNets have been used successfully in settings where RL and MCMC methods have been used in other work, including molecule discovery (Bengio et al. 2021a Malkin et al.
2022 Madan et al. 2022), biological sequence design (Malkin et al. 2022 Jain et al. 2022 Madan al. 2022). and Bayesian structure learning (Deleu et al. 2022). A connection of the theoretical
foundations of GFlowNets (Bengio et al. 202 b) with variational methods was first mentioned by
Malkin et al. 2022) and expanded in Zhang et 2022a
2023).
A concurrent and closely related paper Ziimmermann et al. 2022) theoretically and experimentally explores interpolations between forward and reverse KI. objectives.
EXPERIMENTS
The goal of the experiments is to empirically investigate two main observations consistent with the
above theoretical analysis:
Observation 1. On-policy VI and TB (GFlowNet) objectives can behave similarly in some cases, when both can be stably optimized, while in others on-policy TB strikes a better compromise than either the (mode-secking) REVERSE KL or (mean-seeking) FORWARD KL VI objectives. This claim is supported by the experiments on all three domains below.
However, in all cases, notable differences emerge. In particular, HVI training becomes more stable near convergence and is sensitive to learning rates, which is consistent with the hypotheses about gradient variance in 22.3
Observation When exploration matters, off off-policy TB outperforms both on-policy TB and VI objectives, avoiding the possible high variance induced by importance sampling in off-policy VI. GFlowNets capable of stable off-policy training without importance sampling. This claim is supported by experiments on all domains, but is especially well illustrated on the realistic domains in $4.2 and $ 4.3 This capability provides advantages for capturing more diverse set of modes.
Observation 1 and Observation 2 provide evidence that off- TB is the best method among those tested in terms of both accurately fitting the target distribution and effectively finding modes, where the latter is particularly important for the challenging molecule graph generation and causal graph discovery problems studied below.
4.1 HYPERGRID: EXPLORATION OF LEARNING OBJECTIVES
In this section, we comparatively study the ability of the variational objectives and the GFlowNet objectives to learn multimodal distribution given by its unnormalized density, or reward function, R. We use the synthetic hypergrid environment introduced by |Bengio et al.] (2021a) and further explored by Malkin et al. (2022). The states form a D-dimensional hypeeggg with side length H, and the reward function has 2D flat modes near the corners of the hypergrid. The states form a pointed DAG, where the source state is the origin so O, and each edge corresponds to the action of incrementing one coordinate in a state by I (without exiting the grid). More details about the environment are provided in D We focus on the case where PB is learned, which has been shown to accelerate convergence (Malkin et al. 2022).
In Fig. D compare how fast each learning objective discovers the 4 modes of 128 x 128 grid, with an exploration parameter Ro = 0.001 in the reward function. The gap between the learned distribution PT and the target distribution is measured by the Jensen- divergence (JSD)

---

Published as conference paper at ICLR 2023
W5
00
KL KL
Reverse WS
Trajectories sampled
Trajectories sampled
Target
32
On Pollcy
Off Palicy
Figure I: Top: The evolution of the JSD between the learned sampler P and the target distribution on the 128 x 128 grid, function of the number of trajectories sampled. Shaded areas represent the standard error evaluated across 5 different runs (on-policy left, off-policy right). Bottom: The average (across 5 runs) final learned distribution P for the different algorithms,, along with the target distribution. To amplify variation, the plot intensity at each grid position is resampled from the Gaussian approximating the distribution over the 5 runs. Although WS, FORWARD KL, and REVERSE WS (off-policy) find the 4 target modes, they do not model them with high precision, and produce a textured pattern at the modes, where should be flat.
between the distributions, avoid giving preference we show graphical representations of the learned 2D terminating states distribution, along with the target distribution. We provide in $ details on how PT and the JSD are evaluated and how hyperparameters were optimized separatety for each learning algorithm.
KI. or the other. Additionally,
Exploration poses challenge in this environment, given the distance that separates the different modes. We thus include in our analysis an off off-policy version of each objective, where the behavior policy is different from, but related to, the trained sampler Py (T). The GFlowNet behavior policy used here encourages exploration by reducing the probability of terminating trajectory at any state of the grid. This biases the learner towards sampling longer trajectories and helps with faster discovery farther modes. When off-policy, the HVI gradients corrected using importance sampling weights.
For the algorithms that use score function estimator of the gradient (FORWARD KL, REVERSE WS, and REVERSE KL), found that using global bascline, explained in 2.2, was better than using the more common local baseline most cases (see Fig. D.II. This brings VI methods closer to GFlowNets and thus factors out this issue from the comparison with the GFlowNet objectives.
We see from Fig. that while FORWARD KL and WS the two algorithms that use DKL (P (PPll PF) local
the objective for Pr discover the four modes of the distribution faster, they converge minimum and do not model all the modes with high precision. This is due to the mean-seeking behavior of the forward KL objective, requiring that P. puts non-zero mass terminating states x where R(x) O. Objectives that use the reverse KL train the forward policy (REVERSE KL and REVERSE WS) are mode-seeking and can thus have a low loss without finding all the modes. it converges to lower value of the
The TB GFlowNet objective offers the best of both worlds, ISD, discovers the four modes, and models them with high precision. This supports Observation 1. Additionally, in support of Observation 2, while both the TB objective and the HVI objectives benefit from off-policy sampling. TB benefits more, as convergence is greatly accelerated.
We supplement this study with comparative analysis of the algorithms on smaller grids DD1
4.2 MOLECULE SYNTHESIS
We study the molecule synthesis task from Bengio et al. 2021aa) in which molecular graphs are generated by sequential addition of subgraphs from library of blocks (Jin al. 2020 Kumar

---

Published as a conference paper at ICLR 2023
0.8
0.6
(og to logR(x) correlation
0.0
-0.2
-0.4
Reverse KL onpoliyy
Reverse KL off-policy
-0.8
on-policy
off-policy
10
16
10
10
reward exponent
learning
Figure 2: Correlation between marginal sampling log-likelihood and log-reward on the molecule generation task for different learning algorithms, showing the advantage of off-policy TB (red) against on-policy TB (orange) and both on-policy (blue) and off-policy HVI (green). For each hyperparameter setting the x-axis (ar or B), we take the optimal choice of the other hyperparam- eter (B or a, respectively) and plot the mean and standard error region over three random seeds.
ct al. 2012). The reward function is expressed in terms of fixed, pretrained graph neural network that estimates the strength binding to the soluble epoxide hydrolase protein (Trott Olson
2010 To be precise, R(x) f(x) where f(x) is the output of the binding model molecule and B a parameter that can be varied to control the entropy of the sampling model.
Because the number of terminating states is too large to make exact computation of the target dis- tribution possible, we use performance metric from past work on this task (Bengio et al. 2021a) evaluate sampling agents. Namely, for each molecule x in a held-out set, we compute log P (x), the likelihood of x under the trained model (computable by dynamic programming, see E and evaluate the Pearson correlation of log (x) and log R(x). This value should equal 1 for perfect sampler, log PT (x) and log R(x) would differ by a constant, the log-partition function log 2.
In Malkin et (2022). GFlowNet samplers using the DB and TB objectives, with the backward policy PB fixed to a uniform distribution over the parents of each state, were trained off-policy. sampled from mixture of the (online) forward policy PF and uniform distribution each sampling step, with a special weight depending on the trajectory length used for the termination action.
Specifically, the trajectories used for DB and TB gradient updates
We wrote an extension of the published code of Malkin et al. (2022) with an implementation of the HVI (REVERSE KI.) objective, using reweighted importance sampling correction. We compare the off-policy TB from past work with the off off-policy REVERSE KI., as well as on-policy TB and REVERSE KL objectives. (Note that on-policy TB and REVERSE KL equivalent in expectation in this setting, since the backward policy fixed.) Each of the four algorithms was evaluated with four values of the inverse temperature parameter B and of the learning rate a, for a total of x 4x4 64 settings. (We also experimented with the off- off-policy FORWARD KL. WS objective for optimizing Pr. but none of the hyperparameter settings resulted in an average correlation greater than 0.1.)
The results are shown in Fig. 2 in which, for each hyperparameter (a or B), we plot the performance for the optimal value of the other hyperparameter. We make three observations:
In support of Observation 2, off-policy REVERSE KL. performs poorly compared to its on on-policy counterpart, especially for smoother distributions (smaller values of (3) where more diversity is present in the target distribution. Because the two algorithms agree in the expected gradient, this suggests that importance sampling introduces unacceptable variance into HVI gradients.
In support of Observation the difference between on on-policy REVERSE KL and on-policy TB is quite small, consistent with their gradients coinciding in the limit of descent along the full-batch gradient field. However, REVERSE KL algorithms are more sensitive the learning rate.
In support of Observation 2, off-policy TB gives the best and lowest- variance fit to the target dis- tribution, showing the importance of an exploratory training policy, especially for sparser reward landscapes (higher B).
8

---

Published as a conference paper at ICLR 2023
Comparison of the Sensen-Shannor divergence for Bayesian structure learning, showing the advantage of off-po TB over on-policy TB and on-policy or off-policy HVI. The JSD is measured between the true posterior distribution p(G D) and the learned approximation PG)
Table 2:
Number of nodes
Objective
3
4
5
(Modified) Detailed Balance
5.32 = 4.15 x 10-6
2.05 = 0.70 x 10-5
4.65 + 1.08 x 10
Off-Policy Trajectory Balance
3.70 2.51 10-7
9.35 2.99 x 10
5.44 + 10
On- PoliPolic Trajectory Balance
0,022 + 0.007
+ 0.125 0.027
0.277 = 0.040 0.306 = 0.042 0.656 = 0.009
On-Policy REVERSE KL (HVI)
0,022 + 0.007
Off-Policy REVERSE KL. (HVI)
0.014 + 0.008
0.605 + 0.019
4.3 GENERATION OF DAGS IN BAYESIAN STRUCTURE LEARNING
Finally, we consider the problem of learning the (posterior) distribution over the structure of Bayesian networks, as studied Deleu et al. (002)) The goal of Bayesian structure learning is to approximate the posterior distribution p(G D) over DAGs G, given a dataset of observations D. Following Deleu al. (2022), we treat the generation of a DAG as a sequential decision prob- lem, where directed edges are added one at a time, starting from the completely disconnected graph. Since our goal is to approximate the posterior distribution p(G D), R(G) = pG, D) as the reward function, which is proportional to the former up to a normalizing constant. Details about how this reward is computed, as well as the parametrization of the forward policy PF. are available in |D.3 Note that similarly to $4.2 and following Deleu et al. (022)). we leave the backward policy PB fixed to uniform.
the joint probability
We only consider settings where the true posterior distribution p(G | D) can be computed exactly by enumerating all the possible DAGs G over d nodes (for d 5). This allows to exactly compare the posterior approximations, found either with the GFlowNet objectives or HVI, with the target posterior distribution. The state space grows rapidly with the number of nodes (e.g., there are 29k DAGs over d 5 nodes). For each experiment, we sampled a dataset D of 100 observations from a randomly generated ground-truth graph G* the size of D was chosen to obtain highly multimodal posteriors. In addition to the (Modified) DB objective introduced by Deleu al. (022), we also study the TB (GFlowNet) and the REVERSE KL. (HVI) objectives, both on-policy and off off-policy.
In 2 we compare the posterior approximations found using these different objectives in terms of their Jensen-Shannon divergence (JSD) to the target posterior distribution P(G D). We observe that on the easiest setting (graphs over d 3 nodes), all methods accurately approximate the poste- rior distribution. But as we increase the complexity of the problem (with larger graphs), we observe that the accuracy of the approximation found with Off-Policy REVERSE KL degrades significantly, while the ones found with the off off-policy GFlowNet objectives (Modiiied) DB & TB) remain very accurate. We also note that the performance of On-Policy TB and On-Policy REVERSE KL degrades too, but not significantly; furthermore, both of these methods achieve similar performance across all experimental settings, confirming our Observation 1, and the connection highlighted in $2.2 The consistent behavior of the off-policy GFlowNet objectives compared to the on-policy objec- tives (TB & REVERSE KL) as the problem increases in complexity (i.e., as the number of nodes increases, requiring better exploration) also supports our Observation 2. These observations are fur- ther confirmed when comparing the edge marginals P(X, - X, D) in Fig. D.3 ( DD3). computed either with the target posterior distribution or with the posterior approximations.
5 DISCUSSION AND CONCLUSIONS
The theory and experiments in this paper place GFlowNets, which had been introduced and moti- vated as a reinforcement learning method, in the family of variational methods. They suggest that off- off-policy GFlowNet objectives may be advantageous replacement to previous VI objectives, pecially when the target distribution is highly multimodal, striking an interesting balance between the mode-seeking (REVERSE KL) and mean-seeking (FORWARD KL) VI variants. This work should prompt more research how best to choose the behavior policy in off-policy GFlowNet training. seen as a means to efficiently explore and discover modes.
Whereas the experiments performed here focused on the realm of discrete variables, future work should also investigate GFlowNets for continuous action spaces as potential alternatives to VI in continuous-variable domains. We make some first steps in this direction in the Appendix ($F). While this paper was under review, Lahlou et al. (2023) introduced theory for continuous GelowNets and showed that some of our claims extend to continuous domains.
9

---

Published as a conference paper at ICLR 2023
AUTHOR CONTRIBUTIONS
N.M., X.J., D.Z., and Y.B. observed the connection between GelowNets and variational inference, providing motivation for the main ideas in this work. N.M., X.J., and T.D. did initial experimental exploration. S.I., N.M., and D.Z. contributed the theoretical analysis. S.I.. and N.M. extended the theoretical analysis to subtrajectory objectives. D.Z. reviewed the related work. S.L.. performed experiments the hypergrid domain. N.M. performed experiments the molecule domain and the stochastic control domain. T.D., H.H., and K.E. performed experiments on the causal graph domain. All authors contributed to planning the experiments, analyzing their results, and writing the paper.
ACKNOWLEDGMENTS
The authors thank Moksh Jain for valuable discussions about the project.
This research was enabled in part by computational resources provided by the Digital Research Alliance of Canada. All authors are funded by their primary institution. We also acknowledge funding from CIFAR, Genentech, Samsung, and IBM.
REFERENCES
Peter Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanzhez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 1806. 01261, 2018.
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.
Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer, 2006.
David M. Blei, Michael I. Jordan, Thomas Griffiths, and Joshua B. Tenenbaum. Hierarchical topic models and the nested Chinese restaurant process. Neural Information Processing Systems (NIPS), 2003.
Jörg Bornschein and Yoshua Bengio. Reweighted wake-sleep. International Conference on Learn- ing Representations (ICLR), 2015.
Yuri Burda, Roger Baker Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. International Conference on Learning Representations (ICLR), 2016.
Rewon Child. Very deep VAEs generalize autoregressive models and can outperform them on im- ages. International Conference on Learning Representations (ICLR), 2021.
Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
Adji Bousso Dieng, Dustin Tran, Rajesh Ranganath, John William Paisley, and David M. Blei. Variational inference via X upper bound minimization. Neural Information Processing Systems (NIPS), 2017.
Justin Domke and Daniel Sheldon. Importance weighting and variational inference. Neural Infor- mation Processing Systems (NeurIPS), 2018.
Dan Geiger and David Heckerman. Learning Gaussian networks. In Uncertainty Proceedings 1994, pp. 235-243. Elsevier, 1994.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Kristjanson Duvenaud, Back- propagation through the void: Optimizing control variates for black-box gradient estimation. In- ternational Conference on Learning Representations (ICLR), 2018.
Karol Gregor, Ivo Danihelka, Andriy Mnih, Charles Blundell, and Daan Wierstra. Deep AutoRe- gressive networks. International Conference Machine Learning (ICMI.,, 2014.
10

---

Published as a conference paper at ICLR 2023
José Miguel Hernández-Lobato, Yingzhen Li, Mark Rowland, Thang D. Bui, Daniel Heráandez- Lobato, and Richard E. Turner. Black-box alpha divergence minimization. International Confer- ence on Machine Learning (ICML), 2016.
Luke B. Hewitt, Tuan Anh Le, and Joshua Tenenbaum. Learning learn generative programs with memoised wake-sleep. Uncertainty in Artificial Intelligence (UAI), 2020.
Geoffrey E. Hinton, Peter Dayan, Brendan J. Frey, and R M Neal. The "wake-sleep" algorithm for unsupervised neural networks. Science, 268 5214:1158 61, 1995.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Neural Infor- mation Processing Systems (NeurIPS), 2020.
Matthew D. Hoffman, David Blei, Chong Wang, and John William Paisley, Stochastic variational inference. Journal of Machine Learning Research (JMLR), 14:1303-1347, 2013
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. 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-249, 2020. ISSN 2041-3211.
Michael I. Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning, 37:183 2004.
Diederík Kingma and Max Welling, Auto-encoding variational Bayes, International Conference on Learning Representations (ICLR), 2014.
Jack Kuipers, Giusi Moffa, and David Heckerman. Addendum the scoring of Gaussian directed acyclic graphical models. The Annals of Statisties, 42(4):1689-166 1691, 2014.
Ashutosh Kumar, Arnout Voet, and Kam Y.J. Zhang. Fragment based drug design: from experimen- tal to computational approaches. Current medicinal chemistry, 19(30):51 (1:30):112851144
Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-Gareía, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of con- tinuous generative flow networks. arXiv preprint 12594, 2023.
Tuan Anh Le, Adam R. Kosiorek, N. Siddharth, Yee Whye Teh, and Frank Wood. Revisiting reweighted wake-sleep for models with stochastic control flow. Uncertainty in Artificial Intel- ligence (UAI), 2019.
Tuan Anh Le, Katherine M. Collins, Luke B. Hewitt, Kevin Ellis, N. Siddharth, Samuel J. Gershman, and Joshua B. Tenenbaum. Hybrid memoised wake-sleep: Approximate inference the discrete- continuous interface. International Conference on Learning Representations (ICLR), 2022.
Yingzhen Li, José Miguel Hernández-Lobato, and Richard E. Turner. Stochastic expectation propa- gation. Neural Information Processing Systems (NIPS), 2015.
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. Neural Information Processing Systems (NIPS), 2016.
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. Interna- tional Conference on Learning Representations (ICLR), 2017.
Lars Maaløe, Marco Fraccaro, Valentin Liévin, and Ole Winther. BIVA: A very deep hierarchy of latent variables for generative modeling. Neural Information Processing Systems (NeurIPS), 2019.
Kanika Madan, Jarrid Rector-Brooks. Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GelowNets from partial episodes for improved convergence and stability, arXiv preprint 2209. 12782, 2022.
11

---

Published as a conference paper at ICLR 2023
Nikolay Malkin, Moksh Jain, Emmanuel Bengio. Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Prucessing Systems (NeurIPS), 2022.
Vaden Masrani, Tuan Anh Le, and Frank D. Wood. The thermodynamic variational objective. Neural Information Processing Systems (NeurlPS), 2019.
Thomas P. Minka. Expectation propagation for approximate Bayesian inference. arXiv preprint 1301.2294, 2001.
Thomas P. Minka. Divergence measures and message passing. 2005.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. Inter- national Conference on Machine Learning (ICML). 2014.
Andriy Mnih and Danilo Jimenez Rezende. Variational inference for Monte Carlo objectives. Inter- national Conference on Machine Learning (ICML), 2016.
John William Paisley, David M. Blei, and Michael I. Jordan. Variational Bayesian inference with stochastic search. International Conference Machine Learning (ICML., 2012.
Tom Rainforth, Adam R. Kosiorek, Tuan Anh Le, Chris J. Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. International Conference on Machine Learning (ICML), 2018.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. Artificial Intelli- gence and Statistics (AISTA7'), 2014.
Rajesh Ranganath, Dustin Tran, Jaan Altosaar, and David M. Blei. Operator variational inference. Neurul Information Processing Systems (NIPS), 2016a.
Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. Conference on Machine Learning (ICML), 2016b.
International
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. International Conference Machine Learning (ICML), 2014.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco R. Ruiz, and Ömer Deniz Akyildiz. VarGrad: A low-variance gradient estimator for variational inference. Neural Information Pro- cessing Systems (NeurIPS), 2020.
Lawrence K. Saul, T. Jaakkola, and Michael I. Jordan. Mean field theory for sigmoid belief net- works. Journal of Artificial Intelligence Research, 4:61-76, 1996.
Artem Sobolev and Dmitry Vetrov. Importance weighted hierarchical variational inference. Neural Information Processing Systems (NeurIPS), 2019.
Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther. Ladder variational autoencoders. Neural Information Processing Systems (NIPS), 2016.
Michalis K. Titsias and Miguel -reeilll.. Doubly stochastic variational Bayes for non-
conjugate inference. International Conference on Machine Learning (ICML), 2014.
Oleg Trott and Arthur J Olson. AutoDock Vina: improving the speed and accuracy of docking with a new scoring function, efficient optimization, and multithreading. Journal of Computational Chemistry, 31(2):455-46 2010.
George Tucker, Andriy Mnih, Chris J. Maddison, John Lawson, and Jascha Narain Sohl-Dicksted REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models. Neural Information Processing Systems (NIPS), 2017.
Arash Vahdat and Jan Kautz. Nvae: A deep hierarchical variational autoencoder. Neural Information Processing Systems (NeurIPS), 2020.
Neng Wan, Dapeng Li, and Naira Hovakimyan. f-divergence variational inference. Neural Informa- fion Processing Systems (NeurIPS), 2020.
12

---

Published as a conference paper at ICLR 2023
Dilin Wang, Hao Liu, and Qiang Liu, Variational inference with tail-adaptive f-divergence. Neural Information Processing Systems (NeurIPS), 2018.
Lex Weaver and Nigel Tao. The optimal reward baseline for gradient- reinforcement learning- Uncertainty in Artificial Intelligence (UAI), 2001.
Ronald J Williams. Simple statistical
algorithms for connectionist reinforcement
learning. Machine Learning, 8(3):229-253) 256, 1992.
Cathy Wu, Aravind Rajeswaran, Yan Duan, Vikash Kumar, Alexandre M. Bayen, Sham M. Kakade, Igor Mordatch, and Pieter Abbeel. Variance reduction for policy gradient with action-dependent factorized baselines. International Conference on Learning Representations (ICLR), 2018.
Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. International Conference on Machine Learning (ICML), 2018.
Cheng Zhang, Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in variational in- ference. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41:2008- 2019a.
Dinghuai Zhang, Ricky T. Q Chen, Nikolay Malkin, and Yoshua Bengio. Unifying generative models with GelowNets. arXiv preprint 2209.02606v/ 2022a.
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), 2022b.
Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying generative models with GelowNets and beyond. arXiv preprini 2209,02606, 2023.
Mingtian Zhang, Thomas Bird, Raza Habib, Tianlin Xu, and David Barber. Variational f-divergence minimization. arXiv preprint /907. / 1891, 2019b.
Heiko Zimmermann, Hao Wu, Babak Esmacili, Sam Stites, and Jan Willem van de Meent. Nested variational inference. Neural Information Processing Systems (NeurIPS), 2021.
Heiko Zimmermann, Fredrik Lindsten, Jan- Willem van de Meent, and Christian A. Naesseth, A variational perspective on generative flow networks. arXiv preprint 2210.07992, 2022.
13

---

Published as a conference paper at ICLR 2023
A
CANONICAL CONSTRUCTION OF A GRADED DAG
95
30
82
96
93
Figure A.1: Illustration of the process by which a DAG (left) can turn into a graded DAG (right). Nodes with double border represent terminating states. Nodes with a dashed border represent dummy states added to make the DAG graded.
Fig. A. shows the canonical conversion of DAG into a graded DAG as described in 2.2 Note that this operation is idempotent: applying it to graded DAG yields the same graded DAG.
B
PROOFS
We prove Prop.
Proof For complete trajectory T T, denote by c(T) log ((xx))((
PF(r)
We have the following:
Voc(r) V log (T( VTC(T) = V log PT(T x,) = V log PB(T)
(13)
(14)
Denoting by fi : logg and f2 - log which correspond to the forward and reverse KL
divergences respectively, and starting from
log PF(T)
Pn) PP
= Br-Pr
Br- Pp [c(r)] + 2,
PB(T)
PB(T)
LHVI, (PF. P DkL(PP)PP) Br-PB
log PF(T)
(Br-Pa [c(T)] log 2).
obtain:
V (P VoBr.> [c(T)] = Br- VV log P)()))c())++vvc(c))))) L (PF, PB) -V [c(T)] [c(T)] EE VV log ( Vac(r)
From 13 and
we obtain: [4
By-PF (V (ec(T)) Br (V log P+(()] (T)]= Σ Pogt)VologPP(t) Σ (pP = vo1=0
TeT
TET
Hence, for any scalar Z O, we can write:
B =0=
=0 1B Br-Pr (r)
and similarly
B IV &c(() = B,-P, VV log (T(T) log Z].
Plugging these two equalities back in the HVI gradients above, we obtain:
(PL. PB) = E (V log (F log ZPF(T)
R(x)) xx))
7PF
V &LLIVI (PF.PH) --E--r, (V log PB (B log R(x,)Pu(T xx)
14

---

Published as a conference paper at ICLR 2023
The last two equalities hold for any scalar Z (that does not depend the parameters of Pp. PB. and that does not depend on any trajectory). In partieular, the equations hold for the parameter Z of the Trajectory Balance objective. It thus follows that:
ZPF(T)
2
(P.. PB)
2
BrrPF
Vo log PBBT ( xx))
2 Br-Pa Po..PP/.ooTTTT.Pp.p..
V LHLL (Pp. PB)
BB PT
Vo log R(xx)PB(T xxx) xx+)
ZPF(T)
2 LTB (
Z)]
As immediate corollary, obtain that the expected on-policy TB gradient does not depend estimated partition function Z.
Next, we will prove the identity
which we restate here:
E--Py [V LTB((T] V Piog (PallP)) + 2(log log 2o
(15)
Proof The RHS of (15 equals
Pa(T [x))R(x))) 2(log 7. -
Pr(T)
B--Pp
log
- 2) log
2PF(T)
PB(T Rxx)((x,
PB(7 (x(T(x))((x+))
PF(T)2
=B,-Pr
V
logg
2(log Z log 2) log
(PF(T)
PB(T xxx)
PB(T ((x) 2(log log 2 log Pn(T xx+)
-B--Pr
2V log PT(T xx) log
2Pr(T)
PB(T [xx)(xx))
-2B,-Pp
V log PH(T xx) log
ZPp(T)
-Br-Py (VoLB(V))))())
C A VARIATIONAL OBJECTIVE FOR SUBTRAJECTORIES
In this section, we extend the claim made in Prop. D to connect alternative GFlowNet losses other variational objectives. Prop. iis is thus a partial case of Prop. This provides an alternative proof to Prop. [
The detailed balance objective (DB): The loss proposed in (Bengio et al.) 202 parametrizes GFlowNet using its forward and backward policies Pr and Рn respectively, along with state flow function F, which is positive function of the states, that matches the target reward function on the terminating states. It decomposes
of transition-depender losses:
Pr. Pn. F) log FFs) (s) F(s')Pn(s (s5)
VS-sS E A Lr
where F(s') R(s) if E X.
(16)
The subtrajectory balance objective (SubTB): Both the DB and TB objectives be as special instances of the subtrajectory balance objective (Malkin et al. 2022, Madan et al.) 2022). Malkin et al. (2022) suggested instead of defining the state flow function F for every state a state flow function could be defined on a subset of the state space S, called the hub states. The loss can losses:
be decomposed into a sum of subtrajectory-dependee
LS (T; Pr, Pp. log F(s1)PF(T)
Vr (5).....5p) partial
(17)
F(s,) (x))
where PF(T) is defined for partial trajectories similarly to complete trajectories ②). (r(7] s)
[[(s-y")r (n(s s'). and again fix F(x) R(x) for terminating states x E X). The SubTB objective reduces to the DB objective for subtrajectories of length I and to the TB objective for complete trajectories, in which case use Z. to denote F(so).
15

---

Published as a conference paper at ICLR 2023
A variational objective for transitions: From on, we work with graded DAG 9 (S, in which the state space S is decomposed into layers: S = - with So = {so} and St. X.
HVI provides class of algorithms to learn forward and backward policies on G. Rather than learn- ing these policies (PF and PB) using a variational objective requiring distributions over complete trajectories, nested variational inference (NVI; Zimmermann et al. 2021)), which combines nested importance sampling and variational inference, defines an objective dealing with distributions over transitions, or edges. To this end. makes use of positive functions F& of the states Sk E Sk, for k = =0....,4. L- 1, to define two sets of distributions P& and Pu over edges from Sk to Sk+1:
((ss)) st)
k ==L-1
oc Fe(s)) sk) (K(SE->SE0)
(F+I(5K+1) PB(x sR+1)
otherwise
(18)
Learning the policies Pr, PB and the functions Fg is done by minimizing losses of the form:
LNvI PB, - D ()((lll
(19)
The positive function Fk plays the same role as the state flow function in GFlowNets (in the DB objective in particular). Before drawing the links between DB and NVI, we first propose a natural extension of NVI to subtrajectories.
C.I A VARIATIONAL OBJECTIVE FOR SUBTRAJECTORIES
Consider a graded DAG G = (S, A) where S = - S₁, So [so). Sy. X. Amongst the L + layers / = 0, ., L, consider K + I + special layers, that call junction layers, of which the states are called hub states. We denote by
mg the indices these layers, and we constrain my 0 to represent the layer comprised of the source state only, and = L representing the terminating states X. On each temminnming junction layer ma L, we define a state flow function F& Smp - R- Given any forward and backward policies Pr and Pn respectively, consistent with the DAG G, the state flow functions define two sets of distributions p& and P& over partial trajectories starting from a state Sm, E Sm, and ending in state Sme.t E &mmai (we denote by Ti the set comprised of these partial trajectories, for k = 0. K 1):
VTk = (Nma
Smaal To (x(TT) FR(Sm) PF(S)))((TT))) PP(T)),
(20)
VTę = (Sme 7 -................ Smm-)) ET (T((T)) FR+1 P PB(T Smaa),
(21)
where Fx is fixed the target reward function R.
Lemma ff = Dxx for all k == 0... K- 1, then the forward policy Pf induces a terminating state distribution PJP that matches the target unnormalized distribution (or reward function) R.
Proof Consider a complete trajectory T (S -... And let Tą= = -Smm. _... Sman), for every k < K.
Smz
5mg
Denote by 2e and Zx the partition functions (constant of proportionality in respectively, for every < K. It is straightforward to see that for every 0 < k K:
18
of Pe and pk
24+1 = 7/2 = Σ
FI+t (5me.))
(22)
Smeet ESmpel
K-1
(x()T) 111- Fx(Sm)) PF(T),
(23)
11K-0 Ze
K-1
[1K-0 FR+1 (5mi(5 PB(T Smg).
(T()))
(24)
[K-1 Ze
Because Pu: = P& for all k = O...K K then both right-hand sides of 23 and 24 are equal.
Combining this with (22), we obtain:
Fo( PF(T)
R(x,) PD(T xx).
ETE
(25)
Żo
Exex R(x)
16

---

Published as a conference paper at ICLR 2023
which implies the TB constraint is satisfied for all r E T. Malkin et al. (2022) shows that this is sufficient condition for the terminating state distribution induced by PF to match the target reward function R, which completes the proof.
Similar to NVI,
Lemma to define objective functions for Pp. Pa, Fk. of the form:
K-1
LSub PB. -o Σ
D (p((l)))))
(26)
Note that the SubNVI objective of (26) matches the NVI objective (Zimmermann al. 2021 when all layers are junction layers (i.c. K = L, and my = for & I L). and matches the HVI objective of (5) when only the first and last layers are junction layers (i.e. K 1, mo 0, and m I.).
C.2 AN EQUIVALENCE BETWEEN THE SUBNVI AND THE SUBTB OBJECTIVES
Proposition 2 Given graded DAG 9 as $2.7 with junction layers mo =0, m.........
as
in S.C.1 For any forward and backward policies, and for any positive function F& defined for the hubs, consider Pa and P& defined in (20) and (2I). The subtrajectory variational objectives of (26)
are equivalent to the subtrajectory balance objective (17) for specific choices of the f-divergences. Namely, denoting by 6, the parameters of Pr. Pn respectively:
Bra-pe Pe (V CoLSSUBB(TX; Pb. F)] = 2V (P:llp))
(27)
Er (V (Tk; Pr. Pp. F)] (.
(28)
where F = Fu.k-1, and : logr and f2 -logt.
PalSmy (Pp(TR)
Proof For subtrajectory Tk (Sm² Smann Tk. let c(TR) log Fear PP((T Sumpet
First, note that because 24 and 24 are not functions of σ. 0 (3))):
FRol (Smen) PP Smen) = = log (T()
Voc(TA)= = log
(29)
Z
Fr (Nme (P()
Voc(Te) Vo log
V log (T(k))
(30)
Zx
17

---

Published as a conference paper at ICLR 2023
We will prove (27 The proof of 28 follows the same reasoning, and is left exercise for the reader.
(Pullb)) (pLll1))
log (T((T) (x(TT)
V EDD (PLl())))
= V Σ (x((T)
TAET
2e V log
Σ (x((T)((TT)+
=
2x
T
0, accordingg
Σ Mocc) 4ET
Σ T
( (TT (C(TT)
=
= - nR VV logg
Σ V log (T()))
following 29
TRET
= -E Pr V log PB(Tx M(Tx)
Σ (x((T)) T
=0 Fi+i (Smi.)
Pb(TK. smmnl
= Brx-pp
V log Pa(Te Smeni log
Fa (P
Fx (P((T)
Bra-Pr Px
V log
2
Fitl (8 PT(TE 8mm+t
2 B
As special case of Prop. 2 when the state flow function is defined for so only (and for the termi- nating states, at which it equals the target reward function), i.e. when K = 1, the distribution Po(T) and (F(T) equal, and are the distributions Po(T) and PT(T). We thus obtain the first two equations of Prop. T consequence Prop. 2
D ADDITIONAL EXPERIMENTAL DETAILS
D.1 HYPERGRID EXPERIMENTS
Details about the environment For completeness, we provide more details about the environ- ment, as explained in Malkin et al. (2022). In D-dimension hypergrid of side length H, the state space is partitioned into the non-terminating states S (0, I 1 and terminating states X ST (0. ..... Hl 1)D The initial state is Ogo (0.... 0) S and in addition to the transitions from a non-terminee ter.inatmmmnigggg state to another (by incrementing one coordinate of the state), an "exit" action is available for all E So that leads to terminating state sT e S The reward at DD) is:
terminating state s = (s)
3d H-1 0.5 E (0.25.0.5)
3d
R(s Ro 0.5
+2
1 d=1
005 E (0.3, (0.3.0.4) (31)
H-1
d=l
where Ro is an exploration parameter (lower values indicate harder exploration).
Architectural details The forward and backward policies are parametrized as neural networks with 2 hidden layers of 256 units each. The neural networks take input one-hot representation of state (also called K- K-hot, or multi-hon hot representations), which is HxD vector including exactly D ones and (H 1)D zeros, and output the logits of PF and Pи respectively. Forbidden actions (e.g. when coordinate already maxed out logits to -oo after the forward pass. Unlike Malkin et al. (2022), we do not tie the parameters of Pr and Pg-
are masked out by setting the corresponding
18

---

Published as a conference paper at ICLR 2023
Reverse KL
Reverse W5
Forward
4x10-
local gltoll
On Policy
3x10
10a
1.5
Off Polley
5.25 K0
10a4
10
2. 3x100
10
10
2.25 x30
Figure D.1: A comparison of the the type of baseline used (local or global) for the three HVI algorithms that use score function estimator of the gradient.
Behavior policy The behavior policy is obtained from the forward policy PF by subtracting a scalar from the logits output by the forward policy neural network. The value of is decayed from Gnir to 0 following cosine annealing schedule (Loshchilov & Hutter 2017). and the value is reached at an iteration max. The values of Einit and max were treated as hyperparamters.
Hyperparameter optimization Our experiments have shown that HVI objectives were brittle to the choice of hyperparameters (mainly learning rates), and that the ones used for Trajectory Balance in Malkin et al. (2022) do not perform as well the larger 128 x 128 grid we considered. To obtain a fair comparison between GFlowNets and HVI methods, a particular care was given to the this domain. The optimization was performed two stages:
optimization of hyperparameters
We use a batch size of 64 for all learning objectives, whether on-policy or off-policy, and the Adam optimizer with secondary parameters set to their default values, for the parameters of Pr, the parameters of Pn. and log Z (which is initialized at 0). The learning rates of Pr. PB. log Z. along with a schedule factor y < by which they are multiplied when the JSD plateaus for more than 500 iterations (i.e. 500 x 64 trajectories sampled), were sought after separately for each combination of learning objective and sampling method (on-policy or off-policy), using a Bayesian search with the JSD evaluated at 200K trajectories as an optimization target. The choice of the baseline for HVI methods (except WS, that does not have score function estimator of the gradient) was treated as hyperparameter as well.
2.
All objectives were then trained for 10 trajectories using all the combinations of hyperparame- ters found in the first stage, for 5 seeds each. The final set of hyperparameters for each objective
and sampling mode was then chosen as the one that leads to the lowest area under the JSD curve (approximated with the trapezoids method).
For off- off-policy runs, Tmax was defined fraction 1/n of the total number of iterations (which is equal to 10 10"/64). The value of n and Enir was optimized the same way as the learning rate and the schedule, described above.
In Fig. D. illustrate the differences between the types of baselines considered (global and local) for the algorithms that use a score function estimator of the gradient, both on on-policy and off-policy.
Smaller environments:
The environment studied in the main body of text (128 x 128, with Ro = 10-3) already illustrates some key differences between the Forward and Reverse KI. objec- tives. As sanity check for the HVI methods that failed to converge in this challenging environment, we consider two alternative grids: 64x66 64 and 8x8x8 8, both with an easier exploration parameter (Ro = 0.1), and compare the 5 algorithms on-policy on these two extra domains. Additionally, for the two-dimensional domain (64 64), we illustrate in Fig. D.2 a visual representation of the aver- age distribution obtained after sampling 106 trajectories, for each method separately. Interestingly, unlike the hard exploration domain, the two algorithms with the mode-seeking KL (REVERSE KL and REVERSE WS) converge to a lower JSD than the mean-seeking KL algorithms (FORWARD KL and WS), and are on par with TB.
D.2 MOLECULE EXPERIMENTS
Most experiment settings were identical to those of Malkin et al. (202)), in particular, the reward model / the held-out set of molecules used to compute the performance metric, the GelowNet model
19

---

Published as a conference paper at ICLR 2023
8x8x8x8
64x64
20
Reverse KL
20
Forward Reverse W5
Trajectories sampled
Trajectories sampled
Figure D.2: Top: The evolution of the JSD between the learned sampler P and the target distribu- tion on the x 8 x 8 x 8 grid left and the 64 x 64 grid right. Trajectories are sampled on-policy. Shaded areas represent the standard error evaluated across 5 different runs Bottom: The average (across 5 runs) final learned distribution P for the different algorithms, along with the target distri- bution. To amplify variation, the plot intensity at each grid position is resampled from the Gaussian approximating the distribution over the runs.
architecture (a graph neural network introduced by by [Bngii et al.) (2021a)). and the off-policy exploration rate. All models were trained with the Adam optimizer and batch size 4 for maximum of 50000 batches. The metric was computed after every 5000 batches and the last computed value of the metric was reported, which was sometimes not the value after 50000 batches when the training runs terminated early because of numerical errors.
D.3 BAYESIAN STRUCTURE LEARNING EXPERIMENTS
Bayesian Networks A Bayesian Network is a probabilistic model where the joint distribution over d random variables (X Xa) factorizes according to directed acyclic graph (DAG) G:
[ p(xi Pa)(Xi)).
p(
Xa)
where Pag (X;) is the set of parent variables of X; the graph Each conditional distribution in the factorization above is also associated with set of parameters O E O. The structure G of the Bayesian Network is often assumed be known. However, when the structure is unknown, we can leamn based on a dataset of observation D: this is called structure learning.
Structure of the state space We use the same structure graded DAG G as the one described in (Deleu et al., 2022), where each state of G is itself DAG G, and where actions correspond to adding one edge to the current graph G to transition to new graph G'. Only the actions maintaining the acyclicity of G' are considered valid; this ensures that all the states are well-defined DAGs, DAGs). Similar to
meaning that all the states the hypergrid environment, the action space also contains an extra action "stop" to terminate the generation process, and return the current graph as sample of our distribution; this "stop" action is denoted GT. follow the notation introduced in $2.1
terminating here (we define distribution
Reward function Our objective Bayesian structure learning to approximate the posterior distribution over DAGs p(G D), given a dataset of observations D. Since our goal is to find forward policy Pp for which (G) R(G) (see 2211,,
define the reward function P(D GG) where p(G) is prior over graphs (assumed to be uniform throughout the paper), and ((D G) is the marginal likelihood. Since the marginal likelihood involves marginalizing over the parameters of the Bayesian Network
the
joint distribution R(G) = p(G
( PDD
p(D ( G) do,
is in general intractable. We consider here special class of models, called models, where the marginal likelihood can be computed in closed form; for this class models, the mer-maggina likelihood is also called the BGe score (Geiger & Heckerman, 1994] Kuipers et al. 2014 in the structure learning literature.
20

---

Published as a conference paper at ICLR 2023
Edge marginals
0.5
0.4
RMME
0.2
0.1
Oiphalty
(Modited
Cheboloy
OhrPatay
(Med Eed)
Ghitslay
Figure D.3: Comparison of edge marginals computed using the target posterior distribution and using the posterior approximations found either with the GFlowNet objectives, REVERSE KI. Performance is reported as the Root Mean Square Error (RMSE) between the marginals (lower is better).
For each experiment, we sampled a dataset D of 100 samples from randomly generated Bayesian network. The (ground truth) structure of the Bayesian Network was generated following Erdōs- Rényi model, with about d edges on average (to encourage sparsity on such small graphs with d & 5). Once the structure is known, the parameters of the linear-Gaussian model were sampled randomly from a standard Normal distribution N(0. 1). See (Deleu al., 2022) for more details about the data generation process. For each setting (different values of d) and each objective, we repeated the experiment over 20 different seeds.
Forward policy Deleu et al. (2022) parametrized the forward policy PF using linear transformer, taking all the d possible edges in the graph G as an input, and returning a probability distribution We chose to parametrize PF using simpler neural network architecture, based on graph neural network (Battaglia et al. (018). The GNN takes the graph G embedding, and returns for each node X, a pair of embeddings u, and V;- The probability of adding an edge X; + Xj to transition from G to G' (given that we do not terminate in G) is then given by
over those edges, where the invalid actions were masked
input, where each node of the graph is associated with a (learned)
PG G, GGG exp(u, vj).
assuming that X, - X, is a valid action (i.c., it doesn't 't introduce a cycle in G), and where the normalization depends only on all the valid actions. We then use hierarchical model to obtain the forward policy (G' G), following (Deleu al., 2022):
PF
G))
Recall that the backward policy Ps fixed here, the uniform distribution over the parents of G (i.c. all the graphs were exactly one edge has been removed from G).
(Modified) Detailed Balance objective For completeness, we recall here the modified Detailed Balance (DB) objective (Deleu al., (002) special case of the DB objective (Bengio et al. 2021b, see also (6)) when all the states of G are terminating (which is the case in our Bayesian structure learning experiments):
L(M G':PP. log R(G')PE(G G G) |
R(G)Pp(G' G) GG)
Optimization Following DDeeuu et al. 2022 we used replay buffer for all our off off-policy ob- jectives (Moddifidd) DB, TB, and REVERSE KL). All the objectives were optimized using batch size of 256 graphs sampled either on-policy from Pr. or from the replay buffer. We used the Adam optimizer, with the best learning rate found among (10-6 6 10-6 10-5 10-5,3 3,3x x 10-5 10-4). For the TB objective, we learned log Z using SGD with learning rate of 0. and momentum 0.8.
Edge marginals In addition to the Jensen Lensen-Shanne Shannon divergence (JSD) between the true posterior distribution pG D) and the posterior approximation P.(G) (see $ for details about how this
21

---

Published as a conference paper at ICLR 2023
divergence is computed), also compare the edge marginals computed with both distributions. That is, for any edge X; - xj the graph, we compare
Σ 2) and -
E PF(
p(x; - GiX,EPiG(X))
GiX
The edge marginal quantifies how likely edge X; Xj is to be present in the structure of the Bayesian Network, and of particular interest the (Bayesian) structure learning literature. To measure how accurate the posterior approximation PF is for the different objectives considered here, we use the Root Mean Square Error (RMSE) between p(x; - X, D) and (X, - Xj), for all possible pairs of nodes (X,, X2) in the graph.
Fig. D.3 shows the RMSE of the edge marginals, for different GFlowNet objectives and REVERSE KL (denoted as HVI here for brevity). The results on the edge marginals largely confirm the ob- servations made in $44.3 the off-policy GFlowNet objectives ((Modified) DB TB) consistently perform well across all experimental settings; On-Policy TBB & On-Policy REVERSE KI perform similarly and degrade as the complexity the experiment increases (as d increases); and Off-Policy REVERSE KI. has performance that degrades the most as the complexity increases, where the edge marginals given by (X, Xj) do match the true edge marginals p(X; X, D) accurately.
E METRICS
Evaluation the terminating state distribution F.: When the state space is small enough (e.g- graphs with d s 5 nodes in the Structure learning experiments, or a 2-D hypergrid with length 128, can propagate the flows in order compute the terminating
in the Hypergrid experiments), state distribution PT from the forward policy Pp. This done using a flow function defined recursively:
if 50 otherwise
F(s) (EEPar(s) F(s)Pr(s' (s)
(32)
PT then given by:
(s P(s)Pr(s) s).
(33)
The recursion can be carried out by dynamic programming. by enumerating the states in any topo- logical ordering consistent with the graded DAG G. In particular, computation of the flow at a given terminating state is linear in the number of states and actions that lie trajectories leading to s, and computation of the full distribution P F is linear in JS] + |AA]
Evaluation of the enough, the target distribution pT R/Z" can be evaluated exactly, given that the marginalization is over X only. The JSD is symmetric divergence, thus motivating our choice. The JSD can directly be evaluated as:
divergence (JSD) Similarly, when the state space is small
JSD(P'
2 (P'M) DKL (PF
where M = (P + 1//222
(34)
2P (s)
(s)
Σ $6
P (s) (T()) + (s)
+ ( log
(35)
(T) + Ps)
22

---

Published as a conference paper at ICLR 2023
F EXTENSION TO CONTINUOUS DOMAINS
As a first step towards understanding GelowNets with continuous action spaces, we perform an experiment on a stochastic control problem. The goal of this experiment is to explore whether the observations in the main text may hold in continuous settings well.
We consider environment in which agent begins at the point xo (0, 0) in the plane and makes a sequence of K = 10 steps over the time interval [0. 1]. through points No.1 N0.2,... X1- Fach step from x, to X,-0.1 is Gaussian with learned mean depending on x, and I and with fixed variance; the variance is isotropic with standard deviation 2V Equivalently, the agent samples the
Fuler-Maruyan discretization with interval Ar of the Itô stochastic differential equation
1 dx, 1xx,1) dt +
2 dd..
(36)
where w, is the two-dimensional Wiener process.
The choice of the drift function f determines the marginal density of the final point, X1. We aim to find f such that this marginal density proportional to a given reward function, in this case quantity proportional to the density function of the standard 8gauss ians distribution, shown in Fig. F2 We scale the distribution so that the modes of the 8 Gaussian components are at distance 2 from the origin and their standard deviations are 0.25.
In GFlowNet terms, the set of states is S = {(0.0) (0) U {(x, f) : X..R (0.1,0.2.... 1}}. States with 1 are terminating. There is an action from (x, to (x'r)) if and only if = Ar. The forward policy is given by a conditional Gaussian:
VAr
PF((x' (x,1) N xx f(x
(37)
We impose a conditional Gaussian assumption on the backward policy as well, i.c.,
Pn((x, () (x', | (x-x:;nnxX AI) MDM. A0))) 1#0 0=
(38)
where μa and log o are learned. Notice that all the policies, except the backward policy from time I to time 0, now represent probability densities; states can have uncountably infinite numbers of K
children and parents.
We parametrize the three functions fgg log o as small (two hidden layers, 64 units per layer) MIPs taking as input the position x and an embedding of the time f. Their parameters can be optimized using any of the five algorithms in Table I of the main text Fig. F.1 shows the marginal densities of x, (estimated using KDE) for different in one well-tranned model, as well as some sampled points and paths.
In addition to training on policy, we consider exploratory training policies that add Gaussian noise to the mean of each transition distribution. We experiment with adding standard normal noise scaled by Texp. where Texp {0. 1. 0.2).
Fig. compares the marginal densities obtained using different algorithms with on- on-policy and off- policy training. The algorithms that use a forward KL objective to learn PB namely, REVERSE WS and FORWARD KL are not shown because they encounter NaN values in the gradients early in training, even when using a 10x lower learning rate than that used for all other algorithms (10-3 for the parameters of f, μB, log of and 10-1 for the log Z parameter of the GFlowNet).
These results suggest that the observations made for discrete-space GFlowNets the main text may continue to hold in continuous settings. The first two rows of Fig. F2 show that off-policy explo- ration is essential for finding the modes and that TB achieves better fit the target distribution. Just as in Fig. i although all modes are found by WAKE-SLEEP, they are modeled with lower pre- cision, appearing off-centre and having oblong shape, which is reflected in the slightly higher MMD.
We conjecture (and strongly believe under mild assumptions) but do not prove that the necessary Gel theory continues to hold when probabilities placed by probability densities; the results obtained here evidence support of this conjecture.
23

---

Published as a conference paper at ICLR 2023
5g
target
Figure F.1: Above: KDE (2560 samples, bandwidth 0.25) of the agent's position after steps for i = O, I..... 10 (r 0, O.I..... for model trained with off policy TB, showing close match to the target distribution (also convolved with the KDE kernel for fair comparison), Below: sample of 2560 points from the trained model and the trajectories taken by 128 of the points.
On-Policy
Off-Policy
Algorithm
Texp
Fexp 0.1
Texp 0.2
TB
MMD
0.1111
0.0005
0,0075
REVERSE KL
MMD
0.0027
0.0059
0.0036
WAKE-SLEEP
MMD
0.0008
0.0036
0.0043
Target
Figure F.2: KDE of learned marginal distributions with various algorithms and exploration policies and MMD with Gaussian kemel exp(-l|x y)l?) estimated using 2560 samples.
24


Reviews:
[{"review_id": "zFfO1Oib29", "reviewer": "Reviewer_6mPn", "summary": "", "questions": "", "limitations": "", "rating": 6, "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 studies the relationship between variational inference and GFlowNet. This work finds out these two algorithms are equivalent in the sense of the expected gradients of their learning objectives. Moreover, they demonstrate the superiority of the GFlowNet on off-policy training.", "summary_of_the_review": "Overall this paper is well-written and provides an interesting connection between GFlownet and VI algorithms, and demonstrate the advantage of GFlowNet on off-policy training. ", "strength_and_weaknesses": "Strength:\n- The connection between the variational inference and GFlowNet is a new finding. This paper finds the intrinsic property, variance reduction, that makes the GFlowNet stand out. They also verify that GFlowNet is capable of stable off-policy training without importance sampling.\n- This paper supplies a missing baseline result for hierarchical VI algorithms on modeling discrete random variables.\n\nWeaknesses:\n-   As the author mentioned in the conclusion section, this paper only studies the performance of GFlowNet and VI algorithms on modeling discrete random variables (since GFlowNet was mostly studied on such regime), which could be limited considering that VI algorithms are mostly used for continuous variables. \n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": {}}, {"review_id": "FoG3qUc2wkR", "reviewer": "Reviewer_CLT1", "summary": "", "questions": "", "limitations": "", "rating": 6, "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 introduces a connection between GFlowNets and variational inference (VI) algorithms for hierarchical variational models (HVMs), demonstrating that special cases of training HVMs via VI are equivalent to training GFlowNets via the trajectory balance (TB) objective. This is interesting because GFlowNets can be shown to automatically perform variance reduction in gradient estimation for reinforcement learning (RL), showing that GFlowNets may be helpful for RL settings where VI would have been used instead.", "summary_of_the_review": "The paper provides a theoretical and empirical investigation into the relationship between hierarchical VI and GFlowNets, outlining when they are equivalent and settings when one should be preferred over the other. ", "strength_and_weaknesses": "Strengths:\n- I liked the experiments the authors conducted which compared the empirical behaviors of various training objectives in practice. Understanding the settings in which the objectives differ and figuring out when each method performs better or worse than others is important, and I thought the authors did a good job on this on the synthetic task. It was also nice to see it scaled to slightly harder problems.\n- I thought the connection to HVMs and wake-sleep algorithms was interesting, and could open up new research areas along this direction.\n\nWeaknesses:\n- The paper was a bit hard to follow at times.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": {}}, {"review_id": "-9TeE0BoVJ", "reviewer": "Reviewer_taL3", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "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 authors discuss the relationship between specific families of variational inference approaches (hierarchical variational inference) with generative flow networks as introduced recently by Bengio et al. (2021). The relation refers to equality in the expected gradients of the respective objectives. \nThe authors evaluate this theoretical relationship further in a set of three experiments of various domains, discussing both similarities and differences.\n", "summary_of_the_review": "A well justified, well written, and well evaluated paper.", "strength_and_weaknesses": "## Strengths\n- The paper demonstrates a novel connection between VI and GFlowNet approaches\n- All statements and empirical claims are clearly stated and demonstrated\n- A diverse set of experiments in different domains is provided\n\n## Weaknesses\n- A, rather minor, weakness is the lack of a stronger evaluation of the connection into the continuous domain, demonstrating the usefulness of the proposed new point of view compared to common VI approaches. \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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster