From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training

Authors: Julius Berner, Lorenz Richter, Marcin Sendera, Jarrid Rector-Brooks, Nikolay Malkin
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: diffusion, variational inference, SDEs, PDEs, sampling, stochastic processes, GFlowNets

Abstract:
We study the problem of training neural stochastic differential equations, or diffusion models, to sample from a Boltzmann distribution without access to target samples. Existing methods for training such models enforce time-reversal of the generative and noising processes, using either differentiable simulation or off-policy reinforcement learning (RL). We prove equivalences between families of objectives in the limit of infinitesimal discretization steps, linking entropic RL methods (GFlowNets) with continuous-time objects (partial differential equations and path space measures). We further show that an appropriate choice of coarse time discretization during training allows greatly improved sample efficiency and the use of time-local objectives, achieving competitive performance on standard sampling benchmarks with reduced computational cost.

Decision: reject
Reviewer scores: 5, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1729480344263

thread_type:
official_review

<thread_object>
creator: reviewer_ofv6

<content>
summary:
This paper explores the connection between continuous-time SDEs and their discretization, particularly focusing on the influence of the chosen timestep. It demonstrates that using non-uniformly discretized time with fewer steps can achieve similar performance during inference. Theoretical results are provided to support this approach.

strengths:
* The problem studied in this paper is well-motivated.

* This paper presents extensive results in both theory and experiments.

* The appendix provides a comprehensive complement to the main text.

weaknesses:
* The theoretical results in Section 3 primarily focus on the convergence of the Euler-Maruyama method. Specifically, they show that convergence is ensured as the maximal step size approaches zero. However, these results do not explain why non-uniform discretization would generally be superior to uniform discretization. The advantage of non-uniform discretization—one of the main contributions of this paper—is demonstrated only through experiments

* As previously mentioned, there seems to be a gap between the theoretical and empirical sections of this paper. After reading the introduction, I expected to see concrete theoretical results that justify the use of non-uniform discretization. However, simply showing that convergence is guaranteed as $\Delta t$ approaches zero is unsurprising. The authors might consider adding more discussion on why uniform discretization is not always the optimal choice

* It has been proven that the order of convergence is determined by the step size, and the Euler-Maruyama scheme with uniform discretization has been shown to achieve optimal performance in the general case (see 'Numerical Treatment of Stochastic Equations' by Rümelin, 1982). I wonder if the claim made in this paper contradicts that result.

I would be willing to increase my rating if the authors are able to address my concerns.

questions:
Please see the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730471409054

thread_type:
official_review

<thread_object>
creator: reviewer_hyjk

<content>
summary:
This paper discusses the relationship between continuous and discrete-time stochastic processes and their training. In particular, the main results give a series of propositions on how a discrete-time process can approximate a continuous time process. I have to say I had a hard time understanding the "big picture" of the authors' results.

strengths:
The paper appears to be mathematically rigorous and experiments appear to give credence to the authors' work.

weaknesses:
I found the paper very difficult to read. The notation is dense, not all appears to be defined, some is non-standard and unclear and I found it a little tricky to understand exactly what the authors wanted to do. It may be that the authors have solved in interesting problem in a genuinely useful way but that was unclear from the paper. All but the very expert reader would, in my view, find the paper a difficult read. 

A few specific comments are:

- Abstract could be more informative and precise
- Introduction is quite meandering and I wasn't quite clear on exactly what the authors were trying to do.
- Figures 1 and 2 were placed, in my view, quite early in the paper and were hard to interpret. They needed more textual description, or, considering where they were placed, needed "dumbing down" a little. 
- Equation (1) is somewhat standard but, for completeness, it would have been useful to know what \sigma(t) is (I could guess). Equation (1) is similar to (9) apart from \mu(t). I think the differences between the various forms of \mu(t) needs to be explained in more detail.
- It wasn't clear to me exactly what the reverse arrow meant in terms of policy e.g. the backwards arrow is used on \pi(t) below equation (3) but without any definition as far as I can see. 
- I found Section 2 quite muddled with various different concepts introduced with not too much explanation. I realise there is a page limit, but it was bordering on the unpenetrable. 
- I didn't really understand how the Propositions in Section 3 ended up affecting the Results in Section 4. Perhaps I am dense, but it would be good if the authors could explain this better.

questions:
My main question is this: for a ML practitioner, how will the authors' results help?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730572255420

thread_type:
official_review

<thread_object>
creator: reviewer_k8pm

<content>
summary:
This paper examines training neural stochastic differential equations (SDEs) to sample from Boltzmann distributions without target samples. This work derives asymptotic equivalences by linking discrete-time policies to continuous-time diffusion. The approach is validated on sampling benchmarks.

strengths:
1. The approach of linking discrete-time policy objectives with continuous-time SDE training is a useful idea, albeit heavily reliant on established results.

2. Authors show that this method potentially reduces computational costs for neural SDE training.

weaknesses:
1. Firstly, I think the presentation of this work remains a major bottleneck for readers. Section 2 is preliminary, and it spans from pages 3 to 7. Such a lengthy preliminary section introduces well-known equations and results (e.g., equations (4)-(6) from GFlowNet papers, (9)-(15) from stochastic control and diffusion models, and (16), (17) as standard Euler-Maruyama discretizations). 
These derivations, mostly grounded in existing work, dilute the contributions and add an undue burden for readers. Figures like Figure 3, which illustrate obvious points, seem unnecessary and further contribute to this issue. It is recommended to present additional informative and easy-to-follow diagrams in these sections.

2. The primary theoretical contribution—showing asymptotic convergence from Euler-Maruyama discretization to continuous-time SDEs (Propositions 3.2, 3.3, 3.4)—seems not surprising. The convergence results are probably straightforward applications of established SDE theory, with little added insights or unique techniques. Without further exploration of new derivation techniques or distinctive theoretical angles, the contributions feel like direct applications of existing results.

3. The experiments are conducted on standard synthetic benchmarks, such as Gaussian mixtures and low-dimensional toy distributions. To support this approach, it might be necessary to conduct higher-dimensional Bayesian inference tasks where the Boltzmann distribution is more untractable. Besides, the compared baselines exclude many recent models, such as flow-based generative models. 

3. While efficiency is demonstrated, additional benchmarks comparing computational costs with traditional methods in larger dimensions would be helpful for real-world applications.

questions:
1. Could the authors clarify why so much space is devoted to standard results? Would simplifying or condensing this content help highlight the unique contributions?

2. Beyond applying existing convergence results, what novel techniques, if any, were introduced in proving Propositions 3.2, 3.3, and 3.4?

3. Would more complex or realistic benchmarks alter the experimental outcomes, particularly in high-dimensional or non-Markovian sampling settings?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731352390988

thread_type:
official_review

<thread_object>
creator: reviewer_glpg

<content>
summary:
This paper investigates the training of diffusion samplers and neural stochastic differential equations (neural SDEs) by examining the connection between continuous-time objectives and their discrete-time counterparts. The authors establish that global objectives for discrete-time policies converge to path-space measure divergence objectives in the continuous-time limit, while local constraints asymptotically align with partial differential equations governing the time evolution of marginal densities. This theoretical grounding aims to bridge reinforcement learning (RL) objectives and stochastic control frameworks for diffusion processes. Empirically, the paper demonstrates that training with coarse, non-uniform time steps, particularly with random placements, can achieve substantial computational efficiency gains while retaining strong performance across a range of benchmarks.

strengths:
1. The paper is very well written and easy to follow, with clear exposition of the mathematical derivations and the empirical results.
2. The experimental section is thorough and well designed, exploring the effects of different discretization strategies and their impact on performance in detail. The benchmarks used are diverse and represent a wide range of sampling challenges.
3. The work provides strong empirical evidence that non-uniform time discretization (particularly random placement) improves training efficiency. This observation could be highly relevant for practitioners working with high-dimensional diffusion models. Furthermore, the identification of random time discretization as a performant strategy is novel and supported by robust experimental evidence.
4. The paper effectively summarizes existing methods and objectives for diffusion sampling, offering a clear context for the proposed contributions and situating them within the broader body of work on diffusion models and sampling techniques.

weaknesses:
1. While the theoretical contributions are valuable and provide an interesting link between discrete-time and continuous-time objectives, they are not completely unexpected and partly already present in the literature.
2. In the experimental results, it is noted that the ELBO gap does not converge to zero as the discretization becomes finer but instead appears to stabilize at a positive value. The authors do not give an explanation for this phenomenon. In particular, the lack of a "benchmark" makes difficult to connect these simulations to the numerical results presented in the first part of the paper above.
3. The observed performance gains with randomly placed time steps are well supported by empirical results, but the paper does not provide a theoretical explanation for why this approach works so well. Offering more insight into this phenomenon would enhance the overall impact of the findings.

questions:
1. Is it correct to expect that the ELBO gap should converge to zero as the discretization becomes finer, or are there inherent limitations in the approach that cause the gap to saturate at a positive value? Clarifying this could help contextualize the observed results better.
2. Are there any existing benchmarks or prior work that provide a comparable measure of ELBO gap performance for optimally trained diffusion samplers? How do the proposed methods stack up in this context?
3. Can the authors provide more insight into why random placement of time steps works so (unexpectedly) well? Is there an intuitive or theoretical rationale for this observed behavior?
4. In Theorem 3.4, there seems to be a potential issue as $\vec μ_t$​ appears twice in the statement. Could this be a mistake, or is there a specific reasoning behind this repetition? Clarification would be helpful.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732210238576

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers

comment:
We thank all the reviewers for their comments. The suggestions have helped us improve the paper, and we have uploaded a revised version with the key changes highlighted in orange.

### On the importance of theoretical results and their relevance of theory to the experiments

For diffusion models trained to maximize a variational bound on data log-likelihood, the denoising score-matching objective is equivalent to maximization of a continuous-time ELBO, equivalently, minimization of a KL divergence between reverse and forward path measures. The implications of this fact for the ability to train a diffusion model in one discretization (or with a continuous time parameter) and sample it in another are well understood starting from [Song et al., 2021a](https://arxiv.org/abs/2011.13456) and [Huang et al., 2021](https://arxiv.org/abs/2106.02808).

**Until our work, a similar result for diffusion samplers of unnormalized densities, in particular ones based on "off-policy" divergences such as VarGrad and the GFlowNet-inspired losses, has not been known.** In fact, most past work on diffusion samplers has ignored the fact that exact time-reversal with Gaussian transition kernels is *not generally possible in discrete time*. Instead, a discrete-time objective that cannot, in theory, be taken to 0 by *any* sampler is minimized, yet the consequences of the discretization error *in training* are not studied. The closest results we are aware of in this direction are Proposition E.1 in [Vargas et al., 2024], which concerns trajectory-level Radon-Nikodym derivative but not the other functionals involved in off-policy losses, and Proposition 9 in [Zhang et al., 2023], which establishes a connection between the limit of detailed balance and score matching, but only considers the $\sqrt{h}$ asymptotics for a fixed reverse process.

**We substantially generalize these known results, both for the global (KL, second-moment) and local (detailed balance) divergences:** We show that the discrete-time objectives asymptotically approach continuous-time objects and indicate the order of asymptotics (0th-order in $\Delta t$ for the trajectory-level divergences (Proposition 3.3) and 0.5th-order and 1st-order for the two results on detailed balance (Proposition 3.4)). 

While it is not entirely *unexpected* that such convergences would hold, it is also **not obvious *a priori***. For example, one could imagine that a sampler trained with $N$ discretization steps would acquire a bias, relative to the ideal continuous-time sampler, that depends on $N$ and scales as $O(\sqrt N)$ as $N\to\infty$. We show that this does not happen. The proofs of our results, although they do not require any entirely new proof technique, are not trivial. They require careful application of stochastic calculus results: in particular, we are not aware of the method of proof of convergence of functionals (Proposition B.3) -- via weak convergence (Proposition 3.1) and convergence of Radon-Nikodym derivatives (Lemma B.7) -- being used in relevant literature. For instance, [Vargas et al., 2024] present an RND discretization in their setting (Proposition E.1), but the convergence as well as implications for *divergence* minimization are not shown.

These convergences imply that objectives evaluated with $N$-step discretization, for different $N$, are *all approximating the same continuous-time object*, which justifies training and inference with different numbers of time steps. If we did not have such convergence, there would be no reason to expect that a sampler trained with $N_{\rm train}$ and sampled with $N_{\rm eval}$ discretization steps would have bias approaching $0$ as $N_{\rm train},N_{\rm eval}\to\infty$. 

**The experiments are designed to illustrate the practical implications of the fact that training and inference with different numbers of time steps is theoretically justified.** The observations regarding the choice of training discretization, which allow to greatly reduce the computation cost of training, are **an interesting empirical result in their own right** that will be of interest to the growing community working on diffusion samplers. They also provide an example of a practicable use of (less expensive) local-time objectives, which so far had not been seen to scale well with long trajectories.

Thanks to your feedback, we added an additional discussion at the end of Section 3.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper considers the training of the neural stochastic differential equations to sample from Boltzmann distributions. By drawing connections between the discrete-time policies and continuous-time diffusion, an asymptotic equivalence between them has been established. As pointed out by the reviewer K8pm, the main results (asymptotic convergence from Euler-Maruyama discretization to continuous-time SDEs) are standard, and most of the preliminary results derived in the paper are either known or standard extension of known results, the contribution of the paper is marginal. 

Nevertheless, the work is still having a good potential as the authors have experimentally shown that non-uniform time steps (in particular the random one) can provide significant performance gain, but this observation does not have any theoretical support as the current paper only provides standard results on arbitrary discretization scheme. If theory can be derived to explain this observation. It will resolve the novelty issue from the reviewers. 

Based on the current feedback, the paper is marginally below the standard of ICLR, we have to reject this paper.