Improved off-policy training of diffusion samplers

Authors: marcin  sendera
minsu  kim
yoshua  bengio
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 11:15:06

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.

Main paper: 
 introduction: 
 Approximating and sampling from complex multivariate distributions is a fundamental problem in probabilistic deep learning [e.g., 26, 34, 25, 47, 56] and in scientific applications [3, 51, 37, 1, 31] . The problem of drawing samples from a distribution given only an unnormalized probability density or 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 as Metropolis-adjusted Langevin [MALA; 23, 64, 63] and Hamiltonian MC [HMC; 19, 30] may be slow to mix between modes and have a high cost per sample. While variants such as sequential MC [SMC; 24, 12, 15] and nested sampling [68, 9, 42] 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, continuous-time stochastic processes that gradually evolve a simple distribution to a complex target, are powerful density estimators with proven mode-mixing properties [14] ; as such, they have been widely used in the setting of generative models learned from data [69, 71, 27, 49, 65] . However, the problem of training diffusion models to sample from a distribution with a given blackbox density or energy function has attracted less attention. Recent work has drawn connections between diffusion (learning the denoising process) and stochastic control (learning the Föllmer drift [20] ), leading to approaches such as the path integral sampler [PIS; 87], denoising diffusion sampler [DDS; 77], and time-reversed diffusion sampler [DIS; 8]; such approaches were recently unified by [62] and [78] . Another line of work [41, 85] 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 [45] .Despite the advances in sampling methods and attempts to unify them theoretically [62, 78] , 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 in their definitions of the same target densities (see §B.1). The first main contribution of this paper is a unified library for diffusion-structured samplers. The library has a focus on off-policy methods (continuous GFlowNets) but also includes simulation-based 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 a study of methods for improving exploration and credit assignmentthe propagation of learning signals from the target density to the parameters of earlier sampling steps -in diffusion-structured samplers ( §4). First, our results ( §5.2) suggest that the technique of utilizing partial trajectory information [43, 54] , as done in the diffusion setting by [85] , offers little benefit, and a higher training cost, over on-policy [87] or off-policy [41] trajectory-based optimization. Second, we examine the utility of a gradient-based variant which parametrizes the denoising distribution as a correction to a Langevin process [87] . We show that this inductive bias is also beneficial in the offpolicy (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. 
prior work: 
 Amortized variational inference approaches use a parametric model 𝑞 𝜃 to approximate a given target density 𝑝 target , typically through stochastic optimization [29, 57, 2] . Notably, explicit density models like autoregressive models and normalizing flows have been extensively utilized in density estimation [59, 18, 80, 21, 50] . However, these models impose structural constraints, thereby limiting their expressive power [13, 22, 86] . The adoption of diffusion processes in generative models has stimulated a renewed interest in hierarchical models as density estimators [79, 27, 75] . Approaches like PIS [87] leverage stochastic optimal control for sampling from unnormalized densities, albeit still struggling with scalability in high-dimensional 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 [45, 89, 72, 17] , expanding from specific scientific domains [e.g., 35, 4, 88] to amortized inference over a broader array of latent structures [e.g., 76, 33] . 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 [44, 43, 54, 58, 67, 38, 36] . These challenges have repercussions in the scalability of these methods in more complex scenarios, which this paper addresses in the continuous case.3 Setting: Diffusion-structured sampling Let E : R 𝑑 → R be a differentiable energy function and define 𝑅(x) = exp(-E (x)), the reward or unnormalized target density. Assuming the integral 𝑍 := ∫ R 𝑑 𝑅(x) 𝑑x exists, E defines a Boltzmann density 𝑝 target (x) = 𝑅(x)/𝑍 on R 𝑑 . We are interested in the problems of sampling from 𝑝 target and approximating the partition function 𝑍 given access only to E and possibly to its gradient ∇E.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). 
euler-maruyama hierarchical samplers: 
 Generative modeling with SDEs. Diffusion models assume a continuous-time generative process given by a neural stochastic differential equation [SDE; 74, 53, 66] :EQUATIONwhere x 0 follows a fixed tractable distribution 𝜇 0 (such as a Gaussian or a point mass). The initial distribution 𝜇 0 and the stochastic dynamics specified by (1) induce marginal densities 𝑝 𝑡 on R 𝑑 for each 𝑡 > 0. The functions 𝑢 and 𝑔 have learnable parameters that we wish to optimize, using some objective, so as to make the terminal density 𝑝 1 close to 𝑝 target . Samples can be drawn from 𝑝 1 by sampling x 0 ∼ 𝜇 0 and simulating the SDE (1) to time 𝑡 = 1.The SDE driving 𝜇 0 to 𝑝 target is not unique. However, if one fixes a reverse-time SDE, or noising process, that pushes 𝑝 target at 𝑡 = 1 to 𝜇 0 at 𝑡 = 0, 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 𝑢 of the denoising process given samples from 𝑝 target , and the diffusion rate 𝑔 is available in closed form [27, 71] .Time discretization. In practice, the integration of the SDE ( 1) is approximated by a discrete-time scheme, the simplest of which is Euler-Maruyama integration. The process ( 1) is replaced by a discrete-time Markov chainx 0 → x Δ𝑡 → x 2Δ𝑡 → • • • → x 1, where Δ𝑡 =1 𝑇 is the time increment and and 𝑇 is the number of steps:EQUATIONThe density of the transition kernel from x 𝑡 to x 𝑡+Δ𝑡 can explicitly be written asEQUATIONwhere 𝑝 𝐹 denotes the transition density of the discretized forward SDE. This density defines a joint distribution over trajectories starting at x 0 :EQUATIONSimilarly, a discrete-time reverse processEQUATIONIf the forward and backward processes (starting from 𝜇 0 and 𝑝 target , respectively) are reverses of each other, then they define the same distribution over trajectories, i.e., for allEQUATIONIn particular, the marginal densities of x 1 under the forward and backward processes are then equal to 𝑝 target , and the forward process can be used to sample the target distribution.Because the reverse of a process with Gaussian increments is, in general, not itself Gaussian, (6) can be enforced only approximately, but the discrepancy vanishes as Δ𝑡 → 0 (i.e., increments are infinitesimally Gaussian), an application of the central limit theorem that is key to stochastic calculus [53] .SDE learning as hierarchical variational inference. The problem of learning the parameters 𝜃 of the forward process so as to enforce ( 6) is one of hierarchical variational inference. The backward process transforms x 1 into x 0 via a sequence of latent variables x 1-Δ𝑡 , . . . , x 0 , and the forward process aims to 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 𝑝 target , one can optimize the forward process by minimizing the KL divergence 𝐷 KL ( 𝑝 target • 𝑝 𝐵 ∥𝜇 0 • 𝑝 𝐹 ) between the distribution over trajectories given by the reverse process and that given by the forward process. This is equivalent to the typical training of diffusion models, which optimizes a variational bound on the data log-likelihood (see [70] ). However, in the setting of an intractable density 𝑝 target , unbiased estimators of this divergence are not available. Instead, one can optimize the reverse KL:2 EQUATIONVarious estimators of this objective are available. For instance, the path integral sampler objective [PIS; 87] uses the reparametrization trick to express (7) as an expectation over noise variables z 𝑡 that participate in the hierarchical sampling of x Δ𝑡 , . . . , x 1 , yielding an unbiased gradient estimator, but one that requires backpropagation into the simulation of the forward process. The related denoising diffusion sampler [DDS; 77] applies the same principle in a different integration scheme.Continuous generative flow networks (GFlowNets) [41] express the problem of enforcing (6) as a 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 [45] in the distribution-matching setting and by [83] in the maximum-likelihood setting, while the theoretical foundations for continuous GFlowNets were later laid down by [41] .State and action space. To formulate sampling as a sequential decision-making problem, one must define the spaces of states and actions. In the case of sampling by 𝑇-step Euler-Maruyama integration, assuming 𝜇 0 is a point mass at 0, the state space isS = {(0, 0) ∪ (x, 𝑡) : x ∈ R 𝑑 , 𝑡 ∈ {Δ𝑡, 2Δ𝑡, . . . , 1} ,with the point (x, 𝑡) representing that the sampling agent is at position x at time 𝑡.Sampling begins with the initial state x 0 := (0, 0), proceeds through a sequence of states (x Δ𝑡 , Δ𝑡), (x 2Δ𝑡 , 2Δ𝑡), . . . , and ends at a state (x 1 , 1); states (x, 𝑡) with 𝑡 = 1 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 𝑡 is clear from context. The sequence of statesx 0 → x Δ𝑡 → • • • → x 1 is called a complete trajectory.The actions from a nonterminal state (x 𝑡 , 𝑡) correspond to the possible next states (x 𝑡+Δ𝑡 , 𝑡 + Δ𝑡) that can be reached from (x 𝑡 , 𝑡) by a single step of the Euler-Maruyama integrator. 3Forward 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, 𝑡). In our context, this amounts to a collection of conditional probability densities 𝑝 𝐹 (x 𝑡+Δ𝑡 | x 𝑡 ; 𝜃), representing the density of the transition kernel from x 𝑡 to x 𝑡+Δ𝑡 . GFlowNet training optimizes the parameters 𝜃, which may be the weights of a neural network specifying a density over x 𝑡+Δ𝑡 conditioned on x Δ𝑡 . A policy 𝑝 𝐹 induces a distribution over complete trajectories𝜏 = (x 0 → x Δ𝑡 → • • • → x 1 ) via 𝑝 𝐹 (𝜏; 𝜃) = 𝑇 -1 𝑖=0 𝑝 𝐹 (x (𝑖+1)Δ𝑡 | x 𝑖Δ𝑡 ; 𝜃).In particular, we get a marginal density over terminal states:EQUATIONThe learning problem solved by GFlowNets is to find the parameters 𝜃 of a policy 𝑝 𝐹 whose terminating density 𝑝 ⊤ 𝐹 is equal to 𝑝 target , i.e.,EQUATIONHowever, because the integral ( 8) is intractable and 𝑍 is unknown, auxiliary objects must be introduced into optimization objectives to enforce (9) , as discussed below.Notably, if the policy is a Gaussian with mean and variance given by neural networks taking x 𝑡 and 𝑡 as input, then learning the policy amounts to learning the drift 𝑢(x 𝑡 , 𝑡; 𝜃) and diffusion 𝑔(x 𝑡 , 𝑡; 𝜃) of a SDE (1), i.e., fitting a neural SDE. The SDE learning problem in §3.1 is thus the same as that of fitting a GFlowNet with Gaussian policies.Backward policy and trajectory balance. A backward policy is a collection of conditional probability densities 𝑝 𝐵 (x 𝑡 -Δ𝑡 | x 𝑡 ; 𝜓), representing a probability density of transitioning from x 𝑡 to an ancestor state x 𝑡 -Δ𝑡 . The backward policy induces a distribution over complete trajectories 𝜏 conditioned on their terminal state (cf. ( 5)):𝑝 𝐵 (𝜏 | x 1 ; 𝜓) = 𝑇 𝑖=1 𝑝 𝐵 (x (𝑖-1)Δ𝑡 | x 𝑖Δ𝑡 ; 𝜓),where exceptionally 𝑝 𝐵 (x 0 | x Δ𝑡 ) = 1 as 𝜇 0 is a point mass.Generalizing a result in the discrete-space setting [44] , [41] show that 𝑝 𝐹 samples from the target distribution (i.e., satisfies ( 9)) if and only if there exists a backward policy 𝑝 𝐵 and a scalar 𝑍 𝜃 such that the trajectory balance conditions are fulfilled for every complete trajectory 𝜏 = (x 0 → x Δ𝑡 → • • • → x 1 ):EQUATION) If these conditions hold, then 𝑍 𝜃 equals the true partition function 𝑍 = ∫ x 𝑅(x) 𝑑x. The trajectory balance objective for a trajectory 𝜏 is the squared log-ratio of the two sides of (10) , that is:EQUATIONOne can thus achieve (9) by minimizing to zero the loss L TB (𝜏; 𝜃, 𝜓) with respect to the parameters 𝜃 and 𝜓, where the trajectories 𝜏 used for training are sampled from some training policy 𝜋(𝜏).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 𝑝 𝐹 and the estimate of the partition function 𝑍 𝜃 . For example, for most experiments in §5, we fix the backward policy to a 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-policy, i.e., using trajectories sampled from an arbitrary distribution 𝜋. Because minimizing L TB (𝜏; 𝜃, 𝜓) to 0 for all 𝜏 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 [16] , Thompson sampling [58] , and backward traces from terminal states obtained by MCMC [42] .In the continuous case, [45, 41] proposed to simply add a small constant to the policy variance when sampling trajectories for training. Off-policy optimization is a key advantage of GFlowNets over variational methods such as PIS, which require on-policy optimization [45] .However, when L TB happens to be optimized on-policy, i.e., using trajectories sampled from the policy 𝑝 𝐹 itself, we get an unbiased estimator of the gradient of the KL divergence (7) with respect to 𝑝 𝐹 's parameters up to a constant [61, 45, 89] , that is:E 𝜏∼ 𝑝 𝐹 ( 𝜏 ) [∇ 𝜃 ′ L TB (𝜏; 𝜃, 𝜓)] = 2∇ 𝜃 ′ 𝐷 KL ( 𝑝 𝐹 (𝜏; 𝜃) ∥ 𝑝 target (x 1 ) 𝑝 𝐵 (𝜏 | x 1 ; 𝜓)), where ∇ 𝜃 ′ denotes the gradient with respect to the parameters of 𝑝 𝐹 , but not 𝑍 𝜃 . This unbiased estimator tends to have higher variance than the reparametrization-based 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; 43], which involves modeling a scalar state flow 𝑓 (x 𝑡 ; 𝜃) 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 trajectoriesEQUATIONwhere for terminal states 𝑓 (x 1 ) = 𝑅(x 1 ). This approach has some computational overhead associated with training the state flow, but has been shown to be effective in discrete-space settings, especially when combined with the forward-looking reward shaping scheme proposed by [54] . 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 [52, 61] , rediscovered for GFlowNets by [84] .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 [45] , this is equivalent to minimizing L TB first with respect to log 𝑍 𝜃 to optimality over the batch, then with respect to the parameters of 𝑝 𝐹 .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA]• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It 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 a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
exploration and credit assignment in continuous gflownets: 
 The main challenges in training off-policy sampling models are 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-structured GFlowNets. These techniques will be empirically studied and compared in §5.Partial energies and subtrajectory-based learning. [85] studied the diffusion sampler learning problem introduced by [41] , but replaced the TB learning objective with the SubTB objective. 4 In addition, an inductive bias resembling the geometric interpolation in [46] was used for the state flow function:EQUATIONwhere NN is a neural network and𝑝 ref 𝑡 (x 𝑡 ) = N (x 𝑡 ; 0, 𝜎 2 𝑡𝐼 𝑑 )is the marginal density of a Brownian motion with rate 𝜎 at x 𝑡 . The use of the target density log 𝑅(x 𝑡 ) = -E (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 on the state flow was called forward-looking (FL) by [54] , and we will refer to this method as FL-SubTB in §5.Langevin dynamics inductive bias. [87] proposed an inductive bias on the architecture of the drift of the neural SDE 𝑢(x 𝑡 , 𝑡; 𝜃) (in GFlowNet terms, the mean of the Gaussian density 𝑝 𝐹 (x 𝑡+Δ𝑡 | x 𝑡 ; 𝜃)) that resembles a Langevin process on the target distribution. One writesEQUATIONwhere NN 1 and NN 2 are neural networks outputting a 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 is 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 on policies represents a different way of incorporating the reward signal at intermediate steps in the trajectory and can steer the sampler towards low-energy regions. It contrasts 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 1 independent of x 𝑡 is equivalent to the forward-looking flow (13) in the limit of small time increments Δ𝑡 → 0, i.e., they induce the same asymptotics of the discrepancy in the SubTB constraints (12) over short partial trajectories. Such theoretical analysis can be the subject of future work.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 §C.3). We present an alternative technique that does not induce additional computation cost per training trajectory. LGCP (𝑑 = 1600) To enhance the quality of samples during training, we incorporate local search into the exploration process, motivated by the success of local exploration [82, 32, 39] and replay buffer [e.g., 16] methods for GFlowNets in discrete spaces. Unlike these methods, which define MCMC kernels via the GFlowNet policies, our method leverages parallel Metropolis-adjusted Langevin (MALA) directly in the target space.Algorithm ↓ Metric → Δ log 𝑍 Δ log 𝑍 RW Δ log 𝑍 Δ log 𝑍 RW Δ log 𝑍 Δ logIn detail, we initially sample 𝑀 candidates from the sampler: {x (1) , . . . , x ( 𝑀 ) } ∼ 𝑝 ⊤ 𝐹 (•). Subsequently, we run parallel MALA across 𝑀 chains over 𝐾 transitions , with the initial states of the Markov chain being {x (1) , . . . , x ( 𝑀 ) }. After the 𝐾 burn-in burn-in transitions, the accepted samples are stored in a local search buffer D LS . We occasionally update the buffer using MALA steps and replay samples from it to minimize the computational demands of iterative local search. MALA steps are far more parallelizable than sampler training and need to be made only rarely (as the buffer is much larger than the training batch size), so 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 D LS (uniformly or using a prioritization scheme, §E), sample a trajectory 𝜏 leading to x from the backward process, and make a gradient update on the objective (e.g., TB) associated with 𝜏.When training with local search guidance, we alternate two steps, inspired by [42] , 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 search-guided GFlowNet training. 
experiments: 
 We conduct comprehensive benchmarks of various diffusion-structured samplers, encompassing both GFlowNet samplers and methods such as PIS. For the GFlowNet samplers, we investigate a 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.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 32-dimensional Manywell distribution, and the 1600-dimensional Log-Gaussian Cox processand conditional sampling from the latent posterior of a variational autoencoder (VAE; [40, 60] ). This allows us to investigate both unconditional and conditional generative modeling techniques.We evaluate three algorithm categories:(1) Traditional sampling methods: We consider a standard Sequential Monte Carlo (SMC) implementation and a state-of-the-art nested sampling method (GGNS, [42] ). ( 2) Simulation-driven variational approaches: DIS [8] , DDS [77] , and PIS [87] .(3) Diffusion-based GFlowNet samplers: Our evaluation focuses on TB-based training and the enhancements described in §4: the VarGrad estimator (VarGrad), off-policy exploration (Expl.), Langevin parametrization (LP), and local search (LS). Additionally, we assess the FL-SubTBbased continuous GFlowNet as studied by [85] for a 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 [87] . We aim to learn a Gaussian forward policy 𝑝 𝐹 that samples from the target distribution in 𝑇 = 100 steps (Δ𝑡 = 0.01). Just as in past work [87, 41, 85] , the backward process is fixed to a discretized Brownian bridge with a noise rate 𝜎 that depends on the domain; explicitly,EQUATIONunderstood to be a point mass at 0 when 𝑡 = Δ𝑡. To keep the learning problem consistent with past work, we fix the variance of the forward policy 𝑝 𝐹 to 𝜎 2 . 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 metrics. To evaluate diffusion-based samplers, we use two metrics from past work [87, 41] , which we restate in our notation. Given any forward policy 𝑝 𝐹 , we have a variational lower bound on the log-partition function log𝑍 = ∫ R 𝑑 𝑅(x) 𝑑x: log ∫ R 𝑑 𝑅(x) 𝑑x = log E 𝜏=(•••→x 1 )∼ 𝑝 𝐹 ( 𝜏 ) 𝑅(x 1 ) 𝑝 𝐵 (𝜏 | x 1 ) 𝑝 𝐹 (𝜏) ≥ E 𝜏=(•••→x 1 )∼ 𝑝 𝐹 ( 𝜏 ) log 𝑅(x 1 ) 𝑝 𝐵 (𝜏 | x 1 ) 𝑝 𝐹 (𝜏).We use a 𝐾-sample (𝐾 = 2000) Monte Carlo estimate of this expectation, log Ẑ, as a metric, which equals the true log 𝑍 if 𝑝 𝐹 and 𝑝 𝐵 jointly satisfy (10) and thus 𝑝 𝐹 samples from the target distribution. We also employ an importance-weighted variant, which emphasizes mode coverage over accurate local modeling:log ẐRW := log 𝐾 ∑︁ 𝑖=1 𝑅(x (𝑖) 1 ) 𝑝 𝐵 (𝜏 (𝑖) | x (𝑖) 1 ) 𝑝 𝐹 (𝜏 (𝑖) ), where 𝜏 (1) , . . . , 𝜏 (𝐾 ) are trajectories sampled from 𝑝 𝐹 and leading to terminal states x (1) 1 , . . . , x (𝐾 ) 1 . The estimator log ẐRW is also a lower bound on log 𝑍 and approaches it as 𝐾 → ∞ [10] . In the unconditional modeling benchmarks, we compare both estimators to the true log-partition function, which is known analytically for all tasks except LGCP (leading to discrepancies in past work; see §B.1).In addition, we include a sample-based metric (2-Wasserstein distance); see §C.1. Unconditional sampling. We report the metrics for all algorithms and energies in Table 1 .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 hypotheses from past work about the importance of offpolicy exploration [45, 41] and the importance of improved credit assignment [85] . On the other hand, our results do not show a consistent and significant improvement of the FL-SubTB objective used by [85] 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 performance on 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 a replay buffer (LS) leads to a substantial improvement in performance, surpassing or competing with GFlowNet baselines, non-GFlowNet baselines, and non-amortized sampling methods in most tasks and metrics. 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-3× slower per iteration), indicating that previous observations [87] transfer to off-policy algorithms. Compared to FL-SubTB, 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 or competitive across most tasks and metrics.In §C.3, we study the scaling of the algorithms with dimension, showing efficiency of the proposed LS.Conditional sampling. For the VAE task, we observe that the performance of the baseline GFlowNet-based samplers is generally worse 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 is 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.1 we show decoded samples encoded using the best-performing diffusion encoder.)Our implementation of diffusion-structured generative flow networks includes several additional options that diverge from the modeling assumptions made in most past work in the field. Notably, it features the ability to:• optimize the backward (noising) process -not only the denoising process -as was done for related learning problems in [11, 62, 78] ; • learn the forward process's diffusion rate 𝑔(x 𝑡 , 𝑡; 𝜃), not only the mean 𝑢(x 𝑡 , 𝑡; 𝜃); • assume a 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 a GFlowNet objective.As noted in §5.1, in the main experiments we fixed the diffusion rate of the learned forward process, an assumption inherited from all past work and justified in the continuous-time 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 [42] . We see that when the forward policy'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 §C.2, we include preliminary results using a variance-preserving backward process, as commonly used in diffusion models, in place of the reversed Brownian motion used in the main experiments.The ability to model distributions accurately in fewer steps is important for computational efficiency. Future work can consider ways to improve performance in coarse time discretizations, such as non-Gaussian transitions, whose utility in diffusion models trained from data has been demonstrated [81] . 
conclusion: 
 We have presented a study of diffusion-structured 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 expectation-maximization loop for training latent variable models, as was recently done for discrete compositional latents by [32] , and for sampling Bayesian posteriors over high-dimensional model parameters. The most important direction of theoretical work is understanding the continuous-time limit (𝑇 → ∞) of all the algorithms we have studied. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?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 or 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). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We believe there are no violations of the CoE.Guidelines:• The answer NA means that the authors have not 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., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 a ML problem with no 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 societal impact or 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.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper studies a ML problem with no immediate application to generation of 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: 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?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 or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: Code is available at https://github.com/GFNOrg/gfn-diffusion and will continue to be maintained and extended.Below are commands to reproduce some of the results on Manywell and VAE with PIS and GFlowNet models as an example, showing the hyperparameters: PIS:--mode_fwd pis --lr_policy 1e-3--mode_fwd pis --lr_policy 1e-3 --langevin Gaussian Mixture Model with 25 modes (25GMM). The model, termed as 25GMM, consists of a two-dimensional 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 a grid formed by the Cartesian product {-10, -5, 0, 5, 10} × {-10, -5, 0, 5, 10}, effectively distributing them across the coordinate space.Funnel [28] . The funnel represents a classical benchmark in sampling techniques, characterized by a ten-dimensional distribution defined as follows: The first dimension, 𝑥 0 , follows a normal distribution with mean 0 and variance 9, denoted as 𝑥 0 ∼ N (0, 9). Conditional on 𝑥 0 , the remaining dimensions, 𝑥 1:9 , are distributed according to a multivariate normal distribution with mean vector 0 and a covariance matrix exp(𝑥 0 )I, where I is the identity matrix. This is succinctly represented as𝑥 1:9 | 𝑥 0 ∼ N (0, exp (𝑥 0 ) I).Manywell [51] . The manywell is characterized by a 32-dimensional distribution, which is constructed as the product of 16 identical two-dimensional double well distributions. Each of these two-dimensional components is defined by a potential function, 𝜇(𝑥 1 , 𝑥 2 ), expressed as𝜇(𝑥 1 , 𝑥 2 ) = exp -𝑥 4 1 + 6𝑥 2 1 + 0.5𝑥 1 -0.5𝑥 2 2. VAE [40] . This task involves sampling from a 20-dimensional latent posterior 𝑝(𝑧|𝑥) ∝ 𝑝(𝑧) 𝑝(𝑥|𝑧), where 𝑝(𝑧) is a fixed prior and 𝑝(𝑥|𝑧) is a pretrained VAE decoder, using a conditional sampler 𝑞(𝑧|𝑥) dependent on input data (image) 𝑥.LGCP [48] . This density over a 1600-dimensional variable is a Log-Gaussian Cox process fit to a distribution of pine saplings in Finland.Wrong definitions of the Funnel density. As already noted by [77] , [87] 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 [41] .LGCP. The LGCP benchmark suffers from the lack of a consistent ground truth log 𝑍 to compare against. Previous work has compared the value of the partition function log 𝑍 against a "long run of Sequential Monte Carlo" [87] . We note that this approach produces noisy estimates of the partition function, especially in high-dimensional problems (indeed, SMC has rarely been used in problems with over a 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 𝑍 = 512.6 in one repository and log 𝑍 = 501.8 in another.On FL-SubTB as used in [85] . We make two observations calling into question the main results of [85] .First, the only substantial difference between the algorithm used by [85] and the one from the past work [41] -which first proposed the use of GFlowNet objectives to train diffusion samplers -is the substitution of the FL-SubTB objective [54, 43] for TB [44] . However, [85] elects to compare FL-SubTB with the Langevin parameterization to TB without the Langevin parameterization. Our results in Table 1 show that while the Langevin parameterization is 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 [85] 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 [85] , who confirmed that running their published code does not reproduce the results in the paper but could not provide any further explanation or a working implementation. with 𝜈 ∈ (0, ∞).In particular, we follow the common procedure -use 𝜈 := 1 and 𝛽(𝑡) := (1 -𝑡) 𝛽 𝑚𝑖𝑛 + 𝑡 𝛽 𝑚𝑎𝑥 , 𝑡 ∈ [0, 1], with 𝛽 𝑚𝑖𝑛 = 0.01 and 𝛽 𝑚𝑎𝑥 = 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.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 as for 𝑑 = 32. Due to the large runtime, some runs in dimensions 128 and 512 had to be limited at 12 hours, while in 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 is critical to performance. • The even higher overhead of FL-SubTB as used by [85] .• The relatively high efficiency and low overhead of our newly proposed local search. 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 𝜎 and also used in the definition of the fixed backward process. The base diffusion rate 𝜎 2 (parameter t_scale) is set to 5 for 25GMM and 1 for Funnel and Manywell, consistent with past work.For LGCP, we found that using too small diffusion rate 𝜎 2 (e.g., 𝜎 2 = 1) prevents the methods from achieving reasonable results. We tested different values of 𝜎 2 = {1, 3, 5}, and selected 𝜎 2 = 5, which gives the best results, which follows the findings in Zhang & Chen [87] .For all our experiments, we used a learning rate of 10 -3 . Additionally, we used a higher learning rate for learning the flow parameterization, which is set as 10 -1 when using the TB loss and 10 -2 with the SubTB loss. These settings were found to be consistently stable (unlike those with higher learning rates) and converge within the allotted number of steps (unlike those with lower learning rates).For the SubTB loss, we experimented with the settings of 10× lower learning rates for both flow and policy models communicated by the authors of [85] , 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 to the policy when sampling trajectories for training), which decays linearly over the first half of training, consistent with [41] .We train all our models for 25, 000 iterations except those using Langevin dynamics, which are trained for 10, 000 iterations. This results in 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 [87] . Similar to [87] , we also use an initialization scheme with last-layer weights set to 0 at the start of training. Since the SubTB requires the flow function to be conditioned on the current state x 𝑡 and time 𝑡, we follow [85] and parametrize the flow model with the same architecture as the Langevin scaling model NN 2 in [87] . 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 10 2 and policy network to 10 4 , similarly to [77] :EQUATION)All models were trained with a batch size of 300. In each experiment, we train models on a single NVIDIA A100-Large GPU, if not stated explicitly otherwise.VAE experiment. In the VAE experiment, we used a standard VAE model pretrained for 100 epochs on the MNIST dataset. The encoder 𝑞(𝑧|𝑥) 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 be 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-dimensional output. The output is processed by the sigmoid function to be scaled properly into [0, 1].The goal is to sample conditionally on 𝑥 the latent 𝑧 from the unnormalized density 𝑝(𝑧, 𝑥) = 𝑝(𝑧) 𝑝(𝑥 | 𝑧) (where 𝑝(𝑧) is the prior and 𝑝(𝑥|𝑧) is the likelihood computed from the decoder), which is proportional to the posterior 𝑝(𝑧 | 𝑥). We reuse the model architectures from the unconditional sampling experiments, but also provide 𝑥 as an input to the first layer of the models expressing the policy drift (as well as the flow, for FL-SubTB) and add one hidden layer to process high-dimensional conditions. For models trained with TB, log 𝑍 𝜃 also becomes a MLP taking 𝑥 as input.The VarGrad and LS techniques require adaptations in the conditional setting. For LS, buffers (D buffer and D LS ) must store the associated conditions 𝑥 together with the samples 𝑧 and the corresponding unnormalized density 𝑅(𝑧; 𝑥), i.e., a tuple of (𝑥, 𝑧, 𝑅(𝑧; 𝑥)). For VarGrad, because the partition function depends on the conditioning information 𝑥, it is necessary to compute variance over many trajectories sharing the same condition. We choose to sample 10 trajectories for each condition occurring in a minibatch and compute the VarGrad loss for each such set of 10 trajectories. 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 our experiments (see ablation study in §E.2), although the choice should be investigated more thoroughly in future work.We use rank-based prioritization [73] , which follows a probabilistic approach defined as:EQUATIONwhere rank D buffer (x) represents the relative rank of a sample 𝑥 based on a ranking function 𝑅(x) (in our case, the unnormalized target density at sample x). The parameter 𝑘 is a hyperparameter for prioritization, where a lower value of 𝑘 assigns a higher probability to samples with higher ranks, thereby introducing a more greedy selection approach. We set 𝑘 = 0.01 for every task. Given that the sampling is proportional to the size of D buffer , we impose a constraint on the maximum size of the buffer: We use the number of total iterations 𝐼 = 25, 000 for every task as default. Note as local search is performed to update D LS occasionally that per 100 iterations, the number of local search updates is done 25, 000/100 = 250.|D buffer | =Increasing capacity of buffer. The capacity of the replay buffer influences the duration for which it retains past experiences, enabling it to replay these experiences to the policy. This mechanism 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 experiences-enhances its mode-seeking capability. 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 are encouraged to create a 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 well-specification, 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 on a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.38th Conference on Neural Information Processing Systems (NeurIPS 2024).In the case that 𝜇 0 is a point mass, we assume the distribution x 0 | x Δ𝑡 to also be a point mass, which has density 𝑝 𝐵 (x 0 | x Δ𝑡 ) = 1 with respect to the measure 𝜇 0 .To be precise, the fraction in(7) should be understood as a Radon-Nikodym derivative, which makes sense whether 𝜇 0 is a point mass or a continuous distribution and generalizes to continuous time[8,62].Formally, the foundations in[41] 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 on other spaces, such as any Riemannian manifolds, where other methods do not directly apply.Despite the claimed benefits of FL-SubTB for diffusion samplers, we discovered that[85] modifies critical experimental variables in comparisons and reports irreproducible results; see §B.1. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- The experiments are on synthetic energy functions and MNIST VAE. Including more real-world data or models would be informative.
questions: 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: limitations: Adequately addressed
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.
questions: questions: Is there any benchmark in this field involving a real application and complex high-dimensional distribution?
limitations: limitations: See my comments on the weakness above.
paper_summary: paper_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.
strength_weakness: strength_weakness: 1.	Although 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.
2.	The 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.
3.	In 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. 
4.	Besides, 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?
questions: questions: 1.	The 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.
2.	Could 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: limitations: See Weakness 3.
paper_summary: paper_summary: The paper studies the problem of training diffusion models to sample from a target distribution. The contributions are summarized as follows: 

1. A codebase is provided for the study of diffusion-based samplers, due to the issue of inconsistent experimental settings in previous research; 

2. Exploration in the target space can be enhanced by GFlowNet-based off-policy training objectives and local search with the use of replay buffer. 

3. Experimental results validate the effectiveness of the proposed approach.
strength_weakness: strength_weakness: Please see the below questions.
questions: 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$.

- In terms of C2 variance-preserving noising process, 

  - 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? 

  - 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?

- 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)?

- 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?

- 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.


Potential typos:

- 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}$?

- Eq.8: $d x_{0}$ is missing?

- Eq.10 & 11: $p_{target}(x_{1})$ --> $R(x_{1})$?

- Line 181: $Z_{\theta}$ is missing in the equation?

- Eq.12: $f(x_{m})$ --> $f(x_{m \Delta t})$, as well as $f(x_{n})$?

- Line 209: $I_{d}$ is missing in $p_{t}(x_{t})$?

- Line 290 & 294: $\log$ is missing before $\int$, and $P_{F}(\tau | x_{1})$ --> $P_{F}(\tau | x_{0})$?
limitations: limitations: Yes. Limitations were included.

Comments: title: Additional DDS+DIS results
comment: Thank you for the rebuttal. After considering it, I believe that this paper does present a valuable contribution through its thorough comparisons and the generally-good performance of the proposed method. I still believe that it would be strengthened by experiments on more complex/high-dimensional problems but given that, as the author's point out in the rebuttal, the evaluations tasks are standard in the continuous sampling community, I believe they are sufficient for acceptance. I have raised my score to 6.
Dear Reviewer 6eYp,

This message is to ask if you have any final questions and if our response has affected your assessment of the paper.

Thank you,

The authors.
Dear Reviewer 5SPJ,

This message is to ask if you have any final questions and if our response has affected your assessment of the paper.

Thank you,

The authors.
Much appreciated, and thank you yet again for your detailed reading and finding all the small mistakes that we missed.
Thank you for taking the time to clarify. I am willing to increase my score to 7. Typos or incorrect writing should be corrected upon acceptance.
That task (25GMM) is easy enough that none of the methods in the last group (using either LS or LP) are exhibiting mode collapse -- they sample points around each of the 25 modes.

Note that:
- a sampler that finds 9 of 25 modes (the inner 3-by-3 square) will have a $\log Z^{\rm RW}$ estimation error of about $-\log\frac{9}{25}\approx0.44$ (as TB and VarGrad without any extensions do);
- a sampler that finds only one mode will have an error of about $-\log\frac{1}{25}\approx1.40$, and two modes $-\log\frac{2}{25}\approx1.10$ (and the methods with error >1 do indeed usually find one or two modes).
**On the evaluation metrics:** for example, TB+Expl: 4.01 vs. TB + Expl. + LS (ours): 4.57, together with Figure 1, where you mentioned that our proposed replay buffer with local search is capable of preventing mode collapse. I might conclude that TB+Expl suffers from more serious mode collapse, with giving a lower value?

Other points are clear to me now. Thank you.
Thank you for following up.

**On the evaluation metrics:** we are not sure to understand your question. In fact, "TB + Expl. + LS (ours)" shows **less** serious mode collapse and correspondingly lower $\log Z$ estimation error.

**On the VP setting:**

> Do you mean that the policy at the first step ... is constrained to be a unit Gaussian?

Yes, precisely as you wrote (assuming integer indexing of time steps).

> However, the sampling process starts with simple Gaussians and ends with samples, so shouldn't $\beta$ decrease?

We have mistakenly written the noise schedule in equation (16) **in reverse** (i.e., the time indexing convention for diffusion models trained from data, not the one for diffusion samplers). In fact, $\beta$ decreases linearly from 4 to 0.01 (as we have just confirmed in the code that was run for the experiments).

Thank you again for reading so carefully.
Thank you for your detailed responses, and I apologize for the delayed reply.

### Evaluation metrics

- I agree that if ground-truth samples are available, then we can use sample-based metrics for evaluation. 

- I am still a bit confused. For example, if we look at TB + Expl. + LS (ours), we see that $\Delta \log Z = 0.171$ and $\Delta \log Z^{\mathrm{RW}} = 0.004$, which are the best in the column. We might conclude that it performs the best? It gives the lowest approximation error for $\log Z$, while suffering from more serious mode collapse?

### VP setting

> Good 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.

- Do you mean that the policy at the first step, from $s_{0} = ((0, 0), 0)$ to (x_{1}, 1) is constrained to be a unit Gaussian, i.e., $P_{F}((x_{1}, 1) | ((0, 0), 0)) := P(x_{1})$, while subsequent steps are conditional Gaussians with a known variance, i.e., $P_{F}(\cdot | (x_{t}, t))$?

> 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?

- In terms of the equation below Eq.(16), we see that $\beta$ increases from $0.01$ to $4$, implying that we are adding more and more noise for the forward policy. This approach works for diffusion generative modeling. However, the sampling process starts with simple Gaussians and ends with samples, so shouldn't $\beta$ decrease (i.e., $\beta$ decreases from $4$ to $0.01$)? This is my previous concern.

[Updated] Thank you for providing new results.
As promised, here are the results for DDS and DIS on Manywell.

First, we clarify the difference between the task as considered in DIS and in this work, noting that both have been studied in other papers as well and neither choice seems obviously better.

|    | Ours | DIS (hardest version) |
|---|---|---|
| Dimension | 32 | 50 |
| No. of double wells | 16 | 5 |
| No. of modes | $2^{16}$ | $2^5$ |

We consider Manywell as a product of 16 double well distributions having unnormalized density given by $R(x,y) = \exp(-x^4 + 6x^2 + 0.5x - 0.5y^2)$, following [Midgley et al., "Flow annealed importance sampling bootstrap", ICLR'23], whereas in the DIS paper, the authors consider the more general case of the distribution being a product of some number of double well distributions. However, despite their considering this problem in higher dimensionality, they are stacking a smaller number of double wells (no more than 5), which makes the number of modes in distribution much smaller.

We ran DDS and DIS with our setting (after searching for the best hyperparameters, DDS was still somewhat unstable and often encounters an explosion of the loss and mode collapse, in which case we take results from the preceding checkpoint):

| Algorithm | $\Delta\log Z$ | $\Delta\log Z^{\rm RW}$ | ${\cal W}_2^2$ |
|---|---|---|---|
| DIS (4 runs) | 10.52 ± 1.02 | 3.05 ± 0.46 | 5.98 ± 0.46 |
| DDS (3 runs) | 7.36 ± 2.43 | 0.23 ± 0.05 | 5.71 ± 0.16 |
| VarGrad + LP + LS (ours) | 4.11 ± 0.45 | 0.02 ± 0.21 | 5.30 ± 0.02 |

We see that our best method outperforms both DIS and DDS in both metrics, but both are comparable with the various baselines in Table 1 of the paper.

We hope you find this comparison helpful. Thank you again for your time.
Thank you for following up.

As pointed out in [41, Section 3.1], the referenced paper [1] makes some serious errors in math -- incorrectly doing a change of variables in an integral -- that make the algorithm unsound. Formulating this task as a continuous sampling problem requires access to the Jacobian of the environment’s transition function, which isn’t available.

As noted in the initial response, we have demonstrated the effectiveness of our sampling method on a standard collection of tasks in the diffusion samplers literature. If you see a task of higher difficulty that is standard in the field, please let us know.
Thanks for the detailed responses. However, the reviewer still considers it is necessary to do some experiments on high-dimensional real-world applications to show the applicability of the proposed method or make enough theoretical contribution like DDS anyway. Otherwise, it is hard to judge whether the proposed method is useful and efficient. 

In that case, the reviewer decide to reject this paper and wish the authors could find appropriate complex real-world applications for their method later. Moreover, I guess a potential application for your method is the robot locomotion task, which is considered in CFlowNets [1].

[1] Li Y, Luo S, Wang H, et al. Cflownets: Continuous control with generative flow networks[J]. arXiv preprint arXiv:2303.02430, 2023.