Improved off-policy training of diffusion samplers
Authors: Marcin Sendera
Minsu Kim
Sarthak Mittal
Pablo Lemos
Luca Scimeca
Jarrid Rector-Brooks
Alexandre Adam
Yoshua Bengio
Nikolay Malkin
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion models
amortized inference
stochastic control
GFlowNets
Primary area: probabilistic_methods
TL;DR: We benchmark various approaches to training neural SDEs to match a target distribution and study ways to improve training and credit assignment.

Abstract: We study the problem of training diffusion models to sample from a distribution with a given unnormalized density or energy function. We benchmark several diffusion-structured inference methods, including simulation-based variational approaches and off-policy methods (continuous generative flow networks). Our results shed light on the relative advantages of existing algorithms while bringing into question some claims from past work. We also propose a novel exploration strategy for off-policy methods, based on local search in the target space with the use of a replay buffer, and show that it improves the quality of samples on a variety of target distributions. Our code for the sampling methods and benchmarks studied is made public at [this link](https://github.com/GFNOrg/gfn-diffusion) as a base for future work on diffusion models for amortized inference.

Improved off-policy training of diffusion samplers
Marcin Sendera Mila, Université de Montréal
Minsu Kim Mila, Université de Montréal KAIST
Sarthak Mittal
Mila, Université de Montréal
Jagiellonian University
Pablo Lemos
Luca Scimeca Mila, Université de Montréal
Jarrid Rector-Brooks Mila, Université de Montréal Dreamfold
Mila, Université de Montréal
Ciela Institute
Dreamfold
Alexandre Adam
Yoshua Bengio
Nikolay Malkin
Mila, Université de Montréal Ciela Institute
Mila, Université de Montréal CIFAR
Mila, Université de Montréal University of Edinburgh
{marcin. sendera,
nikolay malkin) quebec cubec
Abstract
We study the problem of training diffusion models to sample from a distribution with given unnormalized density energy function. We benchmark several arffusion-structredd inference methods, including simulation-barl proaches and off effiipllcy methods (continuous generative flow networks). Our results shed light the relative advantages of existing algorithms while bringing into ques- tion some claims from past work. We also propose a novel exploration strategy for off-policy methods, based on local search in the target space with the use of replay buffer, and show that it improves the quality of samples on variety of target distri- butions. Our code for the sampling methods and benchmarks studied is made public at (link) as a base for future work on diffusion models for amortized inference.
variational ap-
Introduction
Approximating and sampling from complex multivariate distributions is fundamental problem in
probabilistic deep learning [e.g., 27, 35, 26, 48, 57] and in scientific applications [3, 52, 1, 32]. The problem of drawing samples from a distribution given only an unnormalized probability density energy is particularly challenging in high-dimensional spaces and when the distribution of interest. has many separated modes [5]. Sampling methods based on Markov chain Monte Carlo (MCMC)-
such Metropolis-adjusted Langevin [MALA: 24, 65, 64] and Hamiltonian MC [HMC; 20, 31] may be slow to mix between modes and have a high cost per sample. While variants such as sequential MC [SMC; 25, 13, 16] and nested sampling [69, 10, 43] have better mode coverage, their cost may grow prohibitively with the dimensionality of the problem. This motivates the use of amortized variational inference, i.e., fitting parametric models that sample the target distribution.
Diffusion models, contimuous-time stochastic processes that gradually evolve a simple distribution to complex target, are powerful density estimators with proven mode-mixing properties [15]; as such, they have been widely used in the setting of generative models learned from data [70, 72, 28, 50, 66). However, the problem of training diffusion models sample from distribution with given black- box density or energy function has attracted less attention. Recent work has drawn connections between diffusion (leaming the denoising process) and stochastic control (learning the Föllmer
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

drift [21]), leading to approaches such as the path integral sampler [PIS; 88], denoising diffusion sampler [DDS; 78]. and time-reversed diffusion sampler [DIS; 8]; such approaches were recently unified by [63] and [79]. Another line of work [42, 86] is based on continuous generative flow networks (GFlowNets), which are deep reinforcement learning algorithms adapted to variational inference that offer stable off-policy training and thus flexible exploration [46].
Despite the advances in sampling methods and attempts to unify them theoretically [63, 79], the field suffers from some failures in benchmarking and reproducibility, with the works differing in the choice of model architectures, using unstated hyperparameters, and even disagreeing their definitions of the same target densities (see $B.1). The first main contribution of this paper is unified library for diffusion-struuturee samplers. The library has a focus on off policy methods (continuous GFlowNets) but also includes simulation-busee variational objectives such as PIS. Using this codebase, we are able to benchmark methods from past work under comparable conditions and confirm claims about exploration strategies and desirable inductive biases, while calling into question other claims on robustness and sample efficiency. Our library also includes several new modeling and training techniques, and we provide preliminary evidence of their utility in possible future work ($5.3).
Our second contribution is study of methods for improving exploration and credit assignment the propagation of learning signals from the target density to the parameters of earlier sampling steps in ciffuusion-truuctureddd structured samplers ($4). First, our results ($5.2) suggest that the technique of utilizing partial trajectory information [44, 55], as done in the diffusion setting by [86], offers little benefit, and a higher training cost, over on-policy [88] or off-policy [42] trajectory-based optimization. Second, we examine the utility of a gradient-
variant which parametrizes the denoising distribution as a correction to Langevin process [88]. We show that this inductive bias is also beneficial in the off- policy (GFlowNet) setting despite higher computational cost. Finally, motivated by recent approaches in discrete sampling, we propose an efficient exploration technique based on local search in the target space with the use of a replay buffer, which improves sample quality across various target distributions.
2 Prior work
Amortized variational inference approaches use a parametric model qo approximate given target density Purget- typically through stochastic optimization [30, 58, 2]. Notably, explicit density models like autoregressive models and normalizing flows have been extensively utilized in density estimation [60, 19, 81, 22, 51J. However, these models impose structural constraints, thereby limiting their expressive power [14, 23, 87J. The adoption of diffusion processes in generative models has stimulated renewed interest in hierarchical models as density estimators [80, 28, 76]. Approaches like PIS [88] leverage stochastic optimal control for sampling from unnormalized densities, albeit still struggling with scalability in nigh--imensional spaces.
Generative flow networks, originally defined in the discrete case by [6, 7], view hierarchical sampling (i.e,, stepwise generation) as a sequential decision-making process and represent a synthesis of reinforcement learning and variational inference approaches [46, 90, 73, 18), expanding from specific scientific domains [e.g., 36, 4, 89] to amortized inference over a broader array of latent structures [e.g., 77, 34]. Their ability to efficiently navigate trajectory spaces via off-policy exploration has been crucial, yet they encounter challenges in training dynamics, such as credit assignment and exploration efficiency [45, 44, 55, 59, 68, 39, 37]. These challenges have repercussions in the scalability of these methods in more complex scenarios, which this paper addresses in the continuous case.
3 Setting: Diffussin--trucctured sampling
Let & R R be differentiable energy function and define R(x) exp(-E(x)), the reward or unnormalized target density. Assuming the integral 7 := Ser R(x) dx exists, & defines a Boltzmann density Prarget (x) = RZ on Rd. We are interested in the problems of sampling from Plarget and approximating the partition function Z given access only to & and possibly to its gradient V8.
We describe two closely related perspectives on this problem: via neural SDEs and stochastic control ($3.1) and via continuous generative flow networks ($3.2).

---

3.1 Euler-Maruyama hierarchical samplers
Generative modeling with SDEs. Diffusion models assume a continuous-tine generative process given by a neural stochastic differential equation [SDE; 75, 54, 67]:
dx, 1; dr + gx,,, r; 0) dw..
where xo follows a fixed tractable distribution μf (such as Gaussian or a point mass). The initial distribution do and the stochastic dynamics specified by (1) induce marginal densities Pr R for each 1 > 0. The functions u and have learnable parameters that we wish to optimize, using some objective, as make the terminal density p1 close Ptarget- Samples can be drawn from pt by sampling Xn ~ and simulating the SDE (1) to time onne== 1.
The SDE driving oi to Ptarget is not unique. However, if one fixes reverse-time SDE, noising process, that pushes Purget at f = to μo at f O, then its reverse, the forward SDE (1). is uniquely determined under mild conditions and is called the denoising process. For usual choices of the noising process, there are stochastic regression objectives for learning the drift u of the denoising process given samples from Purget, and the diffusion rate g is available in closed form [28, 72].
Time discretization. In practice, the integration of the SDE (1) approximated by a discrete-time scheme, the simplest of which is Euler-Maruyama integration. The process (I) is replaced by discrete-time Markov chain Xo - XAr - Na and T is the number of steps:
X1, where Ar = + is the time increment and
no.
Xe+di = X, + (xx, r; 0) + g (x,, V 2, N().
(2)
The density of the transition kernel from X, to Xy+ár can explicitly be written as
PF (xtar x,) = N (x,+ariX, (; (;),
(3)
where pF denotes the transition density the discretized forward SDE. This density defines a joint distribution over trajectories starting at Xo:
pr X1 xn)
I i=0
Pr (x(())))) xiac).
Similarly, discrete-time reverse process XI XI-Ar X1- PB (X--ar xx) defines a joint distribution' via
No with transition densities
T
[ 1=1
-I x1) = PB PB(((1-1)A x(ar).
(5)
the forward and backward processes (starting from μo and Piarget, respectively) are reverses of each other, then they define the same distribution over trajectories, i.e., for all Xy XA
X1,
pr (XAr...... X1 Xn) Ptarger (x)m ( (x, XI-Ar | xi).
In particular, the marginal densities of xj under the forward and backward processes are then equal to Ptarget. and the forward process can be used to sample the target distribution.
Because the reverse of a process with Gaussian increments in general, not itself Gaussian, (6) can be enforced only approximately, but the discrepancy vanishes as At - 0 (i.e., increments are infinites- imally Gaussian), an application of the central limit theorem that key to stochastic calculus [54].
SDE learning as hierarchical variational inference. The problem of learning the parameters 0 of the forward process as enforce (6) is one of hierarchical variational inference. The backward process transforms x, into Xg via a sequence of latent variables Xi-Ar.... Ng. and the forward process aims match the posterior distribution over these variables and thus to approximately enforce (6).
In the setting of diffusion models learned from data, where one has samples from Pturget. one can optimize the forward process by minimizing the KL divergence DRL (PL((tun (Puaget P P PF) between the distribution over trajectories given by the reverse process and that given by the forward process.
In the case that is point mass, we assume the distribution Xo XA to also be point mass, which has density (No X 1 with respect the measure μo-
3

---

This is equivalent to the typical training of diffusion models, which optimizes variational bound on the data logeleilioood (see [71]). However, in the setting of an intractable density Purget unbiased estimators of this divergence are not available. Instead, one can optimize the reverse KL:
DKL Purget" PB)
Xo) =f log Puarge X1) Pa No. .o..... X1-ai x1)
=
Xo) dxAi
dxj.
Various estimators of this objective are available. For instance, the path integral sampler objec- tive [PIS; 88] uses the reparametrization trick express (7) an expectation over noise variables 2 that participate in the hierarchical sampling of ........ X1, yielding an unbiased gradient estimator, but one that requires backpropagation into the simulation of the forward process. The related denoising diffusion sampler [DDS; 78] applies the same principle in a different integration scheme.
3.2 Euler-Maruyama samplers as GFlowNets
Continuous generative flow networks (GFlowNett) [42] express the problem of enforcing (6) reinforcement learning task. In this section, we summarize this interpretation, its connection to neural SDEs, the associated learning objectives, and their relative advantages and disadvantages.
The connection between generative flow networks and diffusion models or SDEs was first made informally by [46] in the distribution-matthingg setting and by [84] in the maximum-likelihod setting, while the theoretical foundations for continuous GFlowNets were later laid down by [42].
State and action space. To formulate sampling as a sequential decision cocision-making problem, one must define the spaces of states and actions. In the case of sampling by T T-step Fuler-Maruyame integration, assuming No is a point mass at O, the state space is
S {(0,0) 0) U ((x, x Rd. . E (Ar, 2AI.... 1)}.
with the point (x, r) representing that the sampling agent is at position x at time t.
Sampling begins with the initial state == (0,0), proceeds through a sequence of states (XAr, Af), (N2Ar, 2Ar),...r........ .... and ends at a state (x1, 1); states (x, with
are called terminal states and their collection is denoted X. From now on, we will often write x, in place of the state (x, when the time t clear from context. The sequence of states Xo XAr - ... Xj is called a complete trajectory.
The actions from a nonterminal state (X,, t) correspond to the possible next states (Xr+dr, / + Ar) that can be reached from (x by a single step of the Euler- -aruyama integrator." 3
Forward policy and learning problem. A (forward) policy is a collection of continuous distributions over the successor states states reachable by a single action of every nonterminal state (x.r). In our
context, this amounts to a collection of conditional probability densities Pr (X X,; 0), represent- ing the density of the transition kernel from x, K,+år- GFlowNet training optimizes the parameters 0, which may be the weights of a neural network specifying a density over X, conditioned on XAr-
A policy pF induces distribution over complete trajectories = (Xo XA
- x1) via
[
PF(T:0)
pF (x X)ar;
1=0
In particular, we get a marginal density over terminal states:
p (x1:0)) f PF (x XAr - x1; 0) dxar dxi-
(8)
The learning problem solved by GFlowNets is to find the parameters of policy Pr whose terminating density PF is equal to Ptarget, i.e.,
( R(x1) Vxi e Rd.
TT be precise, the fraction in (7) should be understood as a Nikoon-i derivative, which makes sense whether bn is point mass or a continuous distribution and generalizes to continuous time [8, 63].
3 Formally, the foundations [42] require assuming reference measures with respect to which the reward and kernel densities are defined. As we deal with Euclidean spaces and assume the Lebesgue measure, readers need not burden themselves with measure theory. We note, however, that this flexibility allows easy generalization to sampling other spaces, such any Riemannian manifolds, where other methods do not directly apply.

---

However, because the integral (8) is intractable and Z is unknown, auxiliary objects must be introduced into optimization objectives to enforce (9), discussed below.
Notably, if the policy is a Gaussian with mean and variance given by neural networks taking x, and r as input, then learning the policy amounts learning the drift (x,, 0) and diffusion 1; r; 6) of SDE (I), i.e. fitting neural SDE. The SDE learning problem in $3.1 is thus the same as that of fitting a GelowNet with Gaussian policies.
Backward policy and trajectory balance. A buckward policy is a collection of conditional prob- ability densities ( X;; k), representing probability density of transitioning from x, to an ancestor state x,-Ar- The backward policy induces a distribution over complete trajectories T conditioned on their terminal state (ef. (5)):
T
1 (B(X ()) 1=1
PB(T x1:4) =
where exceptionally Pa x&r) 1 as Mo is a point mass.
Generalizing a result in the discrete-space setting [45]. [42] show that pF samples from the target dis- tribution (i.e., satisfies (9)) if and only if there exists a backward policy pa and a scalar Ze such that the trajectory balance conditions are fulfilled for every complete trajectory r = (xo - XAx
- X )
Zopr(T:0) 0) = R(x))pB( x):4))
(10)
If these conditions hold, then Ze equals the true partition function Z = R(x) dx. The trajectory balance objective for a trajectory is the squared log-ratio of the two sides of (10), that is:
ZapF(T:0)
CTB(T:0,V) log R(xI)pa(T] x1:
(1I)
One can thus achieve (9) by minimizing to zero the loss LTB(T; 0, ()) with respect to the parameters 0 and v,, where the trajectories I used for training are sampled from some training policy π(T). While it is possible to optimize (11) with respect to the parameters of both the forward and backward policies, in some learning problems, one fixes the backward policy and only optimizes the parameters of p and the estimate of the partition function Zo. For example, for most experiments $5, we fix the backward policy to discretized Brownian bridge, following past work.
Off-policy optimization. Unlike the KL objective (7), whose gradient involves an expectation over the distribution of trajectories under the current forward process, (11) can be optimized off off-policy, i.e. using trajectories sampled from an arbitrary distribution æ. Because minimizing (TB(T: 0, ()) 0 for all r in the support of will achieve (9), π can be taken be any distribution with full support, so as to promote discovery of modes of the target distribution. Various choices motivated by reinforcement learning techniques have been proposed, including noisy exploration or tempering [6], replay buffers [17], Thompson sampling [59], and backward traces from terminal states obtained by MCMC [43]. In the continuous case, [46, 42] proposed to simply add a small constant to the policy variance when sampling trajectories for training. Off- off-policy optimization is a key advantage of GFlowNets over variational methods such PIS, which require on- eo-poiicy optimization [46].
However, when LTB happens be optimized on-policy, i.e., using trajectories sampled from the policy PF itself, we get an unbiased estimator of the gradient of the KL divergence (7) with respect to pF's 's parameters up to a constant [62, 46, 90], that
Br-pr (T) [V or LTa(T; 0, v = 2V or DKL. (KL(PF( (PF(T; (rl (x:)p(( x):()),
where Ver denotes the gradient with respect to the parameters of pr, but not Zp. This unbiased estimator tends to have higher variance than the reparametrizatioo-baaes estimator used by PIS. On the other hand, it does not require backpropagation through the simulation of the forward process and can be used to optimize the parameters of both the forward and backward policies.
Other objectives. The trajectory balance objective (11) is not the only possible objective that can be used to enforce (9), A notable generalization is subtrajectory balance [SubTB; 44], which involves modeling a scalar state flow fx;; 0) associated with each state x, intended to model the marginal density of the forward process at X and enforcing subtrajectory balance conditions for all partial XnAr
trajectories Xmái
X(m+t)Ar
n-1
n1
fX 8)
[
pF ( Nivri 0) = f( 0) (=m
pa (R( Mitrit), iume]
I
(12)
5

---

where for terminal states f(xi)= = R(x)). This approach has some computational overhead associated with training the state flow, but has been shown to be effective in discrete- spacee settings, especially when combined with the forward-looking reward shaping scheme proposed by [55]. It has also been tested in the continuous case, but our experimental results suggest that it offers little benefit over the TB objective in the diffusion setting (see $4.1 and $B.1).
It is also worth noting the off-policy VarGrad estimator [53, 62], rediscovered for GelowNets by [85]. Like TB, VarGrad can be optimized over trajectories drawn off-policy. Rather than enforcing (10) for every trajectory, VarGrad optimizes the empirical variance (over a minibatch) of the log-ratio of the two sides of (10). As noted by [46], this is equivalent to minimizing LTB first with respect to log Ze to optimality over the batch, then with respect to the parameters of pr.
4 Exploration and credit assignment in continuous GFlowNets
The main challenges in training off- aff-policy sampling models exploration efficiency (discovery of high-reward states) and credit assignment (propagation of reward signals to the actions that led to them). We describe several new and existing methods for addressing these challenges in the context of diffusion-structuredd GelowNets. These techniques will be empirically studied and compared in $5.
4.1 Credit assignment methods
Partial energies and subtrajectory- -baee learning. [86] studied the diffusion sampler learning prob- lem introduced by [42], but replaced the TB learning objective with the SubTB objective.* In addition, an inductive bias resembling the geometrie interpolation [47] was used for the state flow function:
log (x;; (1 1) log pf (xx) + l R(x) + ,;;,,,;;),
(13)
is the marginal density of a Brownian motion with rate at x, The use of the target density log R(x,) -x(,,) in the state flow function was hypothesized to provide an effective signal driving the sampler to high-density states at early steps in the trajectory. Such an inductive bias the state flow was called forward-looking (FI.) by [55], and we will refer to this method as FI.-SubTB in $5.
where NN is a neural network and p; (x,) ) = 0x,
Langevin dynamics inductive bias. [88] proposed an inductive bias on the architecture of the drift of the neural SDE u(,,,,; f; 9) (in GFlowNet terms, the mean of the Gaussian density PF PF(xrrrr | x,; that resembles Langevin process on the target distribution. One writes
" Xx 1; NN1 (x,, 0) NNa(1;
(14)
where NN, and NN2 are neural networks outputting vector and a scalar, respectively. The second term in (14) is a scaled gradient of the target energy the drift of a Langevin SDE - and the first term a learned correction. This inductive bias, which we name the Langevin parametrization (LP), was shown to improve the efficiency of PIS. We will study its effect on continuous GFlowNets in $5.
The inductive bias (14) placed policies represents different way of incorporating the reward signal at intermediate steps in the trajectory and can steer the sampler towards low-energy regions. It con- trasts with (13) in that it provides the gradient of the energy directly to the policy, rather than just using the energy to provide a learning signal to policies via the parametrization of the log-state flow (13).
Considerations of the continuous-time limit lead us to conjecture that the Langevin parametrization (14) with NN, independent of x, is equivalent to the forward-looking flow (13) in the limit of small time increments Ar - 0, i.e., they induce the same asymptotics of the discrepancy in the SubTB con- straints ( (12) over short partial trajectories. Such theoretical analysis can be the subject of future work.
4.2
A new method for off-policy exploration with local search and replay buffer
Local search with parallel MALA. The FL and LP inductive biases both induce computational overhead: either in the evaluation and optimization of a state flow or in the need to evaluate the energy gradient at every step of sampling (see SC.3). We present an alternative technique that does not induce additional computation cost per training trajectory.
*Despite the claimed benefits of FL-SubTB for diffusion samplers, we discovered that [86] modifies critical experimental variables in comparisons and reports irreproducible results; $..

---

Table I: Log- Dog-partition function estimation errors for unconditional modeling tasks (mean and standard deviation over runs). The four groups of models are: MCMC- based samplers, simulation-drive variational methods, baseline GelowNet methods with different learning objectives, and methods augmented with Langevin parametrization and local search. See $ for additional metrics.
Energy
25GMM
Funnel
10)
Many welly id = 32)
LGCP (d 16800)
Algorichm
Metric
Aog 2
log 2
4lng 2
2l
Alog 2
Alog Al 2
log 2
log 2
SMC
0 00560.0.000
0.561:0.89
14,90 14,99.i 078
See discousion in $8.1
GGNS (43)
0.016-0002 012
1.033-0:113
0.292x0.151
N/A
[8]
1.125-0.06
0.986.0
1.83943:10
0.004-0005
10.52+1
3.05+3:85
DDS [78]
1,760aouss
0.746.0.30
0.42443:00
0.206
7.36
0.23
471 6..44
489, 20
(88)
1,700-0 101
1.274
0.534+3/00
0.362-0
m.85wm
2.6940.01
414.42+300
LP (88)
L.799unces
0.225
0.58742.012
0:2854
13.19
0.0740.41
471.45amm
487.9221.20
TB [42]
1.1 11176.
1.071 20.
0.690
a.239
400
2.6742/22
336.70 236.70u35.22
379.50,44..0
TB Expl. [42]
0.560sa.30
0,422.0 son
0.749:3.015
0.226.0.13 186
4.01 20us
2.6 2,6883.05
346, 10.55.0.... 54
389 21
VanGand Expl.
06615-230
0,487
0.642-3:000
0.250=o
4001
2.69+3.05
370. 37-4
410.
FL-SubTTB
1.127-0,016
1 10220s00000
0.527
18220 142
3.98
2.72x305
365 20s6
402, 65..66....0
[86]
0.209
0.011
0.56343/011
317
4.25
2.6643.32
483.90+3
TB Expl. LS (ours)
0.17
0.004
0.65343u225
0.285au.00
4.57
0.1943.51
384 90ans
419.55w2.14
TB+ Espl. + LP (owrs)
0.206un.c0s
0.011 20.000
0.66622.615
L.05
7.46.1.0-
1.061.11
452 .2.1.9
477.62a1.79
TB Expl, LS (owey)
190,0.00
0.007
0.7688+3052
0.264-000
4.68 at
0.07+3.17
14:0 25
489 03. 38
VarGiad Expl. LP + LS (ours)
0.207.0.01
0.015a0.1
0.920.a.18
0.256.00477
4.11.0.5
0.02.0.21
168 15.65
467,31141.34
Highlight mean indistinguishuble from best
column with p 0.05 under one-sided Welch unpaired r-test.
To enhance the quality of samples during training, we incorporate local search into the exploration process, motivated by the success of local exploration [83, 33, 40] and replay buffer [e.g., 17] methods for GFlowNets in discrete spaces. Unlike these methods, which define MCMC kernels via the GFlowNet policies, our method leverages parallel Metropolis-adjuste Langevin (MALA) directly in the target space.
+
TB + Expl. TB Txpl +LS True samples
Figure 1: Two-dimensional projections of Many- well samples from models trained by different al- gorithms. Our proposed replay buffer with local search is capable of preventing mode collapse.
In detail, we initially sample M candidates from the sampler: (x())..... "x(")...,,x"" (M) ~ -P)() Subse- quently, we run parallel MALA across M chains over K transitions, with the initial states of the Markov chain being (x(1) ...... x (M)). ). After the Koun-in burn-in transitions, the accepted samples are stored in a local search buffer DLs- We occasionally update the buffer using MALA steps and replay samples from it to minimize the computational demands of iterative local search. MALA steps far more parallelizable than sampler training and need to be made only rarely (as the buffer is much larger than the training batch size), the overhead of local search is small.
Training with local search and replay buffer. To train samplers with the aid of the buffer, we draw a sample x from Dis (uniformly or using prioritization scheme, $E), sample a trajectory r leading to x from the backward process, and make gradient update on the objective (e.g., TB) associated with r.
When training with local search guidance, we alternate two steps, inspired by [43]. who alternate training on forward trajectories and backward trajectories initialized at a fixed set of MCMC samples. Step A involves training with on-policy or exploratory forward sampling while Step B uses samples drawn from the local search buffer described above. This allows the sampler to explore both diversified samples (Step A) and low-energy samples (Step B). See $E for detailed pseudocode of adaptive-step parallel MALA and local -earch-guided GFlowNet training.
5 Experiments
We conduct comprehensive benchmarks of various diffusiin-sttuuctured structured samplers, encompassing both GFlowNet samplers and methods such as PIS. For the GFlowNet samplers, we investigate range of techniques, including different exploration strategies and loss functions. Additionally, we examine the efficacy of the Langevin parametrization and the newly proposed local search with buffer.
5.1 Tasks and baselines
We explore two types of tasks, with more details provided in $B: sampling from energy distributions -a 2-dimensional mixture of Gaussians with 25 modes (25GMM), the 10-dimensional Funnel, the 52-dimensional Manywell distribution, and the 1600-dimensional Log-Gaussian Cox process
7

---

and conditional sampling from the latent posterior of variational autoencoder (VAE; [41, 61]).
This allows us to investigate both unconditional and conditional generative modeling techniques.
We evaluate three algorithm categories:
(1) Traditional sampling methods: We consider standard Sequential Monte Carlo (SMC) imple-
mentation and a state-of-the-an nested sampling method (GGNS, [43]).
(2) Simulation-driven variational approaches: DIS [8], DDS [78], and PIS [88].
(3) Diffusion- based GFlowNet samplers: Our evaluation focuses on TB-based training and the enhancements described in $4: the VarGrad estimator (VarGrad), off-police exploration (Expl.), Langevin parametrization (LP), and local search (LS). Additionally, we assess the FL-SubTB- based continuous GFlowNet studied by [86] for comprehensive comparison.
For (2) and (3), we employ a consistent neural architecture across methods (details in $D).
Learning problem and fixed backward process. In our main experiments, we borrow the modeling setting from [88]. We aim to learn Gaussian forward policy pF that samples from the target distribution in T= 100 steps (Ar = 0.01). Just as in past work [88, 42, 86], the backward process is fixed to discretized Brownian bridge with noise rate that depends on the domain; explicitly,
1- A
1 - o2ArTa)
PB( (x,) X,-A
(15)
understood to be a point mass at O when Ar. To keep the learning problem consistent with past work, we fix the variance of the forward policy pr to o². This simplification is justified in continuous time, when the forward and reverse SDEs have the same diffusion rate. However, in $5.3, we will provide evidence that learning the forward policy's variance is quite beneficial for shorter trajectories.
Benchmarking metries. To evaluate diffusion- [88, 42], which we restate in our notation. Given any forward policy pp, we have variational lower bound on the log-partition function log 2 = Jg0 R(x) dx:
samplers, we use metrics from past work
log ( Rxx) dx = log Rd
R(x))Dn(T xi) PF(T)
R(x1)p(( xi) (F(
E
E
log
r=(- -pp((rr)
T=( (x)---))--P(()))
We use K-sample (K 2000) Monte Carlo estimate of this expectation, log 2, as metric, which equals the true log Z if pr and pa jointly satisfy (10) and thus pF samples from the target distribution. We also employ inportance-weighted variant, which emphasizes mode coverage over accurate local modeling:
R(x ( (P I xx))
log 2RW = log
(F(T(0)
where (K) are trajectories sampled from pr and leading to terminal states x ..... (K) The estimator log 2RW is also lower bound log 7. and approaches it as K - oo [11]. In the uncon- ditional modeling benchmarks, we compare both estimators to the true log-partition function, which known analytically for all tasks except LGCP (leading discrepancies in past work; see I).
In addition, we include a sample-basse metric (2- Wasserstein distance); see SC.11
5.2 Results
Unconditional sampling. We report the metrics for all algorithms and energies in Table 1.
-0.5
-1.0
log2
We observe that TB's performance is generally modest without additional exploration and credit assignment mechanisms, except on the Funnel task, where variations in performance across methods are negligible. This confirms hypothe- ses from past work about the importance of off- policy exploration [46, 42] and the importance of improved credit assignment [86]. On the other hand, our results do not show a consistent
Constant exploration
2.0
Decaying exploration
25
Ground
Exploration rate
Figure 2: Effect of exploration variance on models trained with TB on the 25GMM energy. Explo- ration promotes mode discovery, but should be de- cayed over time to optimally allocate the modeling power to high-likelihood trajectories.
8

---

1+02
1+53
E-04
1+03
L-06
1+57
1+02
+u00
r=
t-02
t=
r.04
tw65
rw06
1eg2
1-00
p-003
1-02
1-23
1-04
1-05
1-64
1-37
1-50
t-05
Floed variance Learnd variance
-6
truth
128
Number of discretization stops
Figure Left: Distribution of No, ... X1 learned (middle) forward policy variance the 25GMM energy. The last step of sampling the fixed- variinvae model adds Gaussian noise of a variance close to that of the components of the target distribution, preventing the the sampler from sharply capturing the modes. The last row shows the policy variance learned as a function of X, at various time steps r (white is high variance, blue is low), showing that less noise is added around the peaks near r 1. The two models' log-partition function estimates are -1.67 and -0.62, respectively. Right: For varying number of steps T, we plot the log 2 obtained by models with fixed and learned variance. Learning policy variances gives similar samplers with fewer steps.
learned by 10-step samplers with fixed (top) and
and significant improvement of the FL-SubTB objective used by [86] over TB. Replacing TB with the VarGrad objective yields similar results.
The simple off-policy exploration method of adding variance to the policy notably enhances perfor- mance the 25GMM task. We investigate this phenomenon in more detail in Fig. 2, finding that exploration that slowly decreases over the course of training is the best strategy.
On the other hand, our local search-guided exploration with replay buffer (LS) leads to substantial improvement in performance, surpassing or competing with GFlowNet baselines, non- -eetWetleeet baselines, and non-amortized sampling methods in most tasks and metries. This advantage is attributed to efficient exploration and the ability to replay past low-energy regions, thus preventing mode collapse during training (Fig. 1). Further details on LS enhancements are discussed in $E with ablation studies in $E.2.
Incorporating Langevin parametrization (LP) into TB or FL-SubTB results in notable performance improvements (despite being 2-3x slower per iteration), indicating that previous observations [88] transfer to off-policy algo- rithms. Compared to FL-SubbB, which aims for enhanced credit assignment through partial energy, LP achieves superior credit assignment leveraging gradient information, akin to partial energy in continuous time. LP is either superior competitive across most tasks and metrics.
Table 2: Log-likelihood estimates test set for a pretrained VAE decoder on MNIST, The latent being sampled 20 20-dimensional. The VAB's train- ing ELBO (Gaussian encoder) was x - 1011
Algorishm Metric
log 2
log 2
GGNS 1431
-82 222402
PIS [88]
102.55120445
-7. 22.75322.821
LP (88)
-99.8904..37
44..326.n.77
TB (42)
-162.733as.s
m61.407
VarGrad
-102.544
46.502-1
Expl. [42)
-148
FL-SubTB
-147 147.992x22.00
19.1196.000
In SC.3, we study the scaling of the algorithms with dimension, showing efficiency of the proposed LS.
+ LP [86]
-111,536x1.00
-47.640.1.11
Expl.
-245.78
-55.378.0.13
TB Expl. (owrs)
-112.45-0.71
-48.87711m mm
TB Expl. + L.P (ours)
-117.26x2.50
159.157.2001
VarGrad + Expl. (ours)
-103 39:339
-7.3118
Conditional sampling, For the VAE task, we observe that the performance of the baseline GFlowNet-based samplers is generally worse
VarGrad Expl.
(owns)
-105.40=n
-8.2355-0m
VarGrad + Expl. (ours)
-99.472-0.255
36.574.0.736
VarGrad + Expl. + LP LS (ours)
-99.783.0.31
16.245un.543
than that of the simulation-based PIS (Table 2). While LP and LS improve the performance of TB, they do not close the gap in likelihood estimation; however, with the VarGrad objective, the performance is competitive with or superior to PIS. We hypothesize that this discrepancy due to the difficulty of fitting the conditional log-partition function estimator, which is required for the TB objective but not for VarGrad, which only learns the policy. (In Fig. D. we show decoded samples encoded using the best-performing diffusion encoder.)

---

5.3 Extensions to general SDE learning problems
Our implementation of diffusion-structii generative flow networks includes several additional options that diverge from the modeling assumptions made in most past work the field. Notably, features the ability to:
optimize the backward (noising) process not only the denoising process as was done for related learning problems in [12, 63, 79];
learn the forward process's diffusion rate g(,,, not only the mean a 1)
assume varying noise schedule for the backward process, making it possible to train models with standard noising SDEs used for diffusion models for images.
These extensions will allow others to build on our implementation and apply it to problems such as finetuning diffusion models trained on images with GFlowNet objective.
As noted in $5.1, in the main experiments we fixed the diffusion rate the learned forward process, an assumption inherited from all past work and justified in the continuous-tine limit. However, we perform an experiment to show the importance of extensions such as learning the forward variance in discrete time. Fig. 3 shows the samples of models on the 25GMM energy following the experimental setup of [43]. We see that when the forward policy's 's variance is learned, the model can better capture the details of the target distributions, choosing a low variance in the vicinity of the peaks to avoid 'blurring' them through the noise added in the last step of sampling.
In SC.2, we include preliminary results using used in diffusion models, in place of the reversed Brownian motion used in the main experiments.
prreerciirreriinggggg backward process, as commonly
The ability to model distributions accurately in fewer steps is important for computational efficiency, Future work consider ways to improve performance in coarse time discretizations, such non- Gaussian transitions, whose utility in diffusion models trained from data has been demonstrated [82].
Conclusion
We have presented a study of diffusion-sttructured samplers for amortized inference over continuous variables. Our results suggest promising techniques for improving the mode coverage and efficiency of these models. Future work on applications can consider inference of high-dimensional parameters of dynamical systems and inverse problems. In probabilistic machine learning, extensions of this work should study integration of our amortized sequential samplers as variational posteriors in an cocectatoonmaximizaattee
loop for training latent variable models, as was recently done for discrete compositional latents by [33]. and for sampling Bayesian posteriors over high sigh-dimensional model parameters. The most important direction of theoretical work understanding the continuous time limit (T + oo) of all the algorithms we have studied.
Note added in final version: In paper that appeared subsequently the publication of this work, Berner et al. [9] have shown connections among the families of diffusion sampling algorithms considered here and analyzed their continuous-time limits.
Acknowledgments
We thank Cheng-Hao Liu for assistance with methods from prior work, well Julius Berner, Víctor Elvira, Lorenz Richter, Alexander Tong, and Siddarth Venkatraman for helpful discussions and suggestions.
The authors acknowledge funding from UNIQUE, CIFAR, NSERC, Intel, Recursion Pharmaceuticals, and Samsung. The research was enabled in part by computational resources provided by the Digi- tal Research Alliance of Canada (https: /ttps////aaaaacaann ca), Mila pattps://miaa. /ttp:/aam.m.abee)... and NVIDIA. The research of M.S. was in part funded by National Science Centre, Poland, 2022/45/N//T66/33377
References
[1] Adam, A., Coogan, A., Malkin, N., Legin, R., Perreault-Levvaseuu L., Hezaveh, Y., and Bengio, Posterior samples of source galaxies in strong gravitational lenses with score-based priors. arXiv preprint
2022.
10

---

[2] Agrawal, A. and Domke, J. Amortized variational inference for simple hierarchical models. Neural Information Processing Systems (NeurIPS), 2021.
[3] Albergo, M. S., Kanwar, G., and Shanahan, E. Flow-based generative models for Markov chain Monte Carlo in lattice field theory. Physical Review D, 100(3)033515, 2019.
[4] Atanackovic, L., Tong, A., Wang, B., Lee, L. J., Bengio, Y., and Hartford, J. DynGFN: Towards bayesian inference of gene regulatory networks with GFlowNets. Neural Information Processing Systems (NeurIPS), 2023.
[5] Bandeira, A. S., Maillard, A., Nickl, R., and Wang, S. On free energy barriers in Gaussian priors and failure of cold start MCMC for high-dimentn lish-iimensional unimodal distributions. Philosophical transactions. Series A, Mathematicall, physical, and engineering sciences, 381, 2022.
[6] Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y. Flow network based generative models for non-iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
[7] Bengio, Y., Lahlou, S., Deleu, T., Hu, J., Tiwari, M., and Bengio, E. GFlowNet foundations. Journal of Machine Learning Research, 24(210):1-55, 2023.
[8] Berner, J., Richter, L., and Ulirich, K. An optimal control perspective on ceffusion-basee generative modeling. arXiv preprint arX1v:221..01364, 2022.
[9] Berner, J., Richter, L.,, Sendera, M., Rector-Brooks, J., and Malkin, N. From discrete-tinge time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training. arXiv preprint arXiv:2501.06114) 2025.
[10] Buchner, J. Nested sampling methods. arXiv preprint arXiv:2101.09675: 2021.
[11] Burda, Y., Grosse, R. B., and Salakhutdinov, R. Importance weighted autoencoders. Interna- tional Conference on Learning Representations (ICLR), 2016.
[12] Chen, T., Liu, G.-H., and Theodorou, E. A. Likelihood training of Schrödinger bridge using aewaaddbackward SDEs theory. International Conference on Learning Representations (ICLR), 2022.
[13] Chopin, N. A sequential particle filter method for static models. Biometrika, 89()):33-5552, 2002.
[14] Cornish, R., Caterini, A., Deligiannidis, G., and Doucet, A. Relaxing bijectivity constraints with continuously indexed normalising flows. International Conference on Machine Learning (ICML), 2020.
[15] De Bortoli, V. Convergence of denoising diffusion models under the manifold hypothesis. Transactions on Machine Learning Research (TMLR), 2022.
[16] Del Moral. P., Doucet, A., and Jasra, A. Sequential Monte Carlo samplers, Journal of the Royal Statistical Society Series B: Statistical Methodology, 68(3):411 100:44111336.. 2006.
[17] Deleu, T., Góis, A., Emezue, C., Rankawar, M., Lacoste-Julien, S., Bauer, S., and Bengio, Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
[18] Deleu, T., Nouri, P., Malkin, N., Precup, D., and Bengio, Y. Discrete probabilistic inference as control multi-path environments. Uncertainty in Artificial Intelligence (UAI), 2024.
[19] Dinh, L., Sohl-Dickstein, J., and Bengio, Density estimation using Real NVP. International Conference Learning Representations (ICLR), 2017.
[20] Duane, S., Kennedy, A., Pendleton, J., and Roweth. D. Hybrid Monte Carlo. Physics Letters B, 195(2):216-222, 1987.
[21] Föllmer, H. An entropy approach the time reversal of diffusion processes, pp. 156-163, 1985.

---

[22] Gao, C., Isaacson, J., and Krause, C. i-flow: High- -ige-nimensional integration and sampling with normalizing flows. Machine Learning: Science and Technology. 144):045023, 2020.
[23] Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D. FFJORD: Free- form continuous dynamies for scalable reversible generative models. International Conference on Learning Representations (ICLR), 2019.
[24] Grenander, U. and Miller, M. I. Representations of knowledge in complex systems. Journal of the Royal Statistical Society: Series B (Methodological), 56(4):549-581, 1994.
[25] Halton, J. H. Sequential Monte Carlo. In Mathematical Proceedings of the Cambridge Philosophical Society, volume 58, pp. 57-78. Cambridge University Press, 1962.
[26] Harrison, J., Willes, J., and Snoek, J. Variational Bayesian last layers. International Conference on Learning Representations (ICLR), 2024.
[27] Hernández-Lobato, J. M. and Adams, R. Probabilistic backpropagation for scalable learning of Bayesian neural networks. International Conference on Machine Learning (ICML), 2015.
[28] Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Neural Information Processing Systems (NeurIPS), 2020.
[29] Hoffman, M., Sountsov, P., Dillon, J. V., Langmore, I., Tran, D., and Vasudevan, S. NeuTra- lizing bad geometry in Hamiltonian Monte Carlo using neural transport. arXiv preprint 193.v:790 2019.
[30] Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. W. Stochastic variational inference. Journal of Machine Learning Research (JMLR), 14:1303-1347, 2013.
[31] Hoffman, M. D., Gelman, A., et al. The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research (JMLR), 15(1):1593-1623, 2014.
[32] Holdijk, L., Du, Y., Hooft, F, Jaini, P., Ensing, B., and Welling, M. Stochastic optimal control for collective variable free sampling of molecular transition paths. Neural Information Processing Systems (NeurIPS), 2023.
[33] Hu, E. J., Malkin, N., Jain, M., Everett, K., Graikos, A., and Bengio, Y. GFlowNet-FM for learning compositional latent variable models. International Conference on Machine Learning (ICML), 2023.
[34] Hu, E. J., Jain, M., Elmoznino, E., Kaddar, Y., Lajoie, G., Bengio, Y., and Malkin, N. Amor- tizing intractable inference in large language models. International Conference on Learning Representations (ICLR), 2024.
[35] Izmailov, P., Vikram, S., Hoffman, M. D., and Wilson, A. G. What are Bayesian neural network posteriors really like? International Conference Machine Learning (ICML), 2021.
[36] Jain, M., Bengio, E.,, Hermandez-Garcia A., Rector- Rector-Brooks, J., Dossou, F., Ekbote, A., Fu, J., Zhang, T., Kilgour, M., Zhang, D., et al. Biological sequence design with gflownets. International Conference on Machine Learning (ICML), 2022.
[37] Jang, H., Kim, M., and Ahn, S. Learning energy decompositions for partial inference of GFlowNets. International Conference on Learning Representations (ICLR), 2024.
[38] Jing, B., Corso, G., Chang, J., Barzilay, R., and Jaakkola, T. Torsional diffusion for molecular conformer generation. Neural Information Processing Systems (NeurIPS), 2022.
[39] Kim, M., Ko, J., Zhang, D., Pan, L., Yun, T., Kim, W., Park, J., and Bengio, Y. Learning to scale logits for temperature-conditooaa GFlowNets. arXiv preprint arXiv:2310.0282,, 2023.
[40] Kim, M., Yun, T., Bengio, F.,, Zhang, D., Bengio, Y., Ahn, S., and Park, J. Local search GelowNets. International Conference on Learning Representations (ICLR), 2024.
12

---

[41] Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. International Conference on Learning Representations (ICLR). 2014.
[42] Lahlou, S., Deleu, T., Lemos, P., Zhang, D., Volokhova, A., Hernandez-Garcia, A., Ezzine, L. N., theory of continuous generative flow networks. International
Bengio, Y., and Malkin, Conference Machine Learning (ICML), 2023.
[43] Lemos, P., Malkin, N., Handley, W., Bengio, Y., Hezaveh, Y., and Perreault-Levasseur
L. Improving gradient-guided nested sampling for posterior inference. arXiv preprint arXiv:2312.03911,000 2023.
[44] Madan, K., Rector- Brooos. J., Korablyov, M., Bengio, E., Jain, M., Nica, A., Bosc, T., Bengio, Y., and Malkin, N. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference Machine Learning (ICML), 2022.
[45] Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Trajectory balance: Improved credit assignment in gflownets. Neural Information Processing Systems (NeurlPS), 2022.
[46] Malkin, N., Lahlou, S., Deleu, T., Ji, X., Hu, E., Everett, K., Zhang, D., and Bengio, GFlowNets and variational inference. International Conference on Learning Representations (ICLR), 2023.
[47] Máté, B. and Fleuret, F. Learning interpolations between Boltzmann densities. Transactions on Machine Learning Research (TMLR), 2023.
[48] Mittal, S., Bracher, N. L., Lajoie, G., Jaini, P., and Brubaker, M. A. Exploring exchangeable dataset amortization for bayesian posterior inference. In ICML 2023 Workshop on Structured Probabilistic Inference { Generative Modeling, 2023.
[49] Møller, J., Syversveen, A., and Waagepetersen, R. Log Gaussian Cox processes. Scandinavian Journal of Statistics, 25(3):451 482, 1998. ISSN 0303-6898.
[50] Nichol, A. and Dhariwal, P. Improved denoising diffusion probabililstic models. International Conference on Machine Learning (ICML), 2021.
[51] Nicoli, K. A., Nakajima, S., Strodthoff, N., Samek, W., Müller, K.-R., and Kessel, P. Asymptot- ically unbiased estimation of physical observables with neural samplers. Physical Review E, 101(2):023304, 2020.
[52] Noé, F., Olsson, S., Köhler, J., and Wu, H. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457))ceawv 147, 2019.
[53] Nüsken, N. and Richter, L. Solving high- nigh-dimmesiooal Hamilton-Jacohe -elllan PDEs using neural networks: perspectives from the theory of controlled diffusions and measures path space. Partial Differential Equations and Applications, 2(4):48, 2021.
[54] Øksendal, B. Stochastic Differential Equations: An Introduction with Applications. Springer, 2003.
[55] Pan, L., Malkin, N., Zhang, D., and Bengio, Y. Better training of GFlowNetts with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023.
[56] Pillai, N. S., Stuart, A. M., and Thiéry, A. H. Optimal scaling and diffusion limits for the langevin algorithm in high dimensions. The Annals of Applied Probability, 22(6), December 2012.
[57] Radev, S. T., Mertens, U. K., Voss, A., Ardizzone, L., and Köthe, U. Bayesflow: Learning complex stochastic models with invertible neural networks. IEEE transactions on neural networks and learning systems, 33(4):1452-1465 2020.
[58] Ranganath, R., Gerrish, S., and Blei, Black box variational inference. Artificial Intelligence and Statistics (AISTATS), 2014.
13

---

[59] Rector-Brooks J., Madan, K., Jain, M., Korablyov, M., Liu, C.-H., Chandar, S., Malkin, N., and Bengio. Thompson sampling for improved exploration GFlowNets, arXiv preprint arXiv:2306. 17693, 2023.
[60] Rezende, D. and Mohamed, Variational inference with normalizing flows. International Conference Machine Learning (ICML), 2015.
[61] Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. International Conference on Machine Learning (ICML), 2014.
[62] Richter, L., Boustati, A., Nüsken, N., Ruiz, F. J. R., and Ömer Deniz Akyildiz. VarGrad: low-variance gradient estimator for variational inference. Neural Information Processing Systems (NeurIPS), 2020.
[63] Richter, L.,, Berner, J., and Liu, G.-H. Improved sampling via learned diffusions. International Conference Learning Representations (ICLR), 2023.
[64] Roberts, G. O. and Rosenthal, J. S. Optimal scaling of discrete approximations to langevin diffusions. Journal of the Royal Statistical Society: Series (Statistical Methodology), 60(1): 255-268, 1998.
[65] Roberts, G. O. and Tweedie, R. L. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli, pp. 341-363, 1996.
[66] Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
[67] Särkkü, S. and Solin, A. Applied stochastic differential equations. Cambridge University Press, 2019.
[68] Shen, M. W., Bengio, E., Hajiramezanali, E., Loukas, A., Cho, K., and Biancalani, T. To- wards understanding and improving GFlowNet training. International Conference on Machine Learning (ICML), 2023.
[69] Skilling, J. Nested sampling for general Bayesian computation. Bayesian Analysis, 1(4):833 859, 2006. doi: 10. 1214/0116662 BA127. URL https://doi org/10. 100. 1214/06-BA127.
[70] Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., and Ganguli. S. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning (ICML), 2015.
[71] Song, Y., Durkan, C., Murray, I,, and Ermon, S. Maximum likelihood training of score-based diffusion models. Neural Information Processing Systems (NeurIPS), 2021.
[72] Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations (ICLR), 2021.
[73] Tiapkin, D., Morozov, N., Naumov, A., and Vetrov, D. Generative flow networks as entropy- regularized RL. arXiv preprint arXiv:2310.1299: 2023.
[74] Tripp, A., Daxberger, E., and Hermández-Lobato, J. M. Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Neural Information Processing Systems (NeurIPS), 2020.
[75] Tzen, B. and Raginsky, M. Neural stochastic differential equations: Deep latent Gaussian models in the diffusion limit. arXiv preprint /905. 2019.
[76] Tzen, B. and Raginsky, M. Theoretical guarantees for sampling and inference in generative models with latent diffusions. Conference on Learning Theory (CoLT), 2019.
14

---

[77] van Kricken, E., Thanapalasingam, T., Tomczak, J., van Harmelen, F., and ten Teije, A. A-NeSI: scalable approximate method for probabilistic neurosymbolic inference. Neural Information Processing Systems (NeurIPS), 2023.
[78] Vargas, F., Grathwohl, W. and Doucet, A. Denoising diffusion samplers. International Conference on Learning Representations (ICLR), 2023.
[79] Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Transport meets variational inference: Controlled Monte Carlo diffusions. International Conference Learning Representations (ICLR), 2024.
[80] Vincent, P. A connection between score matching and denoising autoencoders. Neural compu- tation, 23(7):1661-1677 2011.
[81] Wu, H., Köhler, J., and Noé, F. Stochastic normalizing flows. Neural Information Processing Systems (NeurIPS), 2020.
[82] Xiao, Z.,, Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion GANs. International Conference on Leraning Representations (ICLR), 2022.
[83] Zhang, D., Malkin, N., Liu, Z., Volokhova, A., Courville, A., and Bengio, Y. Generative flow networks for discrete probabilistic modeling. International Conference on Machine Learning (ICML), 2022.
[84] Zhang, D., Chen, R. T. Q., Malkin, N., and Bengio, Y. Unifying generative models with GFlowNets and beyond. arXiv preprint arXiv:2209.0260 02606, 2023.
[85] Zhang, D., Rainone, C., Peschl, M., and Bondesan, R. Robust scheduling with GFlowNets. International Conference on Learning Representations (ICLR), 2023.
[86] Zhang, D., Chen, R. T. Q., Liu, C.-.., Courville, A., and Bengio, Y. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization. International Conference Learning Representations (ICLR), 2024.
[87] Zhang, Q. and Chen, Y. Diffusion normalizing flow. Neural Information Processing Systems (NeurIPS), 2021.
[88] Zhang, Q. and Chen, Y. Path integral sampler: stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022.
[89] Zhu, Y., Wu, J., Hu, C., Yan, J., Hsieh, C.-Y., Hou, T., and Wu, J. Sample-efficient multi- objective molecular optimization with GFlowNets. Neural Information Processing Systems (NeurIPS), 2023.
[90] Zimmermann, H., Lindsten, F,, van de Meent, J.-W., and Naesseth, A. A variational perspective on generative flow networks. Transactions Machine Learning Research (TMLR), 2023.
15

---

A Code and hyperparameters
Code is available at https: htap:s///githbb.... httpg/gttw...o////gggggggg/aaeee wtw.aoorrggnnddffusiioo and will continue to be main- tained and extended.
Below are commands to reproduce some of the results on Manywell and VAE with PIS and GFlowNet models as example, showing the hyperparameters: PIS:
--mode_fud pis --lr_policy 1e-3
PIS + Langevin:
--moe_ffud pis --Ir_policy 1e-3 --langevin
GFlowNet TB:
python train.py
--t_scale 1. --energy many_well ariisarchitectures --zero_init --clipping
--mode_fwd tb --lr polpccy 1e-3 --lr flow 1e-1
GFlowNet TB Expl.:
python train.py
--t_scale 1. --energy many_well --pis_aree ariittrrhit -zero_init --clipping --mode_fwd fwd tb --1r --r_ppoiccy 1e-3 --lr_flow 1e-1 -exploratory -expooration -axppooratioo_factorr factor 0.2
GFlowNet VarGrad Expl.:
python train.py
--t_scale 1. --energy many_well --pis_ priiittrcisssstees -erooinnt --clipping
--mod__fwd tb-avg --rrpoy 1e-3 -ll_fllow 1e-1
exxpooatoory exxpporation -exploration factor 0.2
GFlowNet FL-SubTB:
python train.py
--t_scale 1. --energy many_well ariiisarchitettres --zero_inis --clipping --mode_fwd subtb --Ir_policy 1e-3 --lr_flow 1e-2 --partial_energy -conditional flow_model
FL-SuTTB LP:
python train.py
--t_scale 1. --energy many_well ariiittrchiesssress --zero_inis init --clipping --mode_fwd subtb --lr_poolice 1e-3 --lr_flon flow 1e-2 --partial_ condiitional fooomallfloummmll
--langevin --epochs 10000
GFlowNet TB Expl. + LS:
python train.py --t_scale 1. --energy many many_well --iss_pcciteetttiiee architectures --eer_init --clipping fod tb --lr lrrppolcccy 1e-3 --lr_back 1e-3 -lr_flow 1e-1
-exploratory -exploration_wt wd -wxpporatioo_ffcaeeeeccc factor 0.1
--both_war bboth_was ways --local_searde
-buffer_size 600000 --prioritized rank --rank_weight 0.01
--id_step 0.1 --Id_schedule
0.574
GFlowNet TB Expl. + L.P:
16

---

python train.py
--t_scale 1. --energy many_well --pis_ aris_aarcitectures --zero_inii init --clipping
--mode_ fwd tb --lr pollicy 1e-3 --lr_floo 1e-1
-exploratory --exploration - factor 0.2
--langevin --epochs 10000
GFlowNet TB Expl. + LS (VAE):
python train.py
--energy vae appis_architectures --zero_ init --clipping --mode_fvd fwd cond-tb-avg --mode bud cond-tb-avg --repeats 5 --lr --r_pooicy 1e-3 --lr_flow 1e-1 --lr back 1e-3
-exploratory -=xxpooaaiion_ exploration_ factor 0.1
--both_ways --local_searee
-mmax_iee__ls 500 burunnnn 200
--buffer_size 90000 pprioiitized rank --rank_weight 0.01
--ld_step 0.001 --hd_schedule -argget_acceptance.eee 0.574
GFlowNet TB Expl. + I.P LS (VAE):
python train.py
--energy vae --pis_architectures --zero_ --zero_init --clipping --mode_fwd cond-tb-avg -mode_bwd cond-tb-are r-repeats 5 -Ir__poicyy 1e-3 --1r_flow 1e-1
--lgv_clip 1e2 --ff_cclippppppppp 1e4 --epochs 10000
-exploratory -exploration_wd exploration_factor 0.1
--both_ways --local_search
--Ir -lr_back 1e-3 --max_iter 500 --burn_in in 200
--buffersiize 90000 --prioritized rank -rrakk_wiigt 0.01 --langevin
--ld_step 0.001 --ld_schedule -acgget. 0.574
17

---

B Target densities
Gaussian Mixture Model with 25 modes (25GMM). The model, termed as 25GMM, consists of a woodiiensionall Gaussian mixture model with 25 distinct modes. Each mode exhibits an identical variance of 0.3. The centers of these modes are strategically positioned on grid formed by the Cartesian product [-10,- 5,00, 0,5, 10) x (-10, -5, 0,5, 10), effectively distributing them across the coordinate space.
Funnel [29]. The funnel represents classical benchmark in sampling techniques, characterized by ten-dimensional distribution defined as follows: The first dimension, Xo, follows a normal distribution with mean 0 and variance 9, denoted xo (0. 9). Conditional Xn, the remaining dimensions, X1:9, are distributed according to a multivariate normal distribution with mean vector and a covariance matrix exp(x)) where I the identity matrix. This succinctly represented as X1:9 Xo N (0, exp (xo) I).
Manywell [52]. The manywell is characterized by 32-dimensional distribution, which is con- structed as the product of 16 identical two-dimensional double well distributions. Each of these two-dimensional components is defined by potential function, xxx1, expressed as x2) = exp (-x + 6x} + 0.5x, - 0.5x3)-
VAE [41]. This task involves sampling from 20-dimensional latent posterior p(z|x) px2), where pz2) a fixed prior and p(x[z) a pretrained VAE decoder, using conditional sampler (2//x) dependent on input data (image)
LGCP [49]. This density over (600-dimensional variable is Mog-Gaussian Cox process fit to distribution of pine saplings in Finland.
B.1 Discrepancies in past work
Wrong definitions of the Funnel density, As already noted by [78], [88] uses a different variance of the first component in the Funnel density, 1 instead of 9. This apparent bug in the task definition has been propagated to subsequent work, including [42].
Evaluation on LGCP. The LGCP benchmark suffers from the lack of consistent ground truth log Z to compare against. Previous work has compared the value of the partition function log Z against a "long run of Sequential Monte Carlo" [88]. We note that this approach produces noisy estimates of the partition function, especially high-dimensional problems (indeed, SMC has rarely been used in problems with over thousand dimensions); therefore, it is unclear how long the SMC needs to be run to produce an accurate estimate. We found that two different values are being used in the literature: log Z 512.6 in one repository and log Z 501.8 in another.
On FL--SubTB as used in [86]. We make two observations calling into question the main results of [86].
First, the only substantial difference between the algorithm used by [86] and the one from the past work [42] - which first proposed the use of GelowNet objectives to train diffusion samplers is the substitution of the FL-SubTB objective [55, 44] for TB [45]. However, [86] elects to compare FL-SubTB with the Langevin parameterization to TB without the Langevin parameterization. Our results in Table show that while the Langevin parameterization crucial for the performance of all objectives; FL-SubTB does not provide any consistent benefit over TB or VarGrad.
Second, the results are not reproducible, neither with the published code from [86] run 'out of the box', nor with our reimplementation. In particular, on the LGCP density, the training did not converge within the allotted training time. We have contacted the authors of [86], who confirmed that running their published code does not reproduce the results in the paper but could not provide any further explanation working implementation.
C
Additional results
C.1 Expanded unconditional sampling results
Table C.1 is an expanded version of Table 1, showing Wasserstein distances between sets of K samples from the true distribution and generated by a trained sampler. (Note that ground truth for LGCP is not available.)
18

---

Table C.I: Log-pas log-partition function estimation errors and 2- Wasssesttein distances for unconditional modeling tasks (mean and standard deviation over runs). The four groups of models are: MCMC- based samplers, simulation-driven variational methods, baseline GFlowNet methods with different learning objectives, and methods augmented with Langevin parametrization and local search.
Energy
25GMM (d 2)
Funnel (d 10)
Manywell (f - 32)
Algoritaa
A Z
Alog 2e
W
Alog 2
W
SMC
0.86u210
0.561
14,99_1288 0.292x644
[43]
0.016a004
1.39ua37
0:033381175
0.51
[8]
1.125 PS6
0.986.
4.71
0.830-
0.09%00a
20.7-3.1
10.52+1 on
3,05-0 is
5,08-0
DDS [78]
1.760
0 24660..
7.18
0,424
0,206
29.3
4.362
0.23
5.7 is
[88]
1,769-m 181
1.23400 21w
6.37+0
0.534-c.00
0.262-0aa
22.0-1.
2.69-0:1
6.55
(88[
1.799
0 233
7.16.211
0,587
0,285etaass
22.
13.19 13.19
0.07
6 6,55
11 [42]
1.176un.109
102
4.83us
0.660accam
0.25
22.4.4.0
4,01.am
2.67
TD Expl. (42)
0.560un.252
0:422.0.53
361141
0.769.cans
0.226.
21.3,40
4,01.0m
2,68.003
VarGind Expl.
0.615un241
0.487a0.25
3,89uc35
0.642xcate
0.250.2.112
22.1.4.
4,01.ame
2.69
6.15
F1--ShhTB
1.127=1000
1200w
4.640.00
0.527-cam1
0.18240.0
22.1+1.3
3.98.307
2.72-0.5
6.15
+L. [86]
0.509-a.c0s
ASee
0,563-cam1
0. 0.155-ca17
22.2
4.23+3
2.66-022
6 10
TB 155 (ours)
0.
0.0440
1,250
0.653-cam5
0.285-cam
21.9-1.
4.57
566
TB Expl. LP (ours)
0.206un.sss
0.01 Lo
1229a.0v
0.666.5412
0.051
22.3
7.661.04
1.
5.73a031
(ours)
190-0.00
0000
1.31+0
0,768-cash
0,264-nas
21.8-3.
533
VarGral + Expl. + LP LS (ours)
0.207
0.015
1.33as.8
0.920.s.11
0,256.va
21.2,A.
4.11.
0.02
5.30a0m
Highlight:: mean indistinguishable from minimum in column with p < 0.05 under one-sided Welch unpaired r-test.
Table C.2: Log-pes partition function estimation errors and empirical 2- 2-Sasserstein distances on the 32-dimensional Manywell with Brownian and variance-preserring noising processes.
Backward process- +
Brownian
VP
Objective J Metric
Alog
A log ZRW
W2
A loo Z
llog RR
W?
TB + Expl. LP
7.46=1.74
1.0621.11
5.7320.31
7.55+2.85
1.49
5.68=0.42
TB + Expl. LP + LS
4.68:0.49
0.07
5.33+0.03
4.52+0.21
1.23+0.07
5.75=0.01
VarGrad + Expl.
4.01
2.69=0.06
6.15+0.02
4.04+0.05
2.65+0.08
6.17=0.02
C.2 Bariance-preserving noising process
Following the recent results by [8, 63, 78], we perform additional set of experiments with different successful noise schedule. We replace the Brownian motion by the variance- variance-preserving SDEs from Song et al. [72]. given by Orastein-Uhlenbeed process:
(1)) = 2 and p(x,1) == -B(1xx
(16)
with E oo).
In particular, we follow the common procedure use y := I and B(t) == (1 - 1)Bmin + +fmax. I E [0,
with Bmin 0.01 and Bmax 4.0.
We evaluate three representative methods using this variance-preserving backward process. The results, in Table C.2, are similar to those using the Brownian bridge process. We expect that the choice of noising process gains importance in challenging high-dimensional problems.
C.3 Scalability study
The Manywell energy ($B) is defined in any even number of dimensions and thus allows to study the scaling of the methods with dimension. We evaluate several representative methods in dimension 8, 128. and 512 (in addition to the 32 studied in the main text). All experimental settings are kept the same as for d 32. Due to the large runtime, some runs in dimensions 128 and 512 had to be limited at 12 hours, while dimensions 8 and 32 all run in under 3 hours on a RTX8000 GPU.
These results are shown in Table C.3. We observe:
The overhead of the Langevin parametrization grows with dimension, but critical to performance.
The even higher overhead of FL-SubTB as used by [86].
The relatively high efficiency and low overhead of our newly proposed local search.
19

---

Table C.3: Scaling with dimension on Manywell: log- reg-partition function estimation errors and time per training iteration on RTX8000 GPU.
Dimension
=3
d= 512
TL2ATT & 2
Objective 1. Metric-
A Alog2 Z
A Alg ZAW
Alog Z
A Z
Alog Z
Alog ZMWW
Alog Z
PIS + LP [PS] TB
0.86
0.14
13.19
0:07
58,0
23.7
251
169
0.95
0.70
4,01
2.68
205.6
119.8
1225
957
FL-SubTB 186]
2.66
198
10"
7:46
1,86
259
TB -.
0.44
0.00
4.57
0.19
458.7
139.3
1626
1677
TR 1.P+18
0.25
0.02
4.68
0.07
56.6
14.9
326
209
10
120
Divention
D Experiment details
Sampling energies. In this section, we detail the hyperparameters used for our experiments. An important parameter is the diffusion coefficient of the forward policy, which is denoted by or and also used in the definition of the fixed backward process. The base diffusion rate o² (parameter t_scale) set to 5 for 25GMM and I for Funnel and Manywell, consistent with past work.
For LGCP, we found that using too small diffusion rate o² (e.g., o² 1) prevents the methods from achieving reasonable results. We tested different values of o² = 3,),, and selected  5, which gives the best results, which follows the findings in Zhang & Chen [88].
For all our experiments, we used learning rate of 10-3. Additionally, we used higher learning rate for learning the flow parameterizaaton, which is set as 10 when using the TB loss and 10- with the SubTB loss. These settings were found to be consistently stable (unlike those with higher learning rates) and converge within the allotted number steps (unlike those with lower learning rates).
For the SubTB loss, we experimented with the settings of 10x lower learning rates for both flow and policy models communicated by the authors of [86], but found the results to be inferior both using their published code (and other unstated hyperparameters communicated by the authors) and using our reimplementation.
For models with exploration, we use an exploration factor of 0.2 (that is, noise with a variance of 0.2 is added the policy when sampling trajectories for training), which decays linearly over the first half of training, consistent with [42].
We train all our models for 25, 000 iterations except those using Langevin dynamics, which are trained for 10, 000 iterations. This results approximately equal computation time owing to the overhead from computation of the score at each sampling step.
We use the same neural network architecture for the GFlowNet as one of our baselines [88]. Similar to [88], we also use an initialization scheme with last-layer weights set to O at the start of training. Since the SubTB requires the flow function to be conditioned on the current state x, and time 1, we follow [86] and parametrize the flow model with the same architecture as the Langevin scaling model NN₂ [88]. Additionally, we perform clipping on the output of the network as well as the score obtained from the energy function, typically setting the clipping parameter of Langevin scaling model to 102 and policy network to 10t, similarly to [78]:
= NN (k, 0) + NN2 0) o clip(VInn(x) -10?, -10t, 110*.10).
(17)
All models were trained with a batch size of 300. In each experiment, train models on a single NVIDIA A A100-Large GPU, if not stated explicitly otherwise.
VAE experiment. In the VAE experiment, we used standard VAE model pretrained for 100 epochs on the MNIST dataset. The encoder (2(xx) contains an input linear layer (784 neurons) followed by hidden linear layer (400 neurons), ReLU activation function, and two linear heads (20 neurons each) whose outputs were reparametrized to means and scales of multivariate Normal distribution. The decoder consists of 20-dimensional input, one hidden layer (400 neurons), followed by the ReLU activation, and 784-dimnnsi output. The output is processed by the sigmoid function to be scaled properly into [0. l].
The goal to sample conditionally on x the latent z from the unnormalized density p(z, pxx)p(x z) (where p(z) is the prior and p(x|z) is the likelihood computed from the decoder), which
20

---

is proportional to the posterior p(z x). We reuse the model architectures from the unconditional sampling experiments, but also provide x as an input to the first layer of the models expressing the policy drift (as well as the flow, for FI.-SubTB) and add one hidden layer to process high- nigh-dimensional conditions. For models trained with TB, log Z. also becomes a MLP taking as inpur.
The VarGrad and LS techniques require adaptations in the conditional setting, For LS, buffers (Dbufier and Dis) must store the associated conditions together with the samples z and the corresponding unnormalized density R(z:x), i.e., a tuple of (x, z, R(zix)). For VarGrad, because the partition function depends on the conditioning information x, necessary to compute variance over many trajectories sharing the same condition. We choose to sample 10 trajectories for each condition occurring in minibatch and compute the VarGrad loss for each such set of 10 trajectories.
The VAE model was trained on the entire MNIST training set and never updated on the test part of MNIST. In order to evaluate samplers (with respect the variational lower bound) set of examples, we chose the first 100 elements of MNIST test data. All of the samplers were trained having access the MNIST training data and the frozen VAE decoder. For fair comparison, samplers utilizing the LP were trained for 10, 000, whereas the remaining for 25, 000 iterations. In each iteration, batch of 300 examples from MNIST was given conditions. In each experiment, we train models on single NVIDIA A AlrgeL GPU, if not stated explicitly otherwise.
unique
7 2
7
2
0
4
4
9
7
0
4
4
9 5
9 0
9
0
5
5
9
0
6
9
0
1
5
S 9
0
0
J
(a) Conditioning data (MNIST test (b) VarGrad Expl. L.P samples set)
(c) VAE reconstruction
decoded
Figure D.1: Our sampler (VarGrad + Expl. LP) conditioned by a subset of never-seen data coming from the ground truth distribution (left). The conditional samples were then decoded by the the fixed VAE (middle). For the comparison, we show the reconstruction of the real data by VAE (right). We observed that the decoded samples are visually very similar to the reconstructions making these two pictures almost indistinguishable. Both, decoded samples and reconstruction, are more blurry than the ground truth data, which is caused by a limited capacity of the VAE's latent space.
21

---

E Local search-guided GFlowNet
Prioritized sampling scheme. We can use uniform or prioritized sampling to draw samples from the buffer for training. We found prioritized sampling to work slightly better in experiments (see ablation study $F.2), although the choice should be investigated more thoroughly in future work.
We use rank-based prioritization [74], which follows probabilistic approach defined as:
p(x: Drfffe) & (bufeel + rank
(x)))²
(18)
where rank Doutc, (x) represents the relative rank of sample x based on ranking function R(x) (in our case, the unnormalized target density at sample x). The parameter k is hyperparameter for prioritization, where a lower value of k assigns a higher probability to samples with higher ranks, thereby introducing more greedy selection approach. We set = 0.01 for every task. Given that the sampling proportional the size of Duffer, we impose a constraint on the maximum size of the buffer: DDanter 600, 000 with first-in first out (FIFO) data structure for every task, except we use Dautfe 90, 000 for VAE task. See the algorithm below for detailed pseudocode.
Algorithm 1 GFlowNet Training with Local search
I: Initialize policy parameters for Pp, and empty buffers Douffer. DLs
2: for / 1. do if /%k2 = then
Sample M trajectories (TI...., TM] ~ Pp( (p(ee--gredy)) 5: Update Druffer <- Douffer U (x|T - x]
6:
Minimize L(T; using [T).-..,TM) else
TM) to update Pf
7:
if 100 0 then
9:
Sample (X].....IM} ~ Dbuffer DLS - Local Search( (x]...... .MM DLS) end if
10:
11:
12:
Sample (xf......'' x MM (.uffer ;Dss)
13: Sample (T......Th) M PB(-.. (x')
Minimize ((T';0) using (Tf..... 3 to update PF
15:
end if
16:
end for
We use the number of total iterations / = 25, 000 for every task default. Note local search is performed to update D1.s occasionally that per 100 iterations, the number local search updates done 25, 000/ 100 250.
22

---

E.1 Local search algorithm
This section describes detailed algorithm for local search, which provides an updated buffer DLs. which contains low-energy samples.
Dynamic adjustment of step size 17. To enhance local search using parallel MALA, we dynamically select the Langevin step size (77). which governs the MH acceptance rate. Our objective is to attain an average acceptance rate 0.574, which theoretically optimal for high-dimensional MALA's efficiency [56]. While the user can customize the target acceptance rate, the adaptive approach eliminates the need for manual tuning.
Computational cost of local search. The computational cost of local search is not significant. Local search for iteration of K = 200 requires 6.04 seconds (averaged with five trials in Manywell), where we only occasionally (every 100 iterations) update DLs with MALA. The speed is evaluated using the computational resources of the Intel Xeon Scalable Gold 6338 CPU (2.00GHz) and the NVIDIA RTX 4090 GPU.
Algorithm 2 Local search (Parallel MALA)
input Initial states {x (0)
M/ J. current buffer DLs, total steps K, burn in steps Kburninns initial step size 7o.- amplifying factor fincrease: damping factor fdecreuse: unnormalized target density R output Updated buffer DLs
(0), x
Initialize acceptance counter a =
Set ?70 for k 1: K do
Initialize step acceptance count dg for M do
Sample o N(O, (k-1) +1 log R(xm
Propose x - Xm
(K-1))
+ + 2mper
R(x exp
m7 |m 12-1) X log R(x,,, ||
Compute acceptance ratio min
R(x(--11 exp +
(1-1)
N log R R(x (2 112
Xmm
With probability r. accept the proposal: x1k)) + xm and increment dj ax +1 if & > bum-in then Update buffer: DLs < DLs U (xm) end if end for
Compute step acceptance rate arg ugk/
if cr > Otarge then
+ 17 x fincrease
else if og Ottarget then
1) t 17 x fdecrease end end for
We adopt default parameters: fincrese = 1.1, fdecrase = 0.9, 1j 0.01, K 200, K 100, and Ouarget = 0.574 for three unconditional tasks. For conditional tasks of VAE, we give more iterations of local search: 500, K in 200.
is noteworthy that by adjusting the inverse temperature B into R during the computation of the Haaropolii-Hastingss acceptance ratio r,
facilitate greedier local search strategy aimed at exploring samples with lower energy (i.e., higher density Ptarget). This approach proves advantageous for navigating high-dimensional and steep landscapes, which are typically challenging for locating low-energy samples. For unconditional tasks, we set B 1.
In the context of the VAE task (Table 2). utilize two GFlowNet loss functions: TB and VarGrad.
For local search within TE, we set B = 1, while for VarGrad, we employ / 5. As illustrated in Table 2, employing local search with / 1 fails enhance the performance of the TB method. Conversely, a local search with B = results in improvements at the log Rww metric over the VarGrad Expl. + L.,, even though the performance of VarGrad Expl. L.P surpasses that of TB substantially. This underscores the importance of selecting an appropriate f value, which critical for optimizing the exploration-exploitttio balance depending the target objectives.
23

---

E.2 Ablation study for local search-guided GFlowNets
Increasing capacity of buffer. The capacity of the replay buffer influences the duration for which retains past experiences, enabling helps in preventing mode collapse during training. Table E.1 demonstrates that enhancing the buffer's capacity leads to improved sampling quality. Furthermore, Figure 1 illustrates that increasing the buffer's capacity-thereby encouraging the model to recall past low-energy eperiences-eehhnces its mode-seeking capability.
replay these experiences the policy. This mechanism
Table E.1: Comparison of the sampling quality of each sampler trained with varying replay buffer capacities in Manywell. Five independent runs have been conducted, with both the mean and standard deviation reported.
Buffer Capacity Metric
Alog
A log ZRW
W2
30, 000
40.1
2.73
6.17=0.02
60, 000
4.06+0.05
2.38=0.38
6.1440.04
600, 000
4.57=2.13
0.19-0.29
5.66+0.05
(a) Capacity 30, 000
(b) Capacity 60, 000
(c) Capacity 600, 000
Figure E.1: Illustration of each sampler trained with varying capacities of replay buffers, depicting 2,000 samples. As the capacity of the buffer increases, the number of modes captured by the sampler also increases.
Benefit of prioritization.
Rank-prioritized sampling gives faster convergence compared with no prioritization (uniform pling), as shown in Fig. E.2a
Dynamic adjustment of 1 fixed 17 0.01. As shown in Fig. E.2b, dynamic adjustment to target acceptance rate Otarget = 0.574 gives better performances than fixed Langevin step size showcasing the effectiveness of the dynamic adjustment.
otliou Bine S/Sedkdin
254
1.5
140
pricrtbuttinn
35000
(a) Benefit of prioritization replay buffer sampling..
(b) Benefit of scheduling dynamically.
Figure E.2: Ablation study for prioritized replay buffer and step size n scheduling of local search. Mean and standard deviation are plotted based on five independent runs.

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
See theory and experiment sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, see section 5.3 and conclusion, as well references appendix material where relevant.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model vell-specifiiaatin, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
25

---

No new theoretical results. For exposition of the mathematical basis for our algorithms, we state the assumptions.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: See experiment sections and references to appendix material.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
26

---

Answer: [Yes]
Justification:
We provide code to reproduce nearly all of our experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / /nips. wwtppcccccc/cceeccccccccccccclct hataaiiguddeeeedeuuniiieeeeeee
/Coo//eoomissssseeee crriissssoppiiyyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
See the experiment sections and references to appendix material.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All results tables and plots show standard deviation and indicate significance of the best metric.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
27

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
See experiment sections and references to appendix material.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the ines?
NeurIPS Code of Ethics https:: h
Answer: [Yes]
Justification:
We believe there are no violations of the CoE.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
The paper studies ML problem with immediate societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
28

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper studies ML problem with no immediate application to generation new image or text content, nor other functions that have the potential for misuse, to the best of our knowledge.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We cite the works introducing all datasets we study.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi twwode.cooo..c.. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
29

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mooel as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human studies.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human studies.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "n3BXFQ6tik", "reviewer": "Reviewer_5SPJ", "summary": "The paper presents a variety of improvements to off-policy strategies for training diffusion models to sample from unnormalized densities. Equation 13. These include maintaining a replay buffer (obtained with Langevin sampling) to enable efficient off-policy exploration and incorporating an inductive bias into the neural network which estimates the SDE drift term. They also present a software library containing unified implementation of these techniques and e.g. diffusion model training.", "questions": "See weaknesses. My main concern is the underwhelming results when compared to DIS, DDS, PIS, PIS+LP. Is there any reason why the proposed method(s) should be preferred to these? Or e.g. expected to scale better to more complex problems?", "limitations": "Adequately addressed", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Enabling diffusion models to be efficiently applied to sampling from unnormalized probability distributions is a problem with high potential for impact\n- Thorough experimental analysis and comparison of different alterations to the training procedure. On most problems considered, the authors' contributions are necessary to achieve good results with trajectory balance. \n- The contribution of a software library could be valuable to the community.", "weaknesses": "- The results are not overwhelming - although the proposed contributions are helpful compared to a basic version of TB, there is only one modeling task in Tables 1-2 (25GMM) where they provide a statistically significant improvement over the baselines.\n- The experiments are on synthetic energy functions and MNIST VAE. Including more real-world data or models would be informative.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below we will answer your questions and concerns.\n\n### Statistically significant improvement over the baselines\n\nFirstly, we want to highlight that our method achieved comparable or better results to current SOTA models. In Tables 1 and 2, we highlighted the best results (red) and those that are not statistically-significantly worse according to the Welch test (blue). Whereas the baselines achieve good results usually in one or two tasks, our method and its variations are the best or among the best results in **nearly all tasks simultanously**.\n\n### Experiments on synthetic energy functions\n\n\nWe thank the reviewer for their suggestion on including the real-world sampling example. Note that the LGCP task is already a real-world example (it is derived from observations of pines in Finland). However, we recognize the importance of applications and will add more experiments (e.g., molecular conformations) in the revision. \n\nWe kindly direct you to the response to all reviewers for further discussion of the choice of densities.\n\n### Underwhelming results when compared to SOTA?\n\nSimulation-based methods, such as PIS, are predominantly on-policy, making it challenging to explore low-energy regions in high-dimensional spaces and prone to mode collapse in complex densities (see Figure 1 in the Manywell task). In contrast, training a sampler using off-policy techniques allows for the utilization of powerful methods like local search, which enhances exploration and performance. This off-policy trick has been verified in various problems, including simple energy (25GMM), complex multimodal energies (Manywell), high-dimensional tasks (LGCP), and conditional posterior sampling tasks (VAE). \n\n**Thanks again for your feedback. Please let us know if you have any more questions and if any further clarifications would help you assess the paper.**", "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": "n4w6hwI5Rv", "reviewer": "Reviewer_6eYp", "summary": "This paper focuses on the problem of sampling with distributions defined by a black-box and unnormalized energy function. This work provides a comprehensive review of existing works, including both variational methods and policy-based methods, and offers a codebase and benchmark to replicate and evaluate the existing works. Additionally, this work proposes a method to improve existing policy-based methods via local search and a replay buffer.", "questions": "Is there any benchmark in this field involving a real application and complex high-dimensional distribution?", "limitations": "See my comments on the weakness above.", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The studied problem of sampling from a distribution is an important issue with a long history in statistical inference. The paper provides a good review of recent works on this topic by leveraging diffusion models. The codebase that unifies existing methods is certainly useful to the community for continuing research on this topic.\n\n2. The experiments are comprehensive in baselines, including not just diffusion-based methods but also classical MCMC algorithms. The results clearly show the advantages of diffusion-based methods and the techniques proposed in this work.", "weaknesses": "1. It appears to me that this work only tests the algorithm on relatively simple and manually-constructed scenarios. Are there any real and important applications within the field? I am not very familiar with this field, but I think that only conducting experiments on synthetic datasets makes this topic less practical. I believe the main advantage of the diffusion-based method over classical methods is in modeling complex distributions, making experiments on synthetic examples less meaningful.\n\n2. Additionally, the tested scenarios are all low-dimensional cases. I wonder how this algorithm performs on high-dimensional cases, such as when the energy function is learned through neural networks. For example, is it possible to apply this algorithm to image generation where the energy function is represented by an image classifier? Testing the algorithm on high-dimensional tasks like these would provide a better understanding of its scalability and practicality in more complex and realistic settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments, in particular, for acknowledging the strength of our comprehensive benchmarking. \n\nRegarding your questions about the benchmarks and their dimensionality, first, we kindly direct you to the response to all reviewers for discussion of the choice of target densities.\n\nSecond, the dimensions of the tasks studied are: 2, 10, 20 (with 784-dim conditioning vector), 32, and 1600. As just noted, this is in line with past work on diffusion samplers. But it is also worth noting that sampling from complex densities in high dimension **without any prior knowledge** (such as a pretrained model or known samples from the target) is generally an unsolved problem, which is perhaps why it has not been the subject of much study in this past work. Our methods, which use off-policy training, can be combined with prior knowledge, such as known target samples (such as by placing those samples in the replay buffer), as has already been demonstrated, e.g., in [42]. Simulation-based methods like PIS, DDS, etc. do not have this ability. \n\nThe use of diffusion samplers in settings with prior knowledge, while not the subject of this paper, is indeed a very interesting direction of research and likely important in more difficult applications.\n\n**Thank you again for your review. Please let us know if we can provide any other clarification that would help you in understanding and assessing the paper.**", "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": "s6aMp82sTW", "reviewer": "Reviewer_fepi", "summary": "This paper proposes an off-policy diffusion-based sampler training method to match a target distribution and a corresponding exploration strategy and credit assignment to improve it.", "questions": "1.\tThe biological sequence design seems more appropriate to be the validation benchmark for the proposed model, which is also considered in GFlowNets. So the reviewer wonders how the proposed method is compared with GFlowNets in such tasks.\n2.\tCould the authors explain why they use the log-partition function estimation error as the metrics rather than the log-partition function itself? Similar to MLE (Maximum Likelihood Estimation), the model can be considered better with higher $\\log Z$.", "limitations": "See Weakness 3.", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "1.\tThe proposed idea of this paper is interesting, which connects the Euler-Maruyama sampler and GFlowNets.", "weaknesses": "1.\tAlthough the authors mention that traditional MCMC have high cost in sampling, the proposed method based on neural sde seems to still have this problem. To the reviewer’s knowledge, the solving procedure of neural sde is time-consuming as well.\n2.\tThe experimental target distribution also seems relatively simple. In the reviewer’s opinion, for GMM, we can first sample a mode according to the weights of different modes and then obtain a sample in this mode. Hence, it seems unnecessary to use complex model like diffusion.\n3.\tIn many real-world applications like image generation, the pdf (may be unnormalized) of the target distribution is unavailable and we can only achieve data samples from the target distribution. Hence, the application scenarios of the proposed model are limited. \n4.\tBesides, as mentioned in the conditional sampling case, the proposed method seems to need an extra trained vae to perform sampling. However, the vae can directly do the image generation. In that case, what is the real contribution of the proposed method?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. Below we've tried to address what we believe to a number of misunderstandings.\n\n> Strength: The proposed idea of this paper is interesting, which connects the Euler-Maruyama sampler and GFlowNets\n\nFirst, we'd like to point out that this paper is not the first to connect Euler-Maruyama samplers of SDEs with GFlowNets (this was already done in [41] and a few other papers). The goal of this paper is to study methods that, for the first time, achieve results competitive with simulation-based methods on several continuous sampling benchmarks.\n\nNext, we answer the rest of the points:\n\n### Applicability of diffusion samplers is limited\n\nWe respectfully disagree:\n\nYou are correct that there are many domains where we want to train a generative model from ground truth samples. That is an easier problem than sampling unnormalized densities, which is perhaps one reason for the greater attention it receives.\n\nHowever, many important Bayesian inference problems require sampling distributions that are only available through unnormalized density functions, and sampling methods for such distributions are an active area of research. As noted by the other reviewers:\n> The studied problem of sampling from a distribution is an important issue with a long history in statistical inference (**6eYp**)\n\n> Enabling diffusion models to be efficiently applied to sampling from unnormalized probability distributions is a problem with high potential for impact (**5SPJ**).\n\nA few such problems **beyond** Bayesian deep learning are simulations in statistical physics (e.g., [Nicoli et al., \"Asymptotically unbiased estimation of physical observables with neural samplers\", Physical Review E, 2020]), molecular dynamics problems (e.g., [Noé et al., \"Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning\", Science, 2019]), imaging inverse problems in many scientific fields, etc. In fact, it is these problems that have motivated the benchmarks  that are adopted in the literature for **entirely data-free** inference methods such as ours (e.g., Manywell is a simple statistical physics model).\n\n### Neural SDE integration is slow compared to MCMC\n\nThis is inaccurate. The solution of a neural SDE requires as many neural net evaluations as the number of discretization steps chosen (for us, this is 100, in line with past work). This is *far* less expensive than the time to run MCMC chains to convergence on any nontrivial density.\n\n### Target distributions are simple\n\nWe kindly direct you to the response to all reviewers.\n\n### Why test on a Gaussian mixture?\n\nOf course, you are correct that a Gaussian mixture density does not require a diffusion model to sample. However, 2D Gaussian mixture densities have been used as a benchmark in many past works ([41,85,87], among others) to test the ability to discover separated modes in an easily interpretable setting.\n\n### The proposed method seems to need an extra trained vae to perform sampling\n\nThis question appears to rest on a misunderstanding about what that experiment is showing. The goal there is to *learn an encoder for a trained VAE*, which is an established problem in the diffusion samplers literature (e.g., [87]). The goal is **not** to unconditionally sample new images, which could indeed be done by simply decoding noise samples with the pretrained decoder. Here, we need a pretrained decoder to provide the reward for the diffusion encoder that we train (the reward is the joint likelihood -- the product of the decoder likelihood and the prior density).\n\n### Comparing with GFlowNets for biological sequence design\n\nThere seems to be a misunderstanding here. Biological sequence design is a **discrete** sampling problem in which GFlowNets have indeed been used successfully. The goal of this paper is precisely to study the adaptation of GFlowNets and related methods to **continuous** sampling problems.\n\n### $\\log Z$ estimation error\n\nThere may be some misunderstanding about the meaning of this metric. The metrics are **variational lower bounds** on $\\log Z$, meaning that they are lower than the true log-partition function in expectation. The estimation error is a more meaningful metric than the raw value because the best possible value (0) is known and invariant to additive changes to the target log-density, which do not change the target distribution.\n\nWhen the true $\\log Z$ is not known, one can report the estimate, not the estimation error, as we indeed do in some of the experiments (LGCP in Table 1 and VAE in Table 2).\n\n**Thank you again for your feedback. Please let us know if we have misunderstood anything in your comments and we will be happy to provide further clarification. If we have sufficiently clarified these points and answered your questions, please consider updating the score.**", "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": "MLrSnnfzKL", "reviewer": "Reviewer_FAJh", "summary": "The paper studies the problem of training diffusion models to sample from a target distribution. The contributions are summarized as follows: \n\n1. A codebase is provided for the study of diffusion-based samplers, due to the issue of inconsistent experimental settings in previous research; \n\n2. Exploration in the target space can be enhanced by GFlowNet-based off-policy training objectives and local search with the use of replay buffer. \n\n3. Experimental results validate the effectiveness of the proposed approach.", "questions": "- In terms of Table 1, which proposed methods perform best according to $\\log Z^{LB}$ and $\\log Z^{RW}$? Both evaluation metrics reflect different aspects of performance. For example, $\\log Z^{LB}$ should indicate mode collapse behaviour, with lower values suggesting more serious mode collapse (e.g., PIS+LP: 13.19 vs. TB+Expl: 4.01, as illustrated in Figure 1)? In contrast, lower $\\log Z^{RW}$ values suggest a closer approximation to the ground-truth $\\log Z$.\n\n- In terms of C2 variance-preserving noising process, \n\n  - As far as I undertand from the paper, GFlowNet uses PIS architectures, where the initial state is a point mass at 0. How did you design it under VP settings, where the marginal distribution $p_{t}^{ref}$ is an invariant Gaussian distribution, i.e., the initial state is Gaussian distributed, not a point mass at 0? \n\n  - In VP settings, we gradually add more and more noises to data (data --> Gaussian: $\\beta$ increases). However, the sampling process starts with simple Gaussians and ends with samples, so should $\\beta$ decrease, when Gaussian --> data?\n\n- I am curious if the authors have ever tried the detailed balance objective, due to its superior generalization capability (https://arxiv.org/pdf/2407.03105, different settings though)?\n\n- The results for DIS and DDS over Manywell are still missing compared to the previous version. However, the DIS paper does include the results. Any reasons?\n\n- Missing reference - Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. Both papers share the same goal. It would be better to discuss it in the paper.\n\n\nPotential typos:\n\n- Eq.7: To define the reverse KL, a distribution is missing between $\\int$ and $\\log$? Why does $\\mu_{0}(x_{0})$ appear before $d x_{\\Delta t}$?\n\n- Eq.8: $d x_{0}$ is missing?\n\n- Eq.10 & 11: $p_{target}(x_{1})$ --> $R(x_{1})$?\n\n- Line 181: $Z_{\\theta}$ is missing in the equation?\n\n- Eq.12: $f(x_{m})$ --> $f(x_{m \\Delta t})$, as well as $f(x_{n})$?\n\n- Line 209: $I_{d}$ is missing in $p_{t}(x_{t})$?\n\n- Line 290 & 294: $\\log$ is missing before $\\int$, and $P_{F}(\\tau | x_{1})$ --> $P_{F}(\\tau | x_{0})$?", "limitations": "Yes. Limitations were included.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Sampling from a target distribution can be challenging in high-dimensional spaces, especially when the distribution of interest has many separated modes. This paper explores diffusion models to address this challenge. Unlike existing reverse KLD-based methods, such as PIS and DDS, this paper considers GFlowNet-based training objectives (e.g., trajectory balance, sub-trajectory balance), which enable off-policy training. This means that training trajectories are not necessarily from the current forward process, thus enhancing exploratory capability. Additionally, local search using a replay buffer can further enhance exploration in the target space. In general, the paper is well-written and well-organized.\n\n---\n\n**After rebuttal:** I will increase my score to 7. Typos or incorrect writing should be corrected upon acceptance.\n\n---", "weaknesses": "Please see the below questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments and positive assessment of the paper.\n\n### Evaluation metrics\n\nThank you for your question regarding the different aspects of the methods' behaviours, as measured by various evaluation metrics (here: $\\log \\hat Z^{\\rm RW}$ and $\\log \\hat Z$). We agree with your suggestions about the importance of one of them in measuring mode collapse behaviour and the other with the partition function approximation. \n\nHowever, for measuring both mode collapse and mode coverege, we have also a few other options. If we have access to ground truth samples, sample-based metrics (such as those typically used for assessing generative models trained from data) are available. One of them is $\\mathcal{W}_2^2$, which we already included in Appendix C.1. \n\nOther options are, e.g., EUBO, ESS, MMD, and EMC, which are presented in the very interesting \"Beyond ELBOs\" paper you mentioned (which appeared after the submission deadline, but which we were already aware of). The question of which approach is the best in each metric is already hard. However, we noticed that $\\log\\hat Z^{\\rm RW}$ results are better for the methods using Langevin Parametrization (e.g., together with Local Search), whereas $\\log \\hat Z$ performance is better when the method utilizes Local Search.\n\n### VP setting\n\nGood question and observations! As you noticed, we used the PIS model architectures, but changed the noising (backward) process to add increasing levels of noise and made the initial state Gaussian-distributed and not a point mass at 0 (which can be thought of as a fixed first transition from an abstract initial state to a sample from the Gaussian). The exact procedure is presented in Appendix C.2.\n\n### DB objective\n\nThe DB objective fails to discover multiple modes even on the simplest task (25GMM) and thus performs significantly worse than SubTB, which is in turn worse than TB. We guess that this is due to poor credit assignment over long trajectories, as the learning signal to the policy is mediated by learning of the intermediate flow functions (densities).\n\nThe preprint you mention tests generalization abilities only in a single environment, which (1) is discrete, (2) has much shorter and variable-length trajectories, (3) is simple enough that all learning algorithms converge to the optimum when modes are not masked. We guess that all three of these limit the generalizability of those results to the diffusion case. \n\nHowever, it would indeed be interesting, in future work:\n- to study the regularizing effect of local-in-time objectives like DB in the diffusion sampler setting, as has been done in the standard diffusion setting in [Lai et al., \"FP-Diffusion\", ICML'23];\n- to understand the meaning of these objectives in the continuous-time limit.\n\n### Missing results\n\nWe will include the missing results for DIS and DDS on Manywell. The main reason is that our (and, e.g., PIS's) Manywell setting is slightly different than the one used in the DIS paper. We are currently working on making DDS and DIS work on our Manywell version and plan to share those results before the end of the discussion period.\n\n### Missing reference\nThank you for pointing out this paper. We are aware of it and will happily include it in the revised version of our paper. (Note that that paper appeared only on June 11, after the NeurIPS submission date.)\n\n### Potential typos\n\nThank you for reading carefully! We fixed the small mistakes.\n\n**Thanks you again, and please let us know if we can provide any further information during the discussion period.**", "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:
We thank all the reviewers for the effort they put into reviewing our paper and are grateful for the constructive feedback. We appreciate the reviewers remarking that the paper is well-written and well-organized (FAJh), studies a important problem (5SPJ, 6eYp), and does comprehensive benchmarking to demonstrate the effectiveness of the algorithms (6eYp).

Thank you for noting that the paper well-written and well-organized (FAJh), studies a important problem (5SPJ, 6eYp), and does comprehensive benchmarking (6eYp).

Several reviewers (5SPJ, 6eYp, fepi) raised concerns about the difficulty/dimensionality of the densities studied. We respond to those points now.

First, we note that **our evaluation followed the common procedure in the continuous sampling community** [8,41,77,85,87,etc.], which typically evaluates on the same tasks that we chose, with some variations, or even on a subset of them. For instance, DDS and DIS papers evaluate the methods also only on synthetic data and VAE/NICE (DDS) or sampling one image (DIS). 

Furthermore, **as a secondary contribution, we identified discrepancies and reproducibility issues in past work** on these benchmarks (see Appendix B.1). We believe that this is important for sound progress in this field in the future.

The dimensions of the tasks studied are: 2, 10, 20 (with 784-dim conditioning vector), 32, and 1600, which is indeed less than the dimension on which diffusion models are typically trained from data, e.g., images. However, sampling from densities in high dimension **without any prior knowledge** (such as a pretrained model or known samples from the target) is a harder problem than maximizing [a variational bound on] the likelihood of samples, which is perhaps why it has not been the subject of much study in this past work. 

Our methods, which use off-policy training, can be combined with prior knowledge, such as known target samples (such as by placing those samples in the replay buffer), as has already been demonstrated, e.g., in [42]. Simulation-based methods like PIS, DDS, etc. do not have this ability. The use of diffusion samplers in settings with prior knowledge, while not the subject of this paper, is indeed a very interesting direction of research and likely important in more difficult applications.

Meta-review: This paper studies diffusion-structured samplers for amortized inference over continuous variables. The authors present a detailed survey of existing methods for diffusion-structured sampling, and draw connections between several methods in the literature, placing them in a common framework.

The paper makes two main contributions: 1) it provides a software library unifying existing methods and performs a comprehensive benchmark of different diffusion-structured samplers, investigating unconditional and conditional generative modeling tasks: sampling from various energy distributions with dimensionalities ranging from 2 to 1600, and conditional sampling from the latent posterior of a VAE; and 2) they propose a new exploration strategy for off-policy methods that uses local search and a replay buffer.

The authors show that their proposed exploration method helps prevent mode collapse and achieves competitive or better performance than other methods on several unconditional and conditional modeling tasks. Their benchmark results also identify reproducibility issues with previous work, which is useful for the field.

Reviewers found that the paper is well-written, addresses an interesting and important problem, and provides comprehensive benchmark results. Several noted that the provided software library unifies existing methods and will be valuable to the community.

A few reviewers mentioned that the problems in the experiments are small-scale; however, the authors addressed this, noting that their evaluation setting is in line with previous work.

Overall, this paper provides valuable contributions that will be of interest to the NeurIPS community.
Decision: Accept (poster)