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

Dear Reviewer ofV6,

We appreciate your constructive feedback on our paper.

On theory and experiments

We kindly direct you to the response to all reviewers for a detailed answer.

The theory and experiments serve somewhat complementary purposes: while the theory establishes the (perhaps unsurprising, but certaintly not trivial) fact that training with varying time discretization is well-behaved in the limit, the experiments explore the implications of the fact that such training is theoretically justified. Our exploration of nonuniform discretization yielded interesting empirical results, but they are indeed not explained by the theoretical ones (nor by any theory in existing work that we are aware of).

We suspect the main reason that nonuniform time discretization is helpful is that with uniform discretization, the sampler overfits to the small, fixed set of inputs of the time variable to the neural network computing the drift. This is supported by the evidence that our proposed Equidistant discretization -- in which the distribution over time steps seen during training has nearly full support but all time increments except the first and last have length $\frac1N$ -- gives similar results to the Random discretization, in which increments have different sizes.

We are happy to include extended discussion of this in the paper.

On Euler-Maruyama numerics

There may be a small misunderstanding present regarding the Rümelin's result and its relevance.

First, please correct us if we are wrong, but Rümelin, in the paper you referenced, assumes uniform step size $h$ and proves the orders of convergence of various integrators (Euler-Maruyama, Heun, Runge-Kutta) under this assumption (see the bottom of p.605 here).

Second, such a result would concern the order of convergence of integration, which we are interested in when sampling a trained model. Indeed, in our experiments, we evaluate all models with a uniform time discretization. However, the choice of discretization is important during training, where (as shown by our results) a nonuniform discretization gives a better approximate learning signal than a uniform one. Because the gradient of the continuous-time divergence is not an Itô integral whose Euler-Maruyama approximation is the gradient of the corresponding divergence in the discretization, this does not contradict the mentioned result.

Thank you again for your comments. We are happy to answer any further questions you may have.

Thank you for your response and for adjusting your score.

We believe that we can clarify your concern regarding the theoretical gap.

First, your concern seems to be about inference-time, not training-time, discretization. Thus, we just want to repeat that all our models are evaluated with uniform time discretization.

Now, let us discuss the choice of discretization scheme in each phase:

Uniform discretization at inference time: During inference, i.e., sampling, of all trained models, we use a uniform discretization. Thus, the variance in the performance of the Random and Equidistant training methods during inference is comparable to the variance of the Uniform method.

Regarding the discretization error for SDE inference in general:

• We can leverage results in, e.g., De Bortoli et al., NeurIPS 2021 (Section 2.2) to bound the error between the data distribution and the learned distribution for a given step-size of the Euler-Maruyama scheme and approximation of the optimal policy.
• The dependence of the error on the discrepancy of the learned sampler from the optimal one is also studied in Lee et al., NeurIPS 2022 and Chen et al., ICLR 2023 (Section 3.1).
• Bounds specific to neural network function classes are also proved by Oko et al., ICML 2023.

However, in our case, the 100-step uniform integration we use -- following many past studies on the same set of target densities -- is already quite close to continuous-time integration for the distributions considered (Figure 5 is evidence of this fact). Summarizing, the inference-time discretization scheme is not the main source of error, nor is it the focus of our work.

Randomized discretization at training time: Our results show that the non-uniform discretization schemes applied during training give better performance with uniform-step-size inference. Our theoretical results show that the learning objective under any discretization scheme converges to the same continuous-time object as the maximal step size goes to zero. However, it does not tell us which discretization scheme for training is optimal given a fixed number of steps (cf. our original answer above).

The results suggest that the reason for the better performance of the non-uniform discretization schemes is simply one of generalization over the values $t$ given as input to the drift model $\overrightarrow{\mu}(x,t)$. If one trains in a uniform discretization with different step size than used for evaluation, there will be a generalization gap. However, with the non-uniform training schemes, values of $t$ that occur at evaluation time are also seen during training -- the generalization gap is thus smaller. Said otherwise, our results are compatible with the hypothesis that the time generalization error is more significant than the discretization approximation error at training time in the settings considered.

Does this clarify your concern?

Dear Reviewer hyjK,

Thank you for your review and you valuable feedback. Let us comment on your concerns in the following.

Presentation of the paper

While we put effort in making the paper as accessible as possible, we acknowledge that the presentation can be dense and requires certain background knowledge. Given the space constraints, we tried to strike a balance between providing necessary background information and citing the relevant literature for further details. However, the level of necessary detail naturally depends on the background of the reader. For instance, we also received feedback from another reviewer that we should shorten the background sections.

In general, we need to introduce several concepts (in a mathematically rigorous way), such as measures on tractories, likelihood ratios, and objectives, in both discrete as well as continuous time, since they are required by our novel theory. While we try to offer intuition, such concepts rely on knowledge in (numerical) stochastic analysis, GFlowNets, and optimal control, which is difficult to fully explain in the limited amount of space. However, we purposefully start our presentation in discrete time, since this allows to understand the concepts without background in stochastic calculus.

As suggested by you, we polished our presentation and added several details (e.g., defining $\sigma$, unifying the notation for $\overset{\leftarrow}{\mu}$, explaining the backward policy $\overset{\leftarrow}{\pi}$) and (slightly) simplified Figure 2. Moreover, we added a paragraph explaining the connection between our theoretical results and the experiments at the end of Section 3.

Given that a main contribution of our paper is of theoretical nature (a rigorous foundation for the training of diffusion-based samplers), we hope that you understand that we need to assume certain theoretical background knowledge. Please let us know if you have further suggestions on improving the accessibility of our paper.

Practical implications of our work

On a high-level, we show for the first time that different discretizations of training objectives approximate the same, unique continuous-time objective (in the limit of refining the discretization). Our results have a simple, but profound implication:

One can use different discretizations for training and inference of diffusion-based samplers without incurring bias.

This finding is particularly important given that diffusion-based methods for sampling problems rely on a SDE discretization during training, incurring significant computational costs. Motivated by our theoretical results, we show for the first time that a few, randomized steps during training can achieve similar performance at a fraction of the cost. Our theoretical results (covering local as well as global objectives) hold for virtually all existing diffusion-based samplers and we also obtain consistent empirical outcomes across methods and tasks.

We note that our findings are specific to sampling problems. In generative modeling, where we have access to samples from the target distribution, one leverages score-matching objectives that allow training in continuous time (without SDE discretizations), i.e., directly optimizing the continuous-time ELBO. We present further explanations in our general response and hope that this helps to understand the significance of our work to ML practitioners.

Thank you again for your comments. We are happy to answer any further questions you may have.

Dear Reviewer K8pm,

Thank you for your extensive and constructive feedback. We will comment on your concerns and questions in the following:

On the presentation of our work

Thank your for providing suggestions on the presentation of our paper. Generally, we believe that there is no prior work that connects these "preliminary" results (many of which only appeared in the last two years) and we hope that Figures 1 and 2 can serve as informative diagrams for navigating the paper. The motivation behind the current presentation is to make our theoretical results accessible to as many readers as possible:

1. We start with the discrete-time approach since it requires less background knowledge in stochastic calculus. We agree that these are well-known results from GFlowNet papers, however, we think that many readers might not be familiar with this theory. For instance, GFlowNets themselves have only been invented in 2021 (by Bengio et al.) and the theory for continuous state spaces was developed only in 2023 (by Lahlou et al.). Similarly, the off-policy losses have only been introduced to diffusion samplers in 2023 or 2024. However, based on your suggestion, we tried to shorten the exposition and moved Figure 3 to the Appendix.
2. In the section on the continuous-time setting, we need to introduce concepts such as Nelson's identity and the Fokker-Planck equation since our results are based on them. Moreover, we want to emphasize that the Radon-Nikodym derivative between general forward and backward SDEs, as well as the resulting KL divergence are very recent results that appeared only this year (by [Vargas et al., 2024] and [Richter & Berner, 2024]). Thus we would not call them well-known results.

Please let us know if this clarifies our exposition. Otherwise, we are happy to restructure our paper further to faciliate readability.

Theoretical claims and contribution

We would like to refer you to the general response for a detailed exposition of our contributions. In particular, we detail why our theoretical results provide substantial new insights and are not just straightforward applications of existing theory.

More challenging tasks

We want to emphasize that our tasks include many of the arguably most challenging tasks on which samplers are typically benchmarked (see, e.g., the recent large-scale benchmark by Blessing et al., 2024). In particular, our tasks cover up to $1600$-dimensional problems in Bayesian statistics, $32$-dimensional highly multi-modal ($2^{16}$ modes) problems related to models in molecular dynamics, as we as complex posterior distributions from variational auto-encoders trained on image distributions.

Dear Reviewer K8pm,

We would like to thank you once again for your review and feedback, which we believe makes this paper stronger. We've responded to all the issues that you mentioned and added the additional experiments on flow-based generative models as baselines.

Since the discussion period is approaching its end, could you please let us know if these have satisfactorily addressed your concerns? Please, let us know if you have any additional questions. We look forward to hearing from you.

Thank you,

The authors

Dear Reviewer gLPG,

Thank you your extensive review and appreciating our clear exposition as well as thorough experiments. Let us answer your remaining questions and concerns in the following:

On our theoretical contributions

While perhaps not being completely unexpected, we emphasize that there is also not necessarily a reason to believe that training objectives evaluated at different discretizations converge to a unique continuous-time object. While the required background knowledge in Section 2 is known (as also referenced in our paper), our links between discrete-time and continuous-time objectives are not present in the literature. In fact, previous literature has overlooked potential issues of training in discrete time. We present further clarifications in our general response, also contrasting our results with the setting in generative modeling.

Practical implications and intuition

In the general response, we also elaborate why our theoretical results are crucial for being able to train and evaluate at different discretizations (see also the new paragraph at the end of Section 3). Since our results guarantee that we are approximating the same continuous-time objective with different training discretization, we have not been too surprised by the great performance using only a few time-steps during training. In particular, our results also explain why other considered randomized discretizations (such as "equidistant") offer performance improvements similar to our considered "random" one; see Appendix D.1. Intuitively, one could argue that a fixed uniform discretization makes the model overfit by trying to counteract the discretization error incurred the SDE integrator. Since the time-step $\Delta t_n$ is fixed, the model does not learn the optimizer of the continuous-time objective and cannot generalize to other discretizations during inference.

On the ELBO gap

We note that for the ELBO $\log\widehat Z$ it holds that $\mathbb{E}[\log\widehat Z] = \log Z - D_{\mathrm{KL}}(\widehat{\mathbb{P}},\widehat{\mathbb{Q}}) \le \log Z$, which we made clearer in our metrics in Appendix D.1 now. In particular, this shows that the gap is given by $D_{\mathrm{KL}}(\widehat{\mathbb{P}},\widehat{\mathbb{Q}})$, which is only zero if the forward and reverse discrete-time processes are perfect time-reversals. In general, we can not expect that our SGD-based training finds a global minimum. Beyond that (and independent of the expressivity of our neural networks), a perfect time-reversal with Gaussian transition kernels (as given by Euler-Maruyama discretization) is not generally possible in discrete time (as also mentioned in our paper). While this has been ignored in previous work, our theory guarantees that training with different discretizations nevertheless approximates the same continuous-time objective.

When the number of steps tend to infinity during inference, the gap is thus given by the continuous-time KL divergence $D_{\mathrm{KL}}(\mathbb{P},\mathbb{Q})$ (which only vanishes for a perfectly trained model). Nevertheless, the ELBO gap provides a principled metric since lower values necessarily correspond to a smaller KL divergence. It is typically also the only metric considered for tasks where the groundtruth normalizing constant and samples from the target are not known, see, e.g., [Blessing et al., 2024]. For all other tasks, we additionally provide the error in estimating the normalizing constant (see Table 2 in the appendix).

Typo

Thank your for pointing out the typo in Theorem 3.4, which we have fixed in the revised version.

Thank you again for your comments. We are happy to answer any further questions you may have.

Dear Reviewer gLPG,

We would like to thank you once again for your review and feedback, which we believe makes this paper stronger. We've responded to all the issues that you mentioned, explaining the ELBO gap difference, clarifying our theoretical contributions, and showing the practical implications of our findings.

Since the discussion period is approaching its end, could you please let us know if these have satisfactorily addressed your concerns? Please, let us know if you have any additional questions. We look forward to hearing from you.

Thank you,

The authors

Dear reviewers,

We would like to thank you again for your valuable feedback on the paper.

The end of the discussion period is approaching. We’ve responded to all of your comments and suggestions individually and in the above comment. Could you please let us know if these have satisfactorily addressed your concerns? We look forward to hearing from you.

Thank you,

The authors.