Distributional Diffusion Models with Scoring Rules

We thank the reviewer FCZ5 for their feedback on our work. Please see our response below

|$p_{0|t}(x_0|x_t)$,Dirac…not true

Using $p_{0|t}(x_0|x_t)\approx\delta_{\hat{x}_\theta(t,x_t)}$ in $p(x_s|x_t)$ (see eq.5) leads to

$p(x_s|x_t)\approx p(x_s|x_0=\hat{x}_\theta(t,x_t),x_t),$

where $p(x_s|x_0,x_t)$ is Gaussian (see eq.4). Sampling from these approximations of $p(x_{t_k}|x_{t_{k+1}})$ for some $t_N=1>t_{N-1}>....>t_1>t_0=0$ gives exactly the DDIM algorithm (see DDIM paper [1],eq.10)

|Other metrics,diversity

FID [2] already quantifies the diversity of the sampled images as it is Wasserstein-2 distance between Gaussian approximations to the distributions in the space of Inception V3 coding layer features, and it takes into account the covariances of these distributions. See [3] regarding image metrics for evaluation (i.e., “[the FID score] unlike IS [..] is able to detect intra-class mode dropping”)

A good metric for evaluating diffusion policy in robotics is an open challenge and so far success rates during independent evaluation is the only measure that has proven useful for robotics. Systematic evaluations of various metrics for policies yielded no reliable metrics [4]

|more 2D

We trained unconditional standard diffusion and our distributional models on a more complex 2D distribution – checkerboard, see [5]. We report MMD between sampled and target distributions with different NFEs, smaller MMD is in bold. Our approach outperforms standard diffusion for small NFEs. We will add these results to the paper

|[5-6]

Thank you for pointing out these references. We’ll discuss as follows:

“Scoring rules have corresponding MMD divergences, and approaches exist that use the MMD in diffusion and particle flow models. [5] generates a sequence of distributions from a forward diffusion process. It then generates noise-dependent neural MMDs between clean and noisy data, and performs gradient flow to move particles from noise level $t$ to a lower noise level $s<t$. Compared to ours, [5] does not use a generator since it is a particle flow. Their best CIFAR10 FID is 7.7. [6] first trains a diffusion model, and then refines/distills by coarsening the reverse timesteps, and uses MMD on CLIP features to finetune the DDM denoiser. Their CIFAR10 FID for NFE=10 improves from 13.6 to 3.8, while ours is 3.19.

|m sensitivity

Please see our response to reviewer r8Za about m and computational complexity. For image experiments $m=2,4$ already worked well. We will add an ablation over $m$ for CIFAR10 in the revised version of the paper. For robotics, we ran an ablation over $m$ ($\beta=1,\lambda=0.5$), see results below

|IMQ,RBF,kexp

Proper scoring rules define divergences. For kernel scores, this is the squared MMD. The energy distance kernel is appealing – its loss (13) is similar to standard diffusion loss, and other kernels can recover this loss in certain limits. While this recovery is good for interpretability, any characteristic kernel (including the $kexp$) gives a valid scoring rule. A benefit of (13) is its numerical stability without exponentiation

|additional experiments

Please see our response to reviewer waLv for comparisons with DPM-solver++ and SOTA multi-step distillation.

We started an experiment on ImageNet (64x64), but due to time constraints, will only add results after rebuttal. As for robotics, we studied 10 most challenging tasks from Libero benchmark, which is already a diverse set of language-conditioned robotics tasks. We started an experiment on all tasks from Libero and will report results after rebuttal

|$\Sigma_{s,t}$...why learning necessary?

We need to approximately sample from $p(x_s|x_t)=\int p(x_s|x_0,x_t)p(x_0|x_t)dx_0$ where $p(x_s|x_0,x_t)=N(x_s;\mu_{s,t}(x_0),\Sigma_{s,t})$ is fixed. We obtain a better approximation by learning $p(x_0|x_t)$ with scoring rules instead of Dirac mass as denoiser. Minimizing loss (13) recovers the true $p(x_0|x_t)$, if $p^\theta(x_0|x_t)$ is rich enough. The resulting $p(x_s|x_t)$ is indeed non-Gaussian and its form is fully problem dependent.

We hope that our explanations and updated experiments have addressed your concerns. If so, we would be grateful if you would consider raising your score.

[1]Song J et al.,Denoising Diffusion Implicit Models,ICLR2021

[2]Heusel M et al.,GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium,NeurIPS2017

[3]Borji A,Pros and Cons of GAN Evaluation Measures,2018

[4]Hussenot L, et al.,Hyperparameter Selection for Imitation Learning,2021

[5]Galashov A et al.,Deep MMD gradient flow without adversarial training,ICLR2025

[6]Aiello E et al.,Fast inference in denoising diffusion models via MMD finetuning,IEEE2024

[7]Arbel M et al.,MMD Gradient Flow,NeurIPS2019

We thank reviewer r8Za for their overall positive assessment of our paper. Below, we address your comments.

| Weakness 1

As suggested, we have now compared our method to the DPM-solver++ [1]. We have also compared it to a SOTA (multistep) distillation method [2]. Please see our response to the reviewer waLv above, where we added these results.

Our results show that our approach is competitive with multistep distillation (we compared different amounts of steps) and outperforms DPM-solver++ over 8 NFEs (whose performance degrades as expected as the number of steps increases, due to numerical instability). Two compelling benefits of our approach over multistep distillation are that it does not rely on distillation, so there is no need to train a large (slow) teacher model; and that it does not require specifying additional sampler hyperparameters during training. It is certainly possible for our distributional approach to be combined with a modern distillation method. This is an interesting topic for future research.

| Weakness 2

We thank you for bringing up this point regarding the training time complexity. Indeed, as you suggested, our paper will benefit from adding a complexity analysis of our method. We will add the following section to the appendix of our paper.

Computational complexity

Our method has the same computational complexity during sampling as ordinary diffusion models, and only has an increased computational complexity during training which is detailed below.

Assume that the computational complexity of the forward pass in our diffusion model is $O(F)$ and that the dimensionality of $x$ is $D$.

Diffusion loss case. The loss function in eq. (14) with $m=1$ and $\lambda=0$, can be thought of as a standard diffusion model loss function. It requires $O(nF)$ evaluations to compute $x_{\theta}(t_i,X_{t_i})$ for every element of a batch. Then, to evaluate the loss (the norm), the complexity is $O(nD)$. Therefore the total cost of computing the loss is $O(n (F+D))$. The backwards pass is proportional to the forward pass and has computational complexity $O(nF)$.

Distribution loss case. For $\lambda \neq 0$ and $m>1$, we first need $O(nmF)$ function evaluations to compute $x_{\theta}(t_i, X_{t_i}, \xi_{i,j})$ for $i=1,\ldots,n$ and for $\xi_{i,j}$, $j=1,\ldots,m$. After that, we need $O(nmD)$ operations to compute the first (diffusion-like) terms of the loss. Then, in order to compute the interaction terms, we need $O(nm(m-1)D)$ time. Therefore, the total cost is of computing the loss $O(nm(F+D) + nm(m-1)D)$. Naively, backwards pass will take $O(nmF + nm(m-1)F) = O(nm^2F)$ time. However, since the 2nd term uses the gradients of $x_{\theta}(t_i, X_{t_i}, \xi_{i,j})$, one could precompute these for all $i=1,\ldots,n$ and $j=1,\ldots,m$ and therefore decrease the total backwards cost to $O(nmF)$. We found that in practice, the XLA compiler in jax performs this optimization without explicit coding (see below).

Runs on real hardware. We compared the training times on real hardware. We study the training of CIFAR-10 and we vary $m$. We report steps per second metric where a step corresponds to a full forward+backward step. As hardware, we use A100 GPU (40 Gb of memory) with batch size = 16, H100 GPU (80Gb of memory) with batch size = 64 and TPUv5p (with 95Gb of memory) with batch size = 64 (per device with 4 devices in total). The results below indicate that steps per second decreases proportionally to $m$.

| Could this framework be extended to score-based models with ODE sampling?

If the churn parameter $\epsilon$ is taken to be zero in the DDIM equations, one recovers a deterministic sampler which can be thought of as the discretization of an ODE. This ODE is a special case of the standard probability flow ODE; see Section 4.3 of [3] for a discussion. In particular, the DDIM update coincides with the exponential integrator for the probability flow ODE.

We emphasize that our method can also be extended directly to sample a transport map between any two distributions $p_0$ and $p_1$ and does not require one of the distributions to be Gaussians. This is a framework for which flow matching/stochastic interpolants are typically used. We will add a section in our paper to detail this extension.

We hope that our updated experiments have addressed your concerns. If so, we would be grateful if you could consider increasing your score.

References:

[1] DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models, Cheng Lu et al., 2022

[2] Multistep Distillation of Diffusion Models via Moment Matching, Tim Salimans et al., 2024.

[3] Denoising diffusion implicit models, Jiaming Song et al., 2020.

We would like to thank reviewer waLv for their comments and positive feedback.

As requested, we re-implemented and have added comparisons to a popular numerical solver DPM-solver++ [1] (as suggested by reviewer r8Za) and to the recommended distillation method [2]. We did implement the parallel sampling of diffusion models but haven't had enough time to sweep over hyperparameters to obtain competitive results. We will add this comparison as well as the comparisons below, in the final version of the paper.

Experimental details: For all the sampling methods, we sweep over the “safety” parameter $\eta$, which specifies the time interval to be “[$\eta$,$1-\eta$]”. We consider values $[1e-1,1e-2,1e-3,1e-4]$. Moreover, for all the methods we sweep over the “churn” parameter which controls the stochasticity, we tried $[0.0,0.25,0.5,0.75,1.0]$. We found the stochastic version of DPM-solver++ (with a churn parameter close to $1$) overall worked the best. For ordinary diffusion and our method, we use a “uniform” time schedule. For DPM-solver++ we use the “logsnr” time schedule (see [1]), which led to the best results (we also tried EDM and uniform).

For distillation comparisons, we distill pretrained diffusion models into a student model with a given number of student sampling steps. We train a distilled model for 100k iterations with batch size 256 for pixel-space and with batch size 32 for latent space models. We sweep over EMA decay and learning rate. For the student, we use a stochastic DDIM sampler with a churn parameter $\epsilon$ and we sweep over it with considered values $[0,0.5,1.0]$. At sampling time, we found that using a different churn parameter $\epsilon=0$ for the distilled student, led to the best results.

For all the methods, we present the results with the best hyperparameters.

Results: The results are presented in the table below. They indicate that our approach is competitive with multistep distillation (we compared different amounts of steps) and outperforms DPM-solver++ over 8 NFEs (whose performance degrades as expected as the number of steps increases due to numerical instability). Two compelling benefits of our approach over multistep distillation are that it does not rely on distillation, so there is no need to train a large (slow) teacher model; and additionally it does not require specifying additional sampler hyperparameters during training. It is certainly possible for our distributional approach to be combined with a modern distillation method. This is an interesting topic for future research.

The tables below contain FID scores (averaged across 3 random seeds) for different number of function evaluations (NFEs) and for different methods. In bold, we highlight the method with the lowest FID and in italic, we highlight the method with the second lowest FID.

Dataset name = CIFAR-10 (Conditional)

Dataset name = CelebA (Conditional)

Dataset name = LSUN (Unconditional)

Dataset name = Latent CelebA-HQ (Unconditional)

We hope that our updated experiments have addressed your concerns. If so, we would be grateful if you could consider increasing your score.

References:

[1] DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models, Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, Jun Zhu, 2022

[2] Multistep Distillation of Diffusion Models via Moment Matching, Tim Salimans et al., 2024.