Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization
Authors: Dinghuai Zhang
Ricky T. Q. Chen
Cheng-Hao Liu
Aaron Courville
Yoshua Bengio
Venue: ICLR 2024 poster (2024)
Keywords: probabilistic inference
sampling
stochastic optimal control
gflownets
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: DGFS is an algorithm which learns a stochastic process to sample from unnormalized densities, and can update parameters without full specification of diffusion chains.

Abstract: We tackle the problem of sampling from intractable high-dimensional density functions, a fundamental task that often appears in machine learning and statistics. 
We extend recent sampling-based approaches that leverage controlled stochastic processes to model approximate samples from these target densities.  
The main drawback of these approaches is that the training objective requires full trajectories to compute, resulting in sluggish credit assignment issues due to use of entire trajectories and a learning signal present only at the terminal time.
In this work, we present Diffusion Generative Flow Samplers (DGFS), a sampling-based framework where the learning process can be tractably broken down into short partial trajectory segments, via parameterizing an additional ``flow function''.
Our method takes inspiration from the theory developed for generative flow networks (GFlowNets), allowing us to make use of intermediate learning signals.
Through various challenging experiments, we demonstrate that DGFS achieves more accurate estimates of the normalization constant than closely-related prior methods.

Published as a conference paper at ICLR 2024
DIFFUSION GENERATIVE FLOW SAMPLERS: IMPROVING LEARNING SIGNALS THROUGH PARTIAL
TRAJECTORY OPTIMIZATION
Dinghuai Zhang 222; Ricky T. Q. Chen", Cheng-Hao Liut.,,
Aaron Courville & Yoshua Bengio!.2
Mila- Quebec Al Institute, Université de Montréal, SFAIR, Meta AI, "McGill University
ABSTRACT
We tackle the problem of sampling from intractable high-dimensional density functions, a fundamental task that often appears in machine learning and statistics. We extend recent sampling-basee approaches that leverage controlled stochastic processes to model approximate samples from these target densities. The main drawback of these approaches is that the training objective requires full trajec- tories to compute, resulting in sluggish credit assignment issues due to use of entire trajectories and a learning signal present only at the terminal time. In this work, we present Diffusion Generative Flow Samplers (DGFS), sampling based framework where the learning process can be tractably broken down into short partial trajectory segments, via parameterizing an additional "flow function". Our method takes inspiration from the theory developed for generative flow networks (GtlowNets), allowing us make use of intermediate learning signals. Through various challenging experiments, we demonstrate that DGFS achieves more accu- rate estimates of the normalization constant than closely-related prior methods.
INTRODUCTION
While diffusion models for generative modeling can be trained with efficient simulation- free training algorithms, their application in solving sampling problems-ve where no data given but only an unnormalized density function-remain, challenging. This is very general problem formulation that appears across a diverse range of research fields including machine learning (ML) (Duane et al., 1987: Neal, 1995; Hernández-Lobato & Adams, 2015) and statisties (Neal, 1998; Andrieu et al., 2004: Zhang et al., 2021), I), and has wide application scientific fields such as physics (Wu et al., 2018; Albergo al., 2019; Noé et al., 2019) and chemistry (Frenkel & Smit, 1996; Hollingsworth & Dror, 2018; Holdijk et al., 2022).
Two main lines of work tackle this sampling problems are Monte Carlo (MC) methods and vari- ational inference (VI). MC methods achieve sampling typically via a carefully designed Markov chain (MCMC) or ancestrally sample from sequential targets such as series of annealed distribu- tions (Moral & Doucet, 2002). Despite their wide application, MC methods can easily get stuck in local modes or suffer from long mixing time (Tieleman Hinton, 2009; Desjardins et al., 2010). On the other hand, VI methods (Blei et al., 2016) use a parametric model to approximate the target distribution by minimizing some divergence. However, these models often possess limited expres- siveness, and optimizing the divergence can lead to bad learning dynamics (Minka, 2001).
Recent works have formulated sampling from unnormalized density functions as stochastic optimal control problems, where a diffusion model is trained as stochastic process to generate the target distribution. At present, the difficulty using these algorithms in part due to the need to sequen- tially sample long trajectories during training, while only the final step is actually treated as sample from the model. The popular Hellback-Leibler (KL) problem formulation only provides a learning signal for the terminal state, which can cause credit assignment problems in learning, especially for diffusion models which require fairly long sequential chains for sampling.
"Correspondence di ngguaii zhang@mi quebec

---

Published as a conference paper at ICLR 2024
In this work, we build top of Lahlou et al. (2023), who noted that recent methods that train dif- fusion models given an unnormalized density function (Zhang & Chen, 2022; Vargas et al., 2023a) can be interpreted from a GFlowNet perspective (Bengio et al., 2023). Although most GFlowNet works perform amortized inference on combinatorial discrete domains, previous work (Lahlou et al., 2023) establishes theoretical foundation for continuous space GFlowNett and conducts numerical analysis to show their feasibility. Based on this framework, we propose the Diffusion Generative Flow Sampler (DGFS), an enhanced training method of diffusion models for sampling when given an unnormalized density function. In particular, we formulate a training approach that learns inter- mediate "flow' functions" that aggregate information for training intermediate states. This allows us to further define partial vrajectory-based training objectives, where we have access to only partial specification of the sampling path. We find that this approach significantly reduces gradient noise and helps stabilize convergence. Furthermore, our training framework enables us to receive learn- ing signals at intermediate steps, even before the sequential sampling is completed, in contrast to previous methods that can only receive signals at the terminal steps.
In summary, our contributions are as follows,
We propose DGFS, an effective algorithm that trains stochastic processes to sample from given unnormalized target densities.
We show that, contrary to other diffusion- based sampling algorithms, DGFS can (1) update its parameters without having full trajectory specification; (2) receive intermediate signals even before sampling path is finished.
We show through extensive empirical study that these special advantages enable DGFS to benefit from more stable and informative training signals, with DGFS generating samples accurately from the target distribution.
PRELIMINARIES
2.1 SAMPLING STOCHASTIC OPTIMAL CONTROI.
We aim to sample from D-dimensional target distribution x(-) in R defined by a given unnormal- ized target density pa(-) where (..) ((1 and Z p(x) dx. Following similar setup to prior works such as Tzen & Raginsky (2019); Zhang & Chen (2022); Vargas et al. (2023a), we consider
a sequential latent variable model with joint distribution denoted Q of which samples follow Markov process,
Q(x....... .xn) X,+1 PF(+(x,), Xo "P (xo)-
(1)
Through this sequence model, we use the marginal distribution Q(xx) the terminal step N to approximate xx Here (.) is reference distribution at time = O, and (- denotes forward transition probability (or policy) distribution. This forward policy depends on n, i.e. it is pn+ (-Ix,). but we omit the step index n out of simplicity or assume n is encoded in X,- We thus refer to Q as the forward process in the following text?. In this work, we constrain the transition to be a unimodal Gaussian distribution:
P=(+)x)) = (- + hf(x,,, n), hoo
(2)
where f is control drift parameterized by neural network and h is constant step size. We also define a reference process with the forward policy prf (-Ix,) = N(; ho
(x.......
Xnt1 (-
Po (xu).
Since prff has no nonlinear f, the marginal distributions of Qref at step n are known in closed form, denoted (-). We can then construct target process that corresponds to this reference process,
π (xXN)
P(xo..... xN) == Q
X.)
(4)
(NN)
'n this work, the term "flow" specifically denotes the flow network structure in GFlowNets, which has nothing do with normalizing flow models.
"Note that this is consistent with prior GFlowNet literature, though it is the opposite direction considered by standard diffusion model literature in generative modeling, where "forward process" is colloquially used to denote the process from clean sample white noise.
2

---

Published as a conference paper at ICLR 2024
It can be shown that the marginal distribution of P at the terminal time N is exactly proportional to the target (u(-) by marginalizing out the variables Xo..... XN-1 1 in P; see Theorem of Zhang & Chen (2022). Hence it is natural consider learning the drift f(. -) in Equation 2 by minimizing the KI divergence DxI. (Q)(P) between the forward process and this target process, which is equivalent to the following optimal control problem:
((xo:n)
Xn+1 = Xx, + hx n)) + Vhoen.
xo X0~~ pb)
(5)
min Ee
N
Σ
h
[lf m)ll + log ((xx)
(D(() = -((())
(6)
20
ml)))
where C, ~ (O, Vn. See Appendix B.1 for derivation.
If we keep the total time T hN fixed and push the number of steps N + co, this formulation recovers the cntinuous-time stochastic optimal control formulation seen in prior works (Zhang & Chen, 2022):
2(x(O.T): dx, f(x;,t) dt dd Xo ~ po (xo);
T 1 J' 20
min Eg
II((x 2)1 log (lgg((x))))))
(8)
where the forward process is defined by variance exploding stochastic differential equation (VE- SDE)', and dW, is a standard Wiener process. This formulation appears in standard stochastic optimal control theory (Kappen, 2005). The seminal path integral sampler (PIS) (Zhang & Chen, 2022) proposes to solve such optimal control problems with deep learning to tackle the sampling problem. Its follow-up, the denoising diffusion sampler (DDS) (Vargas al., 2023a), takes simi- lar approach but adopts a variance preserving (VP) SDE rather than the variance exploding ones in Equation 5 and 7. We defer the derivation of VP modeling Appendix B.2. The practical imple- mentations of these prior works are performed in discrete time with constant step sizes, equivalent the discrete formulation outlined above. Hence we also adopt the discrete- -ime formulation in the following sections for simplicity and better amenability to our proposed algorithm.
2.2 IFLOWNETS.
Generative flow networks (Bengio et al., 2023, GFlowNets) are a family of amortized inference algorithms that treat the problem of generation with probability proportional to given rewards in an Markov decision process (MDP) sequential decision decision-making making process.
Let G = (S, A) be a directed acyclic graph (DAG), where the set of vertices s E S are referred to states, and the set of directed edges (s s') A C S S represent the transitions between the states and are called actions. Given a designated initial state So with no incoming edges and a set of terminal states NN without outgoing edges, a complete trajectory can be defined as sequence of states 7 = (So + S1 + + SN) E T. We define the forward policy Pr(s'(s) [s) as a distribution over the children of every non-terminal state s, which induces distribution over complete trajec-
torics, (F( = _ Pp(sn)1|sp)) This in turn defines a marginal distribution over terminating states in the sense that Pr = E PF(T), where we call Pr(-) the GFlowNet's terminating distribution. The aim of GFlowNet learning is to obtain a policy such that Pr(()) oc R(-). where R(.)
R+ is given reward function or unnormalized density, where we do not know the normalizing factor Z = Ex R(x). We also use the trajectory flow function F(r) = (P)(T) to incorporate the effect of the partition function. Accordingly, the state flow function is the marginal- ization of the trajectory flow function (s)) Erss F(T) S + R,, Most previous Gel works are about composite discrete objects such molecules; on the other hand, Lahlou et al. (2023) provide theoretical support and initial numerical demonstration for GFlowNets on continu- ous and hybrid (continuous and discrete) spaces. Based on this, our work aims to further tackle more complex sampling problems in continuous domains, i.e. S C R.D.
To avoid confusion, we would like to remark again that in literature the term "forward process" is usually used describe the noising pracess, which is actually variance shrinking in PIS; Equation 21.
3

---

Published as a conference paper at ICLR 2024
Detailed balance (DB). DB is one of the most important GFlowNet training losses. Apart from the forward policy and the state flow function, DB also requires the practitioners to parameterize an additional backward policy (s)s'' O), which is distribution over the parents of any noninitial state s' E S. The DB objective is defined for single transition (s + s') as follows,
PD (s. ;); = log F(s; P) 2
F(s': Ps(ss'
According to the theory of GFlowNets (Bengio et al., 2023), if the DB loss reaches for any transi- tion pair, then the forward policy samples correctly from the target distribution defined by R(.))
3 DIFFUSION GENERATIVE FLOW SAMPLERS
3.1 AMORTIZING TARGET INFORMATION INTO INTERMEDIATE STEPS
Like many other probabilistic methods, the training objective of PIS is KL divergence between the model and the target distribution (Equation 5). In particular, PIS's KL is calculated at the trajectory
level, which means each computation relies on a complete trajectory т (Xo.... ......... xx) sampled from the current model's forward process. Since the training signal (x(x) is only affected by the terminal state XN. long trajectories can potentially yield less informative credit assignment for the intermediate steps, which manifests as high gradient variance and can sometimes lead to more difficult optimization. Therefore, we propose to exploit intermediate learning signals that do not necessarily rely on the complete trajectory information. What if we could know approximate the target distribution at any step n? Note that the target process (Equation 4) can be decomposed into product of conditional distributions
N-1
P(XD:N) = (xN)
II ===
PB (x+1)))))
(10)
where (-..l.) denotes the backward transition probability (or backward policy) derived from the joint of the target process. Then we have the expression for the marginal distribution at step
N-1
Px() / T(xN)
II (x(x((x(1)) dx. Vn E (O....,N); NN,
(11)
where we set (-) = x-)) for consistency. If we know the form of (-) then we could directly learn from it with shorter trajectories and thus achieve more efficient training. Unfortunately, although know the form of (-I-),), for general target distribution there is generally no known analytical expression for it. As result, we propose to a deep neural network F, F= 0) with parameter 0 as a "helper" to approximate the unnormalized density of the n-th step target p (.).
One naive way to train the (: 0) network is to fit its value the MC quadrature estimate of the integral in Equation 11. However, every training step would then have to include such computation- ally expensive quadrature calculation, making the resulting algorithm extremely inefficient. Instead, we propose turning this quadrature calculation into an optimization problem in an amortized way, and instead train Fx 0) to achieve the following constraint:
N-1
N-1
F. (x,:; 0) II PF x;; = (((x)) II Pn(xi(x(41), t=n
(12)
I=
for all partial trajectories Ny:N- Here, we would parameterize both the function F,, and the distri- bution Pp from Equation 2 by deep neural networks. Once Equation 12 is achieved. integrating X.+1:N on both sides would lead to the fact that F, (xx;;) 0) equals the integral in Equation 11. To put it another way, (=(:;:) amortizes the integration computation into the learning of A. Notice that we only use ua()) in Equation 12 because our problem setting provides only μ.. The unknown normalization constant is thus also absorbed into Fn. We construct a training objective by simply regressing the left hand side of Equation 12 to the right hand side. Furthermore, to ensure training stability, this regression is performed in the log space. In practice, the flow function at different steps share the same set of parameters; this is achieved by introducing an additional step embedding input for the F(, n; 0) neural network. As a result, our method learns two neural networks (the forward policy, which is the same as in PIS, and the flow function) in the meantime.
4

---

Published as a conference paper at ICLR 2024
3.2 UPDATING PARAMETERS WITH INCOMPLETE TRAJECTORIES
If we compare Equation 12 with n. and n + 1, we obtain a formula with no dependence on Jr:
F (x 0) (xn++1//Xx 0) = where the details are deferred to Appendix C.I.
(13)
GFlowNet perspective. This formulation is similar to the GFlowNet DB loss in Equation 9. Ac- tually, is not hard to see that the above modeling is a variant of GFlowNets. As Zhang et al. (2022a) has pointed out and as later adopted by Lahlou et al. (2023), there is a connection between diffusion modeling and GFlowNets. In the GFlowNet language, the target density ju(-) is the termi- nal reward function, the temporal-spatial specification space {(n, x,,) : n X,, R RD) is the GFlowNet state space S, the forward transition probability P=(-I-) of the diffusion process Q is the GFlowNet forward policy, the backward transition probability (--(-)) of the reference process Q is the GFlowNet backward policy, and the helper network (-) is the state flow Fss). Notice here that the backward policy is fixed and does not contain learnable parameters. In our setting, we use the step index to distinguish between intermediate states {(n, x,) n < N, Xn E RD] and ter- minating states {NN xN), XN RD What's more, the goal of GFlowNets, which is to generate terminal states with probability proportional to the reward value, matches the goal of sampling in our setting.
It is thus worthwhile to investigate how GFlowNets could be trained in this context. Apart from the DB constraint, Madan et al. (2022) propose a more general subtrajectory balance constraint: Fm.(Xm: 0) I (x xx; 0) = Fn (x,;; 0) IT== (x/xx/x(41) where Xmen is a subtrajectory. Notice that Equation 12 is actually a special case of this constraint if we specify
N. The
resulting objective for our method is
0) IT (x [x 0) 2
F (Xm:
ESUTTB (Xm:ni 0) = log
(
(14)
(x,: 0) Pa (x) (x
[T=
A training loss based on transitions and more generally partial trajectories makes it possible to update parameters without entire trajectory specification. This can significantly reduce the variance in the optimization, as shown by Madan et al. (2022) in a similar context. To better combine the training effect from subtrajectories with different lengths, the adopted objective is
[(T; 0) = Eosmenen "n-""SuNTT (Xm=). , (xo: XN),
(15)
LosmenEn An-m
where À R4, is a scalar for controlling assigned weights to different partial trajectories. For the modeling of the drift of the forward stochastic process, we use f(: n)/a = N n)), +NN2(n) V log u(-), which includes two different neural networks and gradient information from unnormal- ized target densities in a similar form to Langevin dynamics (Zhang & Chen, 2022). We have con- duct extensive ablation study about the parameterization, training objectives, exploration techniques, and other considerations, and display them in Appendix C.
Improved credit assignment with local signals. In both PIS and DDS algorithms, the only learning signal μ(xx) comes at the end of complete trajectories Xo:N when the density func- tion is queried with XN. This is also the case for general GFlowNet methods, where we set (x + μ(x), Vx E DRD. In this way, the propagation of density information from termi- nal states to early states would be slow, thus affecting the training efficiency. We thus adopt the forward-looking trick proposed by Pan et al. (2023a) to incorporate intermediate learning signals for GFlowNets. The forward-looking technique essentially parameterizes the logarithm of the state flow function as log Fx. 0) = log R., (xx) + NN(x,; 0), where R, (x,,) could be an arbitrary hoowward-looking
estimate of how much energy x,, could contribute the final reward ,((xw), The details of the "NN" neural network are deferred to Appendix C.4. In the proposed DGFS algorithm, take the form of the (log) partial reward to be
log R-)()) = (1 3)
logr + & log ((),
(16)
which combines information from the target density and the reference marginal density.
Without loss of generality, we slightly overuse the notations and simply use x,, denote x,).
5

---

Published as a conference paper at ICLR 2024
R(xm)
)(x,)
16"
Micthod
DGFS
P15
Gedeeitvalonne
10"
10"
10"
Fo(X,)
F,(X,)
Fx(Kx) = ((x,))
10"
200
400
G30
800
1000
Training steps
Figure 1: Illustration of the DGFS algorithm. The proposed method can update from partial trajectories (colored seg- ments) as well as intermediate signals (gray arrows).
Figure 2:
Gradient variance of
DGFS and PIS, explaining the better training effects of DGFS.
We illustrate the resulting DGFS algorithm in Figure Different from prior works, DGFS can update parameters with incomplete subtrajectories (colored curve segments in the figure) with the help of amortized flow functions, as well as intermediate signals at non-terminal step (grey arrows in the figure), which fundamentally improves the credit assignment process.
Estimation of the partition function. For an arbitrary trajectory 7 = {x, J. one could define its (log) importance weight to be S(;; log P(xu.v) log ((x::v; 0), where Q(xo:v; 0) ph (xo) II==0 1x 0). We obtain the following log- tog-partition function (log Z) estimator:
B
log Σ exp (S(r()))) b=1
log B S log Z, 7 (() 2(=;)),
(17)
where B E N the number of particles (i.e., trajectories) used for the estimation. Similar to the importance weighted auto-encoders (Burda et al., 2016), our estimator provides a valid lower bound of the logarithm of the normalizing factor with any forward process Q. In this work, we use the difference between estimated log partition function and the ground truth See the Appendix C.2 for additional analysis. We also remark that this approach is similar to the neural importance sampling in Nicoli al. (2020); Wimsberger al. (2020) in the sense that use the forward process Q as the proposal distribution.
evaluation metric.
3.3 DISCUSSION
Variance of gradient updates. Figure shows that the variance of stochastic gradients with DGFS is much smaller than that with PIS: notice that the comparison is valid as we take the same neural network architecture for both algorithms. We hypothesize that DGFS benefits from the low variance behavior due to the following two reasons: (1) From the perspective of temporal difference learn- ing (Sutton & Barto, 2018; Kearns & Singh, 2000), it has long been discussed that policy learning from shorter trajectory data would lower the variance at the cost of potentially higher bias. Before this work, it was impossible for sampling with stochastic process algorithms to update with partial trajectories, while in the proposed DGFS method, feasible. In Figure 3 we also show that the flow network can be properly learned and thus does not introduce much bias, indicating that DGFS achieves a better bias var-vvaiance trade-off. (2) Probabilis- tic ML researchers have shown that for KL divergence based distribution learning, the gradient will not vanish even if the optimal distribution has been reached; see Roeder al. (2017) and Section 4 of Xu et al. (2022). Later works including Vaitl et al. (2022a;b) propose a different path-gradien estimator tackle this issue with a similar insight with Roeder al. (2017); see the related work section. On the other hand, the DGFS objective does achieve zero gradient with the optimal solution (c.f. Appendix C.3), making a more stable algorithm.
an additional flow network to make this
Modeling considerations. Notice that the DGFS derivation in Section 3 does not rely the specific form of the stochastic process. Therefore, DGFS could take both the variance exploding (VE) (Zhang& & Chen, 2022) variance preserving (VP) (Vargas et al., 2023a) formulation. In our numerical study, we find the former to be more stable, and thus we default to it unless otherwise specified. We defer related ablations the Appendix C.5.

---

Published as a conference paper at ICLR 2024
Table 1: Experimental results on benchmarking target densities. For each algorithm, we report its absolute estimation bias for the log normalizing factor and the corresponding standard deviation. Our method achieves the best performance among the diffusion modeling- based samplers.
MoG
FUNNEL
MANYWELL
VAE
CoX
SMC
0.289
0.307
22.36
14.34
99.61
=8.382
VI-NF
L0.47
0.272
2.677 #0.016
6.961 +2.501
83.49 =2,434
CRAFT
0.348
0.454:t0.488
0.987
0.521+0.200
13.79=2.351
FAB w/ BUFFER
0.003 0.0005
0.0022 10.000
0.0322..000
N/A
0.19
PIS
0.036=00007
0.305
1.391=1.011
2.049
DDS
0.028=0.013
0,416=#0094
10.54
1.740
N/A
DGFS
0.009
0..04:10.014
00.047
10.008
8.974=1.169
Convergence guarantees. Previous theoretical studies (Bortoli, 2022; Chen et al., 2022; Lee et al., 2022) have demonstrated that the distribution of terminal samples generated by diffusion model can converge to the target distribution under mild assumptions if the control term well learned. These findings are also applicable DGFS since these proofs independent of how the net- works are trained. Additionally, Zhang et al. (2022a) have proven that a perfectly learned score term corresponds to zero GFlowNet training loss under mild conditions, following theory of GFlowNets (Bengio et al., 2021; 2023). This provides assurance that a sufficiently well-trained DGFS can accurately sample from the target distribution.
4 RELATED WORKS
Sampling methods. MCMC methods (Brooks et al., 2011) are the most classical sampling al- gorithms, including Langevin dynamics (Welling Teh, 2011). Hamiltonian Monte Carlo (Duane et al., 1987; Hoffman & Gelman, 2011; Chen et al., 2014), and others (Huang et al., 2023). With MCMC operators as building blocks, sequential MC methods (Liu,, 2001; Doucet al., 2001) gen- erate samples ancestrally through a series of annealed targets. On the other hand, people turn to parametric models such as deep networks to learn samplers via variational inference (Rezende & Mohamed, 2015). One the most commonly adopted models is the normalizing flow (Dinh et al., 2014) due to its tractability and expressiveness. Various work have been conducted to improve nor- malizing flow based sampling (Müller et al., 2018; Gao et al., 2020; Wu et al., 2020; Chen al., 2020; Gabri'e e et al., 2021; Arbel et al., 2021; Midgley et al., 2022; de G. Matthews et al., 2022; Caselle al., 2022; Gao et al., 2023). These methods support asymptotically unbiased estima- tion as shown in Nicoli et al. (2020). What's more, Vaitl et al. (2022a) propose to to improve the gradient estimation of normalizing flow parameters with the path-gradient technique; Vaitl et al. (2022b) improve continuous normalizing flows based methods (Chen et al., 2018) with similar technique. These methods have wide applications natural sciences, for example in lattice field theories (Nicoli et al., 2021; 2023) and targeted free energy estimation (Wimsberrgr et al., 2020).
Recently, many researchers have developed powerful diffusion process to model distributions in the generative modeling community (Vincent, 2011; Sohl-Dickstein et al., 2015; Ho et al., 2020a; Song et al., 2020; Du et al., 2022). This inspires people to study how to utilize diffusion process as samplers to learn distribution from given unnormalized target densities (Tzen& & Raginsky, 2019; Zhang & Chen, 2022; Vargas et al., 2021; Geffner & Domke, 2022; Doucet et al., 2022; Vargas et al., 2023a; Richter et al., 2023; Vargas & Nusken, 2023; Vargas et al., 2023c;b). Our work also falls into this category. Specifically, Vargas et al. (2023a) can be seen as a special case of the method proposed in Berner et al. (2022).
GFlowNets. GelowNet (Bengio et al., 2023) is family of generalized variational inference al- gorithms with a focus on treating the sampling process as a sequential decision- makiing process.
'We used the fab-jax results. The FAB method utilizes the annealed importance sampling computation during the training period.
'Due to time constraints, we are unable to reproduce DDS performance using PyTor for this task. Instead, we encountered "not number", probably due to mismatch in implementation details or other unspecified difference between PyTorch and Jax. The DDS results reported in its original paper can be found here.

---

Published as a conference paper at ICLR 2024
It is original proposed sample diversity-sceking objects in structured scientific domains such as biological seugence and molecule design (Bengio et al., 2021: Jain et al., 2022: 2023b;a; Shen et al., 2023; Zhang et al., 2023e; Mernández-Garcih et al., 2023; Kim et al., 2023). However, people have developed GelowNet theories and methodologies to be suitable for generating general objects with composite structures. GFlowNet has rich connection with existing probabilistic MI. meth- ods (Zhang et al., 2022b; Zimmermann et al., 2022; Malkin et al., 2023; Zhang et al., 2022a; Hu al., 2023; Ma et al.). Malkin al. (2023) also show that GFlowNet has the capability of training with off-policy data compared to previous amortized inference methods. As policy learning frame- work under Markov decision processes (MDPs), researchers have also drawn fruitful connections between GetowNet and reinforcement learning (Pan et al., 2023b;c;a). GelowNet has also been demonstrated to have wide application in causal discovery (Deleu et al., 2022; 2023; Atanackovic et al., 2023), feature attribution (Li et al., 2023a), network structure learning (Liu et al., 2022), job scheduling (Zhang et al., 2023a), and graph combinatorial optimization (Zhang et al., 2023b). Most prior works focus structured discrete data space. Lahlou et al. (2023) first provides sound the- oretical framework and initial numerical demonstration for using GFlowNet on continuous space. Our work follows this framework and scales up the performance. A prior study (Li et al., 2023b) initially explored GelowNets in a continuous space, whereas their theory is demonstrated to have non-trivial flaws when considering integration under continuous measure indicated by Lahlou et al. (2023).
EXPERIMENTS
5.1 BENCHMARKING TARGET DISTRIBUTIONS
In this section, we conduct extensive experiments on a diverse set of challenging continuous sam- pling benchmarks corroborate the efffcctiveness of DGFS against previous methods.
Mixture of Gaussians (MoG) is 2-dimensional Gaussian mixture where there are modes designed to be well-separated from each other. The modes share the same variance of 0.3 and the means are located in the grid of (-5, 0, 5} x (-5,0, 5,
Funnel is classical sampling benchmark problem from Neal (2003); Hoffman & Gelman (2011). This 10-dimensional density is defined by u(x) = (N(x(()) 0, (x(()((()))))) 0, exp(x:())I)) II)
Manywell (Midgley al.. 2022; Noé et al., 2019; Wu al., 2020) o2-dimensional distribu- tion defined by the product of 16 same 2-dimensional "double well" distributions, which share the unnormalized density of μ(,,, y) = exp( + 6x + 0.5z 0.5y²).
VAÆ' is a 30-dimensional sampling task where we use the latent posterior of a pretrained variational autoencoder (Kingma Welling, 2014; Rezende al., 2014, VAE) to be the target distribution.
Cox denotes the 1600-dimensional target distribution of the log Gaussian cox process introduced in Møller et al. (1998) which models the positions of pine saplings. Its unnormalized density is (x(x) = NXx; /4, K-1) I (xp(x (x),)), where μt, K, y, a are constants.
Baselines. We benchmark DGFS performance against a wide range of strong baseline methods. As an MCMC methods include the sequential Monte Carlo sampler (Moral & Doucet, 2002, SMC), which is often considered the sfate-of-the-art sampling method. As for normalizing flow methods include the variational inference with normalizing flows (Rezende Mohamed, 2015, VI-NF) approach as well as more recent de G. Matthews et al. (2022, CRAFT) which is combined with MCMC. We also provide results of Midgley al. (2022, FAB) with a replay buffer that utilizes the annealed importance sampling technique to augment the training. Last but not least, we compare with PIS (Zhang& Chen, 2022) and DDS (Vargas et al., 2023a) algorithms which are closest to our framework but trained with the KI. divergence formulation (Equation 5) without the improved training techniques described in Section 3.
Evaluation protocol. We measure the performance of different algorithms by their estimation bias of the log normalizing factor of unnormalized target distributions. Contrary to Zhang & Chen (2022)
Previous works (Zhang& & Chen, 2022; Lahlou al., 2023) unintentionally take the variance of (0) to be 1 for their methods (hence much less difficulty) hut 9 for other methods.
8

---

Published as a conference paper at ICLR 2024
True
PIS
DGFS
DDS
Figure The learned DGFS flow function and the ground truth samples from target process at different diffusion steps. This shows DGFS is able learn flow functions correctly.
Figure Manywell plots. DGFS and
DDS but not PIS recover all modes.
which take the best checkpoint across the run, in this work compute the average of the last ten checkpoints for a more robust comparison. All results are conducted with 5 random seeds.
We display the benchmarking results in Table 1. We fail achieve reasonable performance with DDS on Cox target where we mark it as "N/A". The advantage of the proposed DGFS method be observed across different tasks, indicating it could indeed benefit from its design providing more informative training signals to achieve better convergence and capturing diverse modes even in high dimensional spaces. FAB with buffer produces highly stable results, which can be attributed to both annealed importance sampling, the use of buffers, larger networks, and other tricks. It is also remarkable that DGFS obtains very stable results compared to other baselines with diffusion modeling (PIS, DDS). We defer related experimental details and extra experiments to Appendix C.
5.2 ANALYSIS
First, we investigate the learning of the flow function introduced in Equation 12. In the first row of Figure 3, we visualize the learned flow function (x,) with n = 20, 40.... 100 in the 2D MoG task. In the second row, simulate samples from pr (x,) by sampling according Equation 11, ie.. moving ground truth (-)) samples in the backward direction with policy PB. The consistency between these two rows indicate that DGFS flow function can successfully approximate the inter- mediate marginal (-), which makes it possible to provide an effective training signal for learning from partial trajectory pieces.
For the MoG task, we demonstrate the samples together with ground truth contours in Figure 5. This visualization result matches the log partition function estimation bias in Table where DGFS generates the nine modes more uniformly than other baselines. This corroborates the advantage of DGFS to capture
different modes in a complex landscape. We then show visualization results on the Manywell task in Figure 4. The Manywell distribution is 32-dimensional, joint of the first and third dimension. We plot the samples generated from the true distribution and different algorithms, together with the true density contours. This figure implies that PIS model half of the modes in these two dimensions well, but miss the other two modes, whereas both DGFS and DDS could cover all four modes. Further, DGFS achieves slightly better mode capturing than DDS in the sense that its modes are more separate and thus closer to the true distribution, while for DDS there mode, which is not ideal. We defer other visualizations to the Appendix C.5.
display the
considerable samples appear between the upper right mode and the bottom right
To look into why DGFS out- performs PIS with similar ar- chitecture, we plot the learned drift network magnitude in Fig- ure the Appendix. The re- sult shows that the solution that DGFS achieved has more ho- mogencous magnitude outputs across different diffusion steps. This may explain one reason
DGES
PIS
DDS
Figure 5: MoG visualization of DGFS and other diffusion-basse samplers shows that DGFS could capture the diverse modes well. The contours display the landscape of the target density.

---

Published as a conference paper at ICLR 2024
why DGFS is easier to train than PIS, since the policy probabili- ties at different steps (-x(-Ixa)}, are actually parameterized by single neural network with tem- poral embedding input. This makes the learning task more stationary with DGFS. We also conduct in-depth ablation studies evaluate the design choices of DGFS in Appendix C.5 and C.7.
6 CONCLUSION
We propose the Diffusion Generative Flow Sampler (DGFS), novel algorithm that trains diffusion models to sample from given unnormalized target densities. Different from prior works that could only learn from complete diffusion chains, DGFS could update its parameter with only partial spec- ification of the stochastic process trajectory; what's more, DGFS can receive intermediate signals before completing the entire path. These features help DGFS benefit from efficient credit assignment and thus achieve better partition function estimation bias in the sampling benchmarks.
The DGFS framework presents various opportunities for further research, fueled by its inherent lim- itations. To name a few, how can we better design the intermediate local signals in ways more so- phisticated than the current straightforward approach? Additionally, can we harness DGFS's explo- ration capabilities to succeed in high dimensional tasks, pushing the boundaries of its performance? Can we combine DGFS with a prioritized replay buffer? This would largely reduce the number of querying pa-). From the perspective of application, there are also many interesting research ques- tions. Can we deploy DGFS on scientific biological or chemical tasks such as protein conformation modeling, with equivariant modeling? We expect that exploring these future directions will lead to valuable insights and bring fruit follow-ups of DGFS.
ACKNOWLEDGEMENT
The authors are grateful the discussions with Francisco Vargas, Julius Berner, Lorenz Richter, Laurence Midgley, Minsu Kim, Yihang Chen, and Qinsheng Zhang.
ETHICS STATEMENT
We commit to the ICLR Code of Ethics and affirm that our work only utilizes public datasets for experimentation. While our empirical results are largely based on synthetic data, we acknowledge the potential for misuse and urge the responsible application of the proposed methods with real- world data. We welcome any related discussions and feedback.
REPRODUCIBILITY STATEMENT
We provide detailed algorithmic and experimental description in Section 3 and Appendix C, and we have open sourced the code accompanying this research in this GitHub link.
REFERENCES
Michael S Albergo, Gurtej Kanwar, and Phiala E. Shanahan. Flow-based generative models for markov chain monte carlo in lattice field theory. ArXiv, abs/1904. 12072, 2019. URL https: //api. www.nticschoottoooe.oe.e
http:/CoppsI:::111000
Christophe Andrieu, Nando de Freitas, A. Doucet, and Michael I. Jordan. An introduction to meme for machine learning. Machine Learning, 50:5-43, 2004. URL https://api.
semanticscholas wrw.:a//oooo/oooo....... 383363.
Michal Arbel, Alexander G. de G. Matthews, and A. Doucet. Annealed flow transport monte carlo. ArXiv, abs/2102.07501, 2021. URL https:/ / /api. semanticscholee org/ CorpusTD:23119555
10

---

Published as a conference paper at ICLR 2024
Lazar Atanackovic, Alexander Tong, Jason S. Hartford, Leo J. Lee, Bo Wang, and Yoshua Ben- gio. Dyngfn: Bayesian dynamic causal discovery using generative flow networks. ArXiv, abs/2302.04178, 256662181.
2023. URL https:// //ap:://ai.... wew..ecccoooo.oooooooooooooooo
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), 2021.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, (24):1-76, 2023.
Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. ArXiv, abs/2211.01364, 2022. URL https://api.
semanticscher http://cccccoooooooooooooooostas
wttg/oorpusseeooooee ID: 253255370.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statis- ticians. ArXiv, abs/1601.0066 2016. URL sttps://api..o http:////app.semmat.cccooooooooo
CorpusID: 3554631.
Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypoth- esis. ArXiv, abs/2208.05314, 2022. URL
http://pp...mmmtcscch.....eeeeee .rrg/
CorpusID: 2514122251588
Steve P. Brooks, Andrew Gelman, Galin L. Jones, and Xiao-Li Meng. Handbook of markov chain monte carlo. 2011. URL https: /taps://api..../ /tp.. www.steeemeoeeooeeeee.... www.o//rg//oo.D/// 60438653.
Yuri Burda, Roger Baker Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders.
International Conference on Learning Representations (ICLR), 2016.
Michele Caselle, Flia Cellini, Alessandro Nada, and Marco Panero. Stochastic normalizing flows non-equilibrium transformations, Journal of High Energy Physics, 2022, 2022. URL https: org/ wwtp://pr....ooooooooooooa. ID: 4246445579
//api. www.seeenntnnnolseeeee
Sitan Chen, Sinho Chewi, Jungshian Li, Yuanzhi Li, Adil Salim, and Anru R. Zhang. Sam- pling is as easy as learning the score: theory for diffusion models with minimal data as- sumptions. ArXiv, 0bs/2209.11215, 2022. URL https:://///.. sematisstmtttiscoolll.ee
foXpusID:252438990
Tian Qi Chen, Yulia Rubanova, Jesse Bettencourt, and David Kristjanson Duvenaud. Neural or- dinary differential equations. In Neural Information Processing Systems, 2018. URL https: //api. cormmnitccchhlle
Arrg/CorpusID...
:93104444
Tianqi Chen, Emily B. Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International Conference on Machine Learning. 2014. URL attps://apiz
bemanticscholl www.oolrrrrrrraaoooooooo wwtp//rpps/orrrrroooooooo 3228832.
Yanzhi Chen, Dinghuai Zhang, Michael U Gutmann. Aaron Courville, and Zhanxing Zhu. Neural approximate sufficient statistics for implicit models. ArXiv, abs/2010.10079 2020. http:t //api. semant cascholar www.aarrrrrrooorooreee /Ctp//oruuuuii 22 48400445 62.
Alexander G. de G. Matthews, Michal Arbel, Danilo Jimenez Rezende, and A. Doucet. Continual repeated annealed flow transport monte carlo. ArXiv, abx/2201.13117, 2022. URL https: :D: 2443022223..
semant / //api..a// //api.semant:
cacholar www.entcccchllllc
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.
Tristan Deleu, Mizu - Toomey, Jithendaraa Subramanian, Nikolay Malkin, Laurent Char- lin, and Yoshua Bengio. Joint Bayesian inference of graphical structure and parameters with a single generative flow network. arXiv preprint arXiv:2305. /9366, 2023.

---

Published as a conference paper at ICLR 2024
Guillaume Desjardins, Aaron C. Courville, Yoshua Bengio, Pascal Vincent, and Olivier Delal- leau. Tempered markov chain monte carlo for training of restricted boltzmann machines. In International Conference Artificial Intelligence and Statistics, 2010. URL https://app) semant i wwccolschhoooooooo ar. wttg/oorruseeeeeeeeee Corpus 6382829.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components www.sssesaannnnnoeeeeee estimation. CoRR, abs/1410.8516, 2014. URL http:ttts//ppp...mmaacccccooooo
org/
CorpusID:13998862
A. Doucet, Nando de Freitas, and Neil J. Gordon. Sequential monte carlo methods in prac- tice. In Statistics for Engineering and Information Science, 2001. URL https://agg /ppi.
semanticscholar.o wwwp:/trrssoooouuoooeeeeaa
Corpus ID: 30176573.
A. Doucet, Will Grathwohl, Alexander G. de G. Matthews, and Heiko Strathmann. Score-based diffusion meets annealed importance sampling. ArXiv, abs/2208.07698 2022. URL https:
//api. s
/Cr/psrpsIID: :5119945222.
Weitao Du, Tao Yang, Heidi Zhang, and Yuanqi Du A flexible diffusion model. ArXiv, www.sti wtw.ooaarrrrrrrrroooeoeoooeeeeee CorpusII
Abs/2200.10366 2022. URL https: http:///aa.seenisssteeteeeo
249889644.
Simon Duane, Anthony D. Kennedy, Brian Pendleton, and Duncan Roweth. Hybrid monte carlo. 1987. URL sittps://apps semanticscholee org
Daan Frenkel and Berend Smit. Understanding molecular simulation: from algorithms to applica- tions. 1996. URL https: http:p/:/.ppp..eeatttt...coooa/a wwt.s//chhh.....oooooaaaaaaaaeae :otg//oopsII::22200
62612337.
Marylou Gabri'e, Grant M. Rotskoff, and Eric Vanden-Fijnden. Adaptive monte carlo aug- mented with normalizing flows. Proceedings of the National Academy of Sciences of the United States of America, 119, 2021. URL //atps:://p..... semanticschole org/ hatg//orpuss: 235195952.
Christina Gao, Joshua Isaacson, and Claudius Krause. i-flow: High-dimensional integration and sampling with normalizing flows. Machine Learning; Science and Technology, 1, 2020. URL https: //ap://appppi www.nteetnnteeeeeee semanticscholai org/ 2
Zhenggqi Gao, Dinghuai Zhang, Luca Daniel, and Duane S. Boning. Rare event probability learning by normalizing flows. ArXiv, abs/2310.19167, 2023.
Tomas Geffner and Justin Domke. Langevin diffusion variational inference. In Interna-
lional Conference on Artificial Intelligence and Statistics, 2022. URL nttps://apl
semanticscholes wwwilarrrrrrr..oooooo http/CCorpuDDD:D: Dtt.:ccho......ooooou.......:::. :25155945::211555111
Alex Hernández-Gareía, Nikita Saxena, Moksh Jain, Che Liu, and Yoshua Bengio. Multi- fidelity active learning with gflownets. ArXiv, abs/2306.1 11715, 2023. URL https://api.
semanticscholi wwwpsscccwaa...e.eeeoooooo Corp//s :DDD::25800 259202477.
José Miguel Hernández- and Ryan P. Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In International Conference on Machine Learning, 2015.
URL https::///// /api. paapissemante htt.:ccccccc....oooccccpssssia
wwwp//oppws...aeae.c.ocoo 86651117...
Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusion probabilistic models. ArXiv, abs/2006. 11239, 2020a. URL https://api: wew.ntisscooooor..oooooo wrr.://uusorrr/ooottttooo 219955663.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Neural Infor- mation Processing Systems (NeurIPS), 2020b.
Matthew D. Hoffman and Andrew Gelman. The no-u-tum sampler: adaptively setting path lengths in hamiltonian monte carlo. J. Mach. Learn. Res., 15:1593-1623 1623, 2011. URL https: / appi semant i eschol 1 ar. ....a..oooooooo........ ortp:/oorppaat:
12

---

Published as a conference paper at ICLR 2024
Lars Holdijk. Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling. Stochastic optimal control for collective variable free sampling of molecular transition paths. 2022. URL https: /apppapp.... sewant
org/ Corpus 2599511301.
Scott A. Hollingsworth and Ron O. Dror. Molecular dynamics simulation for all. Neuron, 99:1129- ID 22311344,
1143, 2018. URL https::///// /api. semant www.iicw..c........oooooeo
Edward J Hu, Nikolay Malkin, Moksh Jain, Katie Everett, Alexandros Graikos, and Yoshua Bengio. GFlowNet-EM for learning compositional latent variable models. International Conference on Machine Learning (ICML), 2023.
Xunpeng Huang, Hanze Dong, Hao, Yi-An Ma, and T. Zhang. Reverse diffusion monte carlo. 2023.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Gariia, 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.
Moksh Jain, Tristan Deleu, Jason S. Hartford, Cheng- Hao Liu, Alex Hernánezz-Garíaa, and Yoshua Bengio. Gflownets for ai-driven scientific discovery. ArXiv, abs/2200..0065,, 2023a. URL hitpps:/aage http://pc..settttccccccceeccoooe
wrg/ :ttp/CoorusII::::::: :2:544591:
Moksh Jain, Sharath Chandra Raparthy, Alex Hernandez-Garchn Garcia, Jarrid Rector-Brooks Yoshua Ben- gio, Santiago Miret, and Emmanuel Bengio. Multi-objective GFlowNets. International Confer- ence on Machine Learning (ICML), 2023b.
Hilbert J Kappen. Path integrals and symmetry breaking for optimal control theory. Journal of statistical mechanics: theory and experiment, 2005 2005(11):P11011 2005.
Michael Kearns and Satinder Singh. Bias-variance error bounds for temporal difference up- dates. In Anmual Conference Computational Learning Theory, 2000. URL https://api. semant cestiiischha..... org/ /ttg//orreee IC:r/ouoI: 50:3557
Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search gflownets. ArXiv, abs/2310.02710, 2023.
Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. International Conference Learning Representations (ICLR), 2014.
Salem Lahlou, Tristan Deleu, Pablo Lemos Dinghuai Zhang, Alexandra Volokhova, Alex Germandez-Garcin Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of con- tinuous generative flow networks. International Conference on Machine Learning (ICML), 2023.
Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of baore-based generative model- ing for general data distributions. ArXiv, abs/2209. 2022. URL https://api.
semant cscho ar. org/ Corpus ID: 252531877.
Wenqian Li, Yinchuan Li, Zhigang Li. Jianye Hao, and Yan Pang. Dag matters! gflownets enhanced explainer for graph neural networks. ArXiv, abs/2303.00448, 2023a. URL https: semanti cscholacllllll org Co:
Yinchuan Li, Shuang Luo, Haozhi Wang, and Jianye Hao. CFlowNets: Continuous control with generative flow networks. International Conference Learning Representations (ICLR), 2023b.
Dianbo Liu, Moksh Jain, Bonaventure F. Dossou, Qianli Shen, Salem Lahlou, Anirudh Goyal, Nikolay Malkin, Chris C. Emezue, Dinghuai Zhang, Nadhir Hassen, Xu Ji, Kenji Kawaguchi, and Yoshua Bengio. Gflowout: Dropout with generative flow networks. International Conference on Machine Learning, 2022. URL https: /api.
org/CorpusII:
253097963.
Jun S. Liu. Monte carlo strategies in scientific computing. semanti cschol www.arroo//oooouttttti 2ttps//oorpuu1i:22222222222 62226424.
2001. URL https: // /api.
13

---

Published as a conference paper at ICLR 2024
Jiangyan Ma, Emmanuel Bengio, Yoshua Bengio, and Dinghuai Zhang. Baking symmetry into gflownets.
Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bose, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2022.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GelowNets and variational inference. International Conference on Learning Representations (ICLR), 2023.
Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf. and José Miguel Hernaneez-obatto Flow annealed importance sampling bootstrap. ArXiv, Abs/2208.01899 2022. 251280102.
URL https://app
Thomas P. Minka, Expectation propagation for approximate Bayesian inference. arXiv preprint arXiv:1301.2994 2001.
Jesper Møller, Anne Randi Syversveen. and Rasmus Waagepetersen. Log gaussian COx processes. Scandinavian Journal of Statistics, 25, 1998. URL attps://api.
semanticscholas wtg./orpus/or..roo 10:054437720
Pierre Del Moral and A. Doucet. Sequential monte carlo samplers. Journal of the Royal Statistical Society: Series (Statistical Methodology), 68, 2002. URL https://agt
120747782
Thomas Müller, Brian McWilliams, Fabrice Rousselle, Markus H. Gross, and Jan Novák. Neural importance sampling. ACM Transactions on Graphics (TOG), 38:1 - 19, 2018. URL https:
Radford M. Neal. Bayesian learning for neural networks. 1995. URL https://api.
semanticsen wtt.ttooccoo...oooooorrraaaaaaaa
Radford M. Neal. Annealed importance sampling. Statistics and Computing, 11:125-139, 1998.
URL https: /api. peapp.semmaaae semanti wtt.cccccoa.......rrrpppeppopepp
//api.
Radford M. Neal. Slice sampling. The Annals of Statistics, 31(3), Jun 2003. ISSN 0090-5364. doi: 10. 1,.111//////05556246111 URL http:// /ttp///xx//// doi / 1 2144 /aaaa//////10555////
Kim A. Nicoli, Shinichi Nakajima, Nils Strodthoff, Wojeiech Samek, Klaus-Robert Müller, and Pan Kessel. Asymptotically unbiased estimation of physical observables with neural samplers. Physical review: E, 101 2-1:023304, 2020. URL https: /tpp//app...//// wwmanticsccooll.................
CorpusID:21109644
Kim A. Nicoli, Christopher Anders, Lena Funcke, Tobias Hartung, Karl Jansen, Pan Kessel, Shinichi Nakajima, and Paolo Stornati. Estimation of thermodynamic observables in lattice field theories with deep generative models. Physical review letters, 126 3:032001, 2021. URL https: //api. icscholll org/ erg/CorpusID : 220514403.
Kim A. Nicoli, Christopher J. Anders, Tobias Hartung, Karl Jansen, Pan Kessel, and Shinichi Nakajima. Detecting and mitigating mode-collapse for flow-based sampling lattice field theories. ArXiv, abs/2302, 14082, 2023. URL https:/ //api./pppp..pp semant icscholarc orrg
CorpuSID:2572333
Frank Noé, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365. 2019. URL https://agg
semant csshhllarre www.la...rr....e...e...e........
Carg//orpussii ID: 54458652.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023a.
14

---

Published as a conference paper at ICLR 2024
Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative aug- mented flow networks. International Conference on Learning Representations (ICLR), 2023b.
Ling Pan, Dinghuai Zhang, Moksh Jain, Longbo Huang, and Yoshua Bengio. Stochastic generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2023c.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing wwwanticscoola.............c.... flows. ArXiv, abs/1 505.0500000700 2015. URL httpps//app..smmttiiccooo...
CorpusID:125544
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, Julius Berner, and Guan-Horn Liu. Improved sampling via learned diffu- sions. ArXiv, abs/2307.01198, 2023. URL https:/ /app./app... semanticscholar org CorpusID:259316:0
Geoffrey Roeder, Yuhuai Wu, and David Kristjanson Duvenaud. Sticking the landing: Simple, lower-variance gradient estimators for variational inference. In NIPS, 2017. URL https://
api. www.nticsccccllarr....ooaocc.aa. ottp/oorpssI::3332110.
Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towards understanding and improving gflownet training. ArXiv, abs/2305.071700, 2023. URL https: hopp:://ap.... wwwaaannisscolla.orogo.. 258676487.
Jascha Narain Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep un- supervised learning using nonequilibrium thermodynamics. ArXiv, has//553.00385, 2015. URL pettss//aapis semanticscholat org/ Cotg/oorpssID. 1414888155
Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Er- mon, and Ben Poole. Score-based generative modeling through stochastic differential equa- tions. ArXiv, Abss2211.1345k 2020. URL https://pp..emmatiicchoo..a
Io:
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Tijmen Tieleman and Geoffrey E. Hinton.
Using fast weights to improve persistent contrastive divergence. In International Conference on Machine Learning, 2009. URL https://api.
semanticscholen wwreaaar..rrrrrreeeeeeeeeeeeeeee :t:p://ooosss::::::::::
Belinda Tzen and Maxim Raginsky, Theoretical guarantees for sampling and inference in generative models with latent diffusions. In Annual Conference Computational Learning Theory, 2019. URL https: api. semant icscholar. org/ Corpus :
Lorenz Vaidl, Kim A. Nicoli, Shinichi Nakajima, and Pan Kessel. Gradients should stay path: better estimators of the reverse- and forward kl divergence for normalizing flows. Machine Learn- ing: Science and Technology, 2022a. URL https: /api..ppp..// semanticscholoo. orrg/
CorpusID:230 2
Lorenz Vaitl, Kim A. Nicoli, Shinichi Nakajima, and Pan Kessel. for continuous normalizing flows. ArXiv, abs/2206.090n 09001,, 2022b. URL https: /ppp.
Path-gradient estimators
semanticscho ar. wttpoooooss:::4499000 oorpraacccccooooooooooooooooaoao
Francisco Vargas and Nikolas Nusken. Transport, variational inference and diffusions: with ap- plications to annealed flows and schrödinger bridges. ArXiv, abs/2307.01050, 2023. URL hittps://epp semant cscchl ar. www.oorg//oouuuooe 259317007..
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark A. Girolami, Neil D. Lawrence, and Nikolas Nusken. Bayesian learning via neural schrödinger-fölle föllmer flows. Statistics and Comput- ing, 33, 2021.
15

---

Published as a conference paper at ICLR 2024
Francisco Vargas, Will Grathwohl, and A. Doucet. Denoising diffusion samplers. ArXiv, abs/2302.13882 13834, 2023a. URL https://agg wwwanticscholattttt.tteeeeecceee ott.op/eep//ppp...ooiiiiico.ooaa hrtp://pps//rrD/// 257220165.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nusken. Transport meets variational inference: Controlled monte carlo diffusions. 2023b.
Francisco Vargas, Teodora Reu, and Anna Kerekes. Expressiveness remarks for denoising diffusion models and samplers. ArXiv, 00022355006003 2023c. URL https://api.
semanticscholars ID 11:
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Compu- tation, 23:1661-16744 2011. URL https://ap........ wemanticsccoolleee..e... Org/CorppsID: 5560643.
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient langevin dy- namies. In International Conference on Machine Learning, 2011. URL https://api.
semanticscholar.co www.:aa..IoCC/oCCIII.II.. CorpusID:21789E::
Peter Wirnsberger, Andrew J. Ballard, George Papamakarios, Starrt Abercrombie, Sébastien Racanière, Alexander Pritzel, Danilo Jimenez Rezende, and Charles Blundell. Targeted free en- ergy estimation via learned mappings. The Journal of chemical physics, 153 14:144112, 2020. cscholar. URL https: /api..p..../ semant iesmanticccholaa wwt.or///ooooooo ArggCCorpusDDy 210082968.
Dian Wu, Lei Wang, and Pan Zhang. Solving statistical mechanics using variational autore- gressive networks. Physical review letters, 122 8:080602, 2018. URL https://api.
Hao Wu, Jonas Köhler, and Frank No'e. Stochastic normalizing flows. ArXiv, abs/2002.06607,
2020. URL https: api semant wrm.mnntccchhoa.... wttp/ooorpu11.::111333399999
Winnie Xu, Ricky T. Q. Chen, Xuechen Li, and David Duvenaud. Infinitely deep bayesian neural networks with stochastic differential equations. In International Conference on Artificial Intelli- gence and Statistics, pp. 721-738. PMLR, 2022.
David W. Zhang, Corrado Rainone, Markus F. Peschl, and Roberto Bondesan. Robust scheduling with gflownets. ArXiv, abs/2302.05446, 2023a. URL https: /api. www.stemmannnnhoaeeeeeeeo
org/ Prrg/CrrpusID: 256827133.
Dinghuai Zhang, Jie Fu, Yoshua Bengio, and Aaron C. Courville, Unifying likelihood-free inference with black-box optimization and beyond. In International Conference Learning Representa- tions, 2021. URL https: /app.. semant www.emtiinncoiiiia..eeee cwcholar....oooooe wwwoor//ooorouo Corpus 246800129
Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying generative models with GelowNets and beyond. arXiv preprint (tXix:2209902666:22 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, Hanjun Dai, Nikolay Malkin, Aaron Courville, Yoshua Bengio, and Ling Pan. Let the flows tell: Solving graph combinatorial optimization problems with gflownets. ArXiv, abs/2305. 17010, 2023b. 258947700.
URL https: //api. semant semanticscholar. cr..pppur...///ros.I..:..
Dinghuai Zhang. L. Pan, Ricky T. Q. Chen, Aaron C. Courville, and Yoshua Bengio. Distributional gflownets with quantile flows. arXiv preprint arXiv:2302.05599, 2023c.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: a stochastic control approach for sam- pling. International Conference on Learning Representations (ICLR), 2022.
Heiko Zimmermann, Fredrik Lindsten, J.-W. van de Meent, and Christian Andersson Naesseth. A variational perspective on generative flow networks. ArXiv, abs/2210.07992, 2022. URL https: //api. semant i cacholas wewaerorrggoooee www.oor///ooooo::::ooo
16

---

Published as a conference paper at ICLR 2024
A NOTATIONS
Symbol
Description
RD
D-dimensional real space
N N_
number of steps in the discrete-time stochastic process GFlowNet state space, equals {(n, Xa) n E (O,1,..,,N)... N),x, ERP): (n, xn) is often writen as Xn, for simplicity s')
S
A
action transition space (edges
s
state
S0
initial state, equals (0, 0) in VE modeling
complete trajectory (Xo ..... + XN), or Xp:N
learnable parameter of DGFS
F2 RR.. - R
state flow function at step forward policy (distribution over children next step) backward policy (distribution over parents / previous step) forward process defined Equation reference process defined in Equation target process defined in Equation 4 marginal distribution of reference process at step y
PF
PB
Q
Quei
P
P .R.
+ R.
h R+
discretized step size for the stochastic process drift term in the VE stochastic process (Equation 5) variance coefficient in the stochastic process standard Wiener process in continuous-tims stochastic process reward function (unnormalized target density) intermediate scalar, equal ExER" R(x)
dW,
R. R. + R4
Rn : R + R+
signal defined in Equation 16
Z E R
p: RDD R.+
same thing as R(-)
RDD + R+
normalized density, equal to pl-)/Z
B DETAILS ABOUT STOCHASTIC PROCESSES
B. VARIANCE EXPLODING (VE) PROCESS
Song et al. (2020) and Zhang & Chen (2022) take variance exploding stochastic process modeling. In Zhang & Chen (2022), the authors use Gisarnov theorem to derive an optimal control objective for KL divergence minimization. In this work, present discrete-time derivation for completeness.
The reference process starts at phl pho()), which is a Dirac distribution at O, and writes
(xo:N))) : X.+1 Xn + Vhoen, en ~ NO0, I), Xo = O,
(18)
N
N-1
+ (xo:N) = (Xo)
II ==0
PF (x.+1/xx) ( (x) II 1==
(19)
Therefore, (-) (:; hnno According to Bayesian formula,
Pp((x pf (x,.) (Xa+1 xx) N(xx; O, IIn Xn, Xn, 1
(20)
set (xn+1)
N(x 0, h(n 1)o21)
-N(xx n Xatl, +
Io
(21)
The target process is then
n (xx))
Pxxu...... .., xN) Qro ..... .....,,,,
m(xx) II
Pa (x,, Xn+1).
(22)
ret (XN)
==0
17

---

Published as a conference paper at ICLR 2024
The trajectory level log probability difference between model and target is
Q
log
log
+ log
(23)
Qree
P
It is obvious that the second term is 2 E. (xN) =(xN]
thus only need to look the first term
N-1
PF Xnx 0)
N-1
M( + hf(xx, n2), ho21)
log
log
Σ
log
(24)
Quet
PF (x (x,)
Nn; h)21)
n=0
n=
N-1
Vh
[x n2 0)1 + 20
0² f(x,, n: 0) En
(25)
n=0
Notice that under the expectation of 2
EQ [f(x,: 0)Te) em]
f(xx, mi 0) [Em] O.
(26)
Accordingly,
DKL. (olP) Eg log 2
Eg
Σ
h
n1 [x;;;;;;;)1* + log
1/ (xN)
(27)
2a
*(XN)
which leads us Equation 6.
B.2 VARIANCE PRESERVING (VP) PROCESS
Ho et al. (2020b) and Vargas et al. (2023a) take this VP formulation of diffusion process. derivation is similar to the last section, just with different variance schedule.
The
We first define the backward policy
(x(Xn|Xx+H (x+1) N(xx: /1 11:
(28)
Then The target and reference process
N
P(Xo.N) = T(XN)
II
Pi (xn (x+1);
(29)
N
nHd (x = PN (xN)
II Pa((x|Xat1))
(30)
==
where (+) o²)), This guarantees that (-) 0, Vn. From using Bayesian formula Ore!
PF (xn+ 1|x,) p-tt (Xn+1) PB xx+1)
(xx++1 /1 Baxn, Br
(31)
(x)(())) (x)
Inspired by this, let us define the (learnable) forward process to be
PF (x 0) = -N(x V B,x, f(X,, n. 0), B2,2)),
(32)
N-1
@(xo:N) 15 (xo)
II
Pr(x X);
(33)
==0
Similar to the derivation in VE,
Eg
log
Eo
N- Σ
log
PF (x Xni 0)
Σ
Eq
1
[( ne; 0)133
(34)
Qref
pr (xn+ (xn)
2.3n0²
Therefore,
DKL (Q)P)
Feo
log P
Q
F
N
Σ
[[f(x,,; n; 0)12 + log
PN
p (XN)
(35)
28
T((N)
which is the learning objective of DDS algorithm.
18

---

Published as a conference paper at ICLR 2024
C ALGORITHMIC DETAILS
C.I RE-DERIVING DETAILED BALANCE CONSTRAINT
Writing down Equation 12 with n and + 1, we have
N-1
N-1
(x 0)
II
Pe (x 0); = =((N)
II
Ps (x(xx(x(4)),
(36)
N
N N-1
F Fu+1 (x 0)
II Pr (xI+1 x1:0) mxxx)
II (b(x((x)1)
(37)
1m++1
1m++1
Fn (x,, (
Pa(xn Xn+1)
(38)
F 1x=-1;0)
where the last equation is obtained by comparing the first two equations both sides, which res- onates the GFlowNet DB criterion.
C.2 log Z ESTIMATOR ANALYSIS
The bound in Equation 17 could be derived in the following way:
log log f
P(xx:N)
[u(xN) dxn = log
P(Xu:N) dx log Ee
(39)
[(xx:N)
1
B
P(x
(b) )
1
B
P(x
(b)
loogE
$ EQ
log B
(40)
B
0(x
((x(.N)
(6)
bml
hml
B
= log
1
exp
(((T(1)).
7(0) Q
(41)
B
Notice that this bound is valid without any assumption on Q. According to analysis in Burda et al. (2016), this lower bound estimator is monotonically increase B goes larger; what's more, it is tight bound which means that would converge to log when B - co. We remark that this is actually also the estimator used in Zhang & Chen (2022).
C.3 GRADIENT VARIANCE ANALYSIS
In Figure 2, conduct experiment Funnel and report the gradient variance of DGFS and PIS. The variance of DDS is in similar scale to PIS thus we ignore it. According to the analysis in Section 3.3, there are two reasons of why DGFS is more robust algorithm: (1) DGFS utilize short trajectory which is beneficial for efficient credit assignment (2) DGFS does not take the KL divergence minimizing formulation which has non-zero gradient variance. To disentangle these two factors, we run a variant of DGFS which only utilize full trajectories to update the parameters in Figure 7. This variant is shown have gradient variance higher than the original DGFS but still much lower than PIS. This confirms that both factors are valid reasons for why DGFS has lower gradient variance.
For general probabilistic methods with KL between model (-(:) and target distribution p(-).
Ve DKLLL (g0(x)llp(x)) = E, Vo log go (x) log Yo (xx) p(x)
+ Ey Vo log (x))
(42)
When the model reaches its optimal solution, i.e., qo p, the first term is as contains gx(x) 0. However, the second term Vg log (x) only zero under expectation of qo- That log P(X]
to say, there will be a stochastic non-zero gradient term with KL formulation. This observation is the core motivation of Roeder et al. (2017).
19

---

Published as a conference paper at ICLR 2024
True
10"
2.
Godet vattane
10"
DGFS
DDS
0.5
202
400
800
1000
Training stops
Diflasion
Figure 7: DGFS, PIS, and DGFS trained with only full paths.
Gradient variance of
Figure 6: Drift network magni- tude of DGFS and PIS.
Figure 8: Manywell results.
Algorithm DGFS training algorithm
Require:
DGFS model with parameters 0, re- ward function (l-), variance coefficient ō. repeat
Sample with control ((:,;;)) and ē; 0 + VeC(r; 0) (as per Equation 15);
on-policy (bias=0.23)
off-policy (bias=0.09)
Update 9 with Adam optimzier; until some convergence condition
Figure 9: Demonstration of DGFS's off off-policy exploration ability in a modified MoG+ task.
C.4 DGFS DETAILS
Here we describe details about DGFS algorithm. As shown in Algorithm 1, DGFS first interact with the environment collect training data. Next, it updates its parameters 0 with methods described in Section 3. Notice that the training of DGFS is very general and does not have to use on-policy trajectories. Therefore, we could use different variance coefficient ő to rollout. This could poten- tially bring better exploration ability and thus better mode capturing performance. In this work, we assume DGFS does on-policy rollout unless otherwise specified.
We use PyTorch to implement DGFS algorithm. For PIS implementation, we take its original Py- Torch code base. For DDS, we build our own implementation based on PIS code base. For other methods, we take the DeepMind jax code base.
DGFS contains two sets of deep network parameters: one the same as the neural network in PIS / DDS, the other is the flow function to amortize the intermediate marginal densities. The former one for the control network, and express it in the way of fffn.nn)/// = NN, (). +NN2(n) These use Fourier features to embed the step index input. We also ablate diffusion- based sampling methods without gradient information in network parametrization C.6. The second one is scalar output network NN(, n) with similar structure. Same to PIS, layers after the embedding layer. We set to be 2. For all experiments, we train with batch size of 256 and Adam optimizer. We have not tuned too hard on the learning rate, but simply use I x 10-4 and 1 x 10-3 for the policy network (i.e., drift network) and the flow function network. The training keeps for 5000 optimization steps although almost all experiments converge the first 1000 steps. Similar to PIS, we set the number of diffusion steps N to 100 for all experiments. We set the diffusion step size based methods. We set a to 1. Zhang & Chen (2022) set evaluation particle number B = 6000 for the funnel task and B 2000 for other tasks for unknown reasons. In this work, we set the number of evaluation particles B to be 2000 for all tasks for consistency.
log uu.).
two 64-dimensiioal hidden
0.05 for the MoG and Cox task and 0.01 for other tasks for all diffusion
C.5 MORE EXPERIMENTS
20

---

Published as a conference paper at ICLR 2024
0.6
02
0.0
DGPS
wo inteeedlaa w/'P mod:ling
w
wi W/ w/-i 500
DGFS varians
Figure 10: Ablation on DGFS design choices on the VAE task.
Table Experimental results benchmarking target densities, where the diffusion-baseer samplers are not allowed to directly access gradient information of the target log density in the drift network design. Note the PIS-NN and DDS-NN here still require know the gradient information due to their KI. formulation, while the DGFS-NN here is real black-box algorithm and only need to know zeroth order information from the target.
MoG
FUNNEL
MANYWELL
VAE
Cox
PIS-NN (w/ GRADIENT)
1.943+0.259
0.374
2.712+0.017
0.368+0.001
85.36+3.30
DDS-NN (w/ GRADIENT)
2.183:10.03
0.351 10.022
4.016
1.740
N/A
DGFS-NN
0.023:10.001
0.354
2.659
0.496:10100
109.0
We conduct ablation studies on a series of design choices of DGFS. In Figure 10, we compare the performance of DGFS and a few its variants on the VAE task. In the VAE task, the unnormalized target density log uu(-) is defined by the log latent prior log plus the log decoder likelihood log p(xutal-)) which equals the latent posterior log (xts) plus a constant value. The VAE trained on the MNIST dataset, and Xohs is fixed image data. We first remove the usage of the intermediate signals (i.e., the parameterization
Equation 16), which leads to slightly worse performance. We then try to train DGFS with VP modeling in Appendix B.2. For this VP modeling, we also take the form of Equation 16 for getting intermediate signals. This gives a bias 0.86. Notice that this is still better than both PIS and DDS, demonstrating the validity of the proposed DGFS training method. We also ablate the weighting coefficient and number of diffusion steps N. We can their changes lead to small variations but the overall convergence remains good.
flow function
Regarding the rumning time the proposed algorithm, DGFS has the same inference procedure as PIS, and thus having the same inference speed as PIS. The training overhead of DGFS slightly (20%) higher than PIS due to the introduction of the additional flow function network. In the Funnel task, the training time of DGFS and PIS for one batch are 0.54s and 0.65s respectively.
We then show some missing results from the main text. We put the visualization of the joint of the second and third dimension Manywell task in Figure 8. We can that PIS misses one important mode and DDS could not handle the mode separation well enough.
We tried to use MNIST- preetrided normalizing flow as the target density. However, we find well-trained normalizing flow almost always assign coo to the log likelihood for samples, which makes impossible for any sampler success. This actually makes sense as the normalizing flow is only trained on finite size image dataset and never gets out-of-domail inputs.
C.6 REMOVE GRADIENT INFORMATION FROM PARAMETRIZATION
In the network design, PIS explicitly use the score function gradient the log unnormalized den- sity function-as one term in the drift network parametrization. Our work and DDS follow its modeling. In this section, we show performance where the score function are removed from the ar- chitecture of all three diffusion- basse samplers in Table 2. We attach "-NN" postfix to their names to distinguish between their original algorithms. Notice PIS-NN and DDS-NN still require the knowl- edge of the score function implicitly based on their KI. formulation. On the other hand, DGFS-NN
21

---

Published as a conference paper at ICLR 2024
Experimental results on benchmarking target densities, DGFS+ denotes DGFS with addi- tional off off-policy exploration trick in the initial stage of training.
Table 3:
MoG
FUNNEL
MANYWELL
VAE
Cox
DGFS
0.019+0.000
0.274
10.047
0.180+0.000 083
8.974=1.169
DGFS+
0.012:0.002
0.322=0.029
0.932+001223
0.162+0:015
10.23 1.237
Table 4:
Experimental results on benchmarking target densities for methods in Lahlou et al. (2023), with or without off-policy exploration technique. We also present performance of TB and its off- policy variant; see details in Section C.8.
MoG
FUNNEL
MANYWELL
VAE
Cox
LAHLOU ET AL.. (2023)
1.891 L..004
0.398 000061
10.653
4.563 11.029
2728.151.54
LAHIOUU ET AL. (2023)+
00.19
0.373=0.020
6.257+2.41
2.529=0.737
2722.+31.4
TB
0.018:00.002
0.309=0.023
1.218:+0.991
1.097=0.239
138. 6.7.850
TB+
0.009.1..00
0.373 0.003
1.21010.192
0.842 00.248
240.4
is a real black-box method, i.e. only requires zeroth order information of the unnormalized density function. Despite using less information, DGFS-NN still very competitive in most of the tasks.
C.7 OFF-POLICY EXPLORATION ABILITY OF DGFS
The usage of the KL divergence in previous works restricts the training be on-policy, as the training trajectory has to come from the model itself. Using off-policy samples is possible only with importance sampling, which typically results in high variance. This limits the exploration ability of these methods. On the other hand, it is first shown in Malkin et al. (2023) that GFlowNets can potentially benefit from the use of off-policy exploration. To be more specific, as has been shown by similar analysis in prior works Malkin et al. (2023); Lahlou et al. (2023), the objectives in Equation 9 or Equation 14 do not require the training samples to follow any particular distribution (only to have full support), which means our method supports off-policy training without importance sampling. Therefore,
possible to employ off-policy exploration by using a larger variance coefficient in the forward process during the rollout stage. In theory, this could help DGFS to better capture diversity in complex multimodal target distributions.
In algorithmic perspective, we could use slightly larger variance coefficient ō > in the rollout stage of Algorithm to ensure that information of more modes fall into the training data, which would drive the training of DGFS put effort these distant modes. To demonstrate this property, we carefully create a more difficult Mixture of Gaussian task and name it as MoG+. We intentionally put few very distant modes in this MoG+ task, as shown by the contours Figure 9. We set the exploration variance coefficient ā = 2 contrary to a 1, which achieves better mode coverage and lower normalizing factor estimation bias with the same neural network capacity.
We conduct an initial experiment to investigate this possibility in DGFS. We test strategy that linearly anneals the from 1.1 to during the first 1000 steps of training, and denote it as DGFS+ in Table 3. Our experiment results show that this simple strategy can improve the performance under some scenarios, and can be treated as a potential hyperparameter to extend the application scope of DGFS.
C.8 OTHER ABLATION STUDY
We run the method in Lahlou et al. (2023) in this section. As shown in Table 4, this previous method and its ooolcy variant are far from comparable the baselines in Table 1. We also run the trajectory balance method with gradient information (denoted as TB) and its variant with off-policy exploration (denoted TB+) and present in the same table. These methods achieve relatively better performance, but still worse than that of DGFS.
22


Reviews:
[{"review_id": "BGz7rQVYvE", "reviewer": "Reviewer_1u4d", "summary": "The paper proposes an algorithm for better sampling from intractable high-dimensional density functions. The sampling procedure is formulated as a time sequence. The improvement then stems from the usage of partial trajectories instead of merely relying on the end-time variable. Empirical results show the superior performance of the proposed method.", "questions": "See weakness above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I haven't done any research in the sampling area, nor do I have a strong background of the sampling methods and their advantages and disadvantages. However, from reading the paper, I have gained an understanding of the problem to be solved in the paper. Moreover, the methodology makes sense to me. That being said, the paper is well written with a clear logic flow, even for people who is new to the field.\n\nThe experiments cover a wide range of data, including high-dimensional settings. The reported results look very promising and indicating the strength of the proposed method.", "weaknesses": "I am not very confident in evaluating the novelty of the proposed method. From the introduction, it seems like a combination of the existing ideas in constructing a forward Markov chain and exploitation of the detailed balance in GFlowNets. I would appreciate a discussion regarding the novelty.\n\nWhat is the running time comparison to existing baseline methods?\n\nThe reported results in table 1 uses a different metric for baseline methods. I am curious how does the proposed method compares to baseline methods using the results at the best checkpoint?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fc6bfynelU", "reviewer": "Reviewer_oYBq", "summary": "In this paper, a novel training strategy for diffusion-based samplers is proposed, where the sampler is trained with learning signals for incomplete trajectories. Such a signal is amortized with a novel flow function. Training the sampler with these intermediate learning signals results in reduced variance with respect to similar modes trained only on full trajectories, and shows improved results on a wide set of benchmarks.", "questions": "In the experiments, there is no comparison with GFlowNets methods. Is that because GFlowNets perform poorly on continuous sampling benchmarks? And on the other hand, can DGFS be used on discrete data space? And if so, how does it perform?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The paper combines recent improvements to the credit assignment problem for partial trajectories in GFlowNets with recent diffusion-based sampler methods. The introduction of the amortized flow function is novel and well-supported with both theoretical and empirical results. The experimental section presents strong empirical results, as well as insightful analysis of the gradient variance and learned drift function which confirm the claims made by the authors. Overall, the proposed method presents strong performance on challenging tasks, making it a potentially high-impact contribution to the scientific community.", "weaknesses": "In my opinion, the paper could improve in terms of clarity and in separating the previous methods from the proposed one. The paper is heavily based on previous contributions, such as denoising diffusion sampler (DDS) and path integral sampler (PIS), as well as recent contributions in GFlowNets like [1,2]. While the authors provide brief explanations of these previous methods, I think it would be beneficial to first give a clearer introduction of what such methods do (especially for DDS and PIS), and then clearly outline how the ideas are combined in DGFS.\n\nThe off-policy training strategy is mentioned but could also benefit from additional explanations and perhaps an ablation study, for example, the empirical benefits of training DGFS off-policy vs on-policy. From the text, it is not immediately clear whether DGFS can only be trained off-policy, or if it has the possibility to be trained off-policy as opposed to on-policy. \n\n[1] Madan, Kanika, et al. \"Learning GFlowNets from partial episodes for improved convergence and stability.\" International Conference on Machine Learning. PMLR, 2023.\n\n[2] Pan, Ling, et al. \"Better training of gflownets with local credit and incomplete trajectories.\" arXiv preprint arXiv:2302.01687 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xXdH35OA4K", "reviewer": "Reviewer_Tb6D", "summary": "The paper considers the problem of sampling from a distribution given its unnormalized density. Recent work proposes training a diffusion process sampler to match the target distribution. At present such methods are trained by minimizing the Kullback-Leibler (KL) divergence computed using terminal states of the diffusion process as samples. As a result, complete trajectories need to be sampled during training, while only the terminal states get direct learning signal. In this paper, authors propose learning an additional neural net to provide learning signal for intermediate steps. Authors note that the resulting objective matches the detailed balance objective in GFlowNet literature, and tap into this literature to further improve the learning objective. Authors evaluate the method on several target distributions, showing improved performance compared to common baselines, including existing diffusion-based samplers (PIS/DDS).", "questions": "At the end of page 5 authors say that the method does not \"require the training sample to follow any particular distribution (only to have full support)\". Have authors measured the effect of choosing the sampling policy? Is it worthwhile to try to explore intelligently?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "Significance: diffusion models are state of the art in generative modelling, where we can train them efficiently using the step-wise de-noising objective. Finding similarly effective objectives for training diffusion models as samplers is an important research direction.\n\nMotivation/soundness: the method is well-motivated and theoretically sound. Authors start with a clear problem with existing diffusion-based samplers (weak training signal at intermediate steps), and propose to define intermediate targets for the diffusion process using a neural network trained in parallel. The learning objective for the additional network neatly avoids having to estimate the integral in Eq. (11). The developed connection with GFlowNets allows the method to benefit from developments in this literature.\n\nResults: authors evaluate the propose method on several benchmarks, including the high-dimensional Cox distribution. A good selection of baselines is used, including MCMC, methods based on normalizing flows, and diffusion-based methods. The proposed method demonstrates improved performance across the board, as well as less variance in its results when compared to existing diffusion-based methods.\n\nQuality: the write-up is of high quality, with only a few minor typos, clean mathematical notation, and high-quality figures.", "weaknesses": "Originality/novelty: continuous GFlowNets and their connection to diffusion models have been explored by Lahlou et al. (2023) and Zhang et al. (2022a/2023b). The same authors have explored training diffusion models using alternative consistency-based objectives. While the setting in this paper is different (other work considers generative modeling, not sampling), it raises questions about the novelty of the method. Additional discussion with crisp comparisons to Lahlou et al./Zhang et al. would help.\n\nSignificance of the GFlowNet connection: Figure 10 in the appendix suggests that the effect of the forward-looking trick is relatively minor. It would be useful to see the effect of using Eqs. (14-15) instead of the originally proposed objective in Eq. (12). In other words, the practical effect of the GFlowNet connection (and associated improvements) is not clear from the results. \n\nClarity: Sections 2 and 3 are difficult to parse in places. The GFlowNet description in Section 2.2 is extremely dense. The separation of \"learning from intermediate steps\" vs. \"learning with incomplete trajectories\" in Section 3 is confusing. More generally, Section 3 could be re-structured to separate proposed methods, existing methods, and discussion more clearly. Section 4 is dense, with a lot of the references already introduced earlier in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yRB7P3v0uI", "reviewer": "Reviewer_hxZf", "summary": "This paper introduces Diffusion Generative Flow Samplers (DGFS), a neural-based sampling algorithm that utilizes diffusion models. DGFS draws heavy inspiration from GFlowNets (typically constrained to discrete data spaces) and offers a practical and scalable solution for sampling in continuous spaces. The primary message conveyed by this paper is that DGFS demonstrates the ability to leverage partial trajectory segments, thereby enabling a more efficient approach to the learning problem when compared to existing methods, which require learning over complete trajectories (from start to finish), as seen in denoising diffusion samplers (DDS). Consequently, DGFS can sample from target unnormalized probability distributions by updating its parameters during a time-dependent stochastic training process. In other words, it can do so without requiring a full trajectory specification and can accommodate intermediate signals injected between time steps before the sampling process is complete. To accomplish this, a neural network denoted as $F_n(\\theta)$ is trained to approximate the unnormalized density of the $n-th$ step.\n\nThe authors argue that DGFS should lead to more stable training, reduce variance in gradient computations, and ultimately provide access to informative intermediate training signals.", "questions": "- On the bottom of page 5, the authors state, \"On the other hand, notice that our objectives in Equation 9 or Equation 14 do not require the training samples to follow any particular distribution (only have full support).\" While I believe that in the asymptotic regime of infinitely many samples this holds, it may not be true for practical cases with a finite number of samples. In such cases, amortized variational inference schemes might encounter issues, as discussed in recent works [4]. I'm interested in whether this assumption can be relaxed or if the authors can provide arguments for the general validity of the 'full support' assumption in their case.\n\n- In the discussion about Variance Gradient Updates (section 3.3A), the authors cite works like Roederer et al. (2017), which reported that the variance of the gradient does not necessarily vanish even when the optimal distribution is achieved. However, in Refs. [1,2], authors use information-theoretical arguments to demonstrate that the variance can indeed approach zero when the target distribution is perfectly learned, using the so-called path gradient method for normalizing flows. It would be interesting to see a comparison of DGFS with such a method (specifically designed for continuous normalizing flows in [2]). Additionally, given the discussion in Sec. 3.3A, it would be beneficial to include these references and expand the discussion.\n\n- The concept of injecting information between a sequence of transformations strongly reminds me of simulated annealing and annealed importance sampling (as seen in Stochastic Normalizing Flow [4,5]). Could the authors expand on this, highlighting the significant similarities and differences, if any? There appears to be a notable overlap worth exploring.\n\n- I'm interested in seeing how the lower bound estimator for the log partition function used in this work compares to the asymptotically unbiased estimator for the partition function proposed in [7,8].\n\n- As pointed out in the **Weaknesses** section, I'm curious about how the results for the many-well task compare to [3], which has reported notable results in data-free training for multimodal target densities for quantum chemical tasks as well as the many-well problem. \n\n**Side remark**: While I acknowledge the significant contribution and potential impact of the paper, I believe additional robust experimental evidence and a broader discussion, as outlined above, would strengthen its acceptance. Adding new baselines, particularly [2,3], to the comparison in Figure 2 and Table 1 would be a valuable addition to the paper.\n\n**Minor:**\n\n- Page 2: below eq (4) the authors state: \"It can be shown that the marginal distribution of P at the terminal time N is exactly proportional to the target μ(·).\". While it might be self-evident to some, it might be helpful to provide a hint for an explicit derivation in the appendix.\n\n- Page 3: Above eq (7), the authors mention, \"[…] stochastic optimal control formulation seen in prior works.\" It would be beneficial to include explicit references to these prior works.\n\n- Page 7: to the best of my knowledge, what the authors call \"Hamilton Monte Carlo\" is more often found in the literature as \"Hamiltonian Monte Carlo\".\n\n- Page 7: In the related work section, the authors might consider citing some of the references listed below, being relevant within the context of NF-based samplers. Notably, some of these references have shown how to derive an asymptotically unbiased estimator for the normalizing constant [7,8,9], also known as the partition function in the realm of Boltzmann distributions. This is a quantity of interest for the present work as well.\n\n**References:**\n\n- [1] [Vaitl, Lorenz, et al. \"Gradients should stay on path: better estimators of the reverse-and forward KL divergence for normalizing flows.\" Machine Learning: Science and Technology 3.4 (2022): 045006](https://iopscience.iop.org/article/10.1088/2632-2153/ac9455/pdf)\n- [2] [Vaitl, Lorenz, et al. \"Path-gradient estimators for continuous normalizing flows.\" International Conference on Machine Learning. PMLR, 2022.](https://proceedings.mlr.press/v162/vaitl22a/vaitl22a.pdf)\n- [3] [Midgley, Laurence Illing, et al. \"Flow annealed importance sampling bootstrap.\" arXiv preprint arXiv:2208.01893 (2022).](https://arxiv.org/pdf/2208.01893)\n- [4] [Wu, Hao, Jonas Köhler, and Frank Noé. \"Stochastic normalizing flows.\" Advances in Neural Information Processing Systems 33 (2020): 5933-5944.](https://proceedings.neurips.cc/paper/2020/hash/41d80bfc327ef980528426fc810a6d7a-Abstract.html)\n- [5] [Caselle, Michele, et al. \"Stochastic normalizing flows as non-equilibrium transformations.\" Journal of High Energy Physics 2022.7 (2022): 1-31.](https://arxiv.org/pdf/2201.08862.pdf)\n- [6] [Nicoli, Kim A., et al. \"Detecting and Mitigating Mode-Collapse for Flow-based Sampling of Lattice Field Theories.\" arXiv preprint arXiv:2302.14082 (2023).](https://arxiv.org/pdf/2302.14082)\n- [7] [Nicoli, Kim A., et al. \"Asymptotically unbiased estimation of physical observables with neural samplers.\" Physical Review E 101.2 (2020): 023304.](https://link.aps.org/accepted/10.1103/PhysRevE.101.023304)\n- [8] [Wirnsberger, Peter, et al. \"Targeted free energy estimation via learned mappings.\" The Journal of Chemical Physics 153.14 (2020).](https://pubs.aip.org/aip/jcp/article/153/14/144112/316574)\n- [9] [Nicoli, Kim A., et al. \"Estimation of thermodynamic observables in lattice field theories with deep generative models.\" Physical review letters 126.3 (2021): 032001.](https://link.aps.org/pdf/10.1103/PhysRevLett.126.032001)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, scientifically sound, and highly rigorous, featuring a compelling theoretical framework. The main text is effectively complemented by the appendix materials, which offer additional mathematical details and experimental results.\n\n- I particularly appreciate the table that discusses notation on the first page of the appendix. I believe this should be considered standard practice, as it greatly aids readers in navigating the mathematical content.\n\n- The paper extends the theory of GFlowNets to address sampling problems in continuous space. This extension holds significant relevance for numerous applications, especially in the field of physical and chemical sciences, where unnormalized target densities in continuous space are frequently encountered.\n\n- The concepts of leveraging information from partial trajectory paths and training intermediate layers offer benefits from both theoretical and practical standpoints.", "weaknesses": "- **Limited Experiments:** I consider the experiments presented in the paper to be its primary weakness. While the theory is well-presented and comprehensible, the experiments reported in the paper appear to fall short in terms of comparison with more sophisticated baseline methods, particularly in the context of flow-based samplers.\n\n- **Related Works:** Expanding the related work section to include discussions of prior research addressing similar problems, such as gradient variance reduction and efficient sampling of multimodal densities [1, 2, 3], would add significant value. More details are provided below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)