VarDiU: A Variational Diffusive Upper Bound for One-Step Diffusion Distillation

TL;DR

VarDiU introduces an unbiased gradient estimator for diffusion distillation, improving sample quality and training stability by directly optimizing a variational upper bound, surpassing previous methods like Diff-Instruct.

Contribution

The paper proposes VarDiU, a novel variational diffusive upper bound that provides an unbiased gradient estimate for one-step diffusion distillation, enhancing performance and training efficiency.

Findings

Achieves higher generation quality than previous methods.

Enables more efficient and stable training process.

Demonstrates effectiveness through comparison with Diff-Instruct.

Abstract

Recently, diffusion distillation methods have compressed thousand-step teacher diffusion models into one-step student generators while preserving sample quality. Most existing approaches train the student model using a diffusive divergence whose gradient is approximated via the student's score function, learned through denoising score matching (DSM). Since DSM training is imperfect, the resulting gradient estimate is inevitably biased, leading to sub-optimal performance. In this paper, we propose VarDiU (pronounced /va:rdju:/), a Variational Diffusive Upper Bound that admits an unbiased gradient estimator and can be directly applied to diffusion distillation. Using this objective, we compare our method with Diff-Instruct and demonstrate that it achieves higher generation quality and enables a more efficient and stable training procedure for one-step diffusion distillation.

Tables

Table 1. Table 1 : Comparison across different settings. ( ↑ \uparrow ) denotes the higher is better and ( ↓ \downarrow ) denotes the lower is better. Best performance is in bold. Second best is underlined. “True” indicates the samples from the target distribution. Diff-inst ( k k ) means Diff-instruct trained with k k score steps. We can find VarDiU achieve best results with true and empirical score and remain competitive with learned score.

Metrics	True	Diff-Inst (1)	Diff-Inst (5)	Diff-Inst (10)	VarDiU Gaussian	VarDiU NSF
Log-Density ($\uparrow$)	-6.65	$-10.17\pm 1.77$	$-8.09\pm 0.86$	$-7.86\pm 0.56$	$\underset{¯}{-7.00\pm 0.02}$	$-6.99\pm 0.01$
Log-MMD ($\downarrow$)	/	$-5.59\pm 0.36$	$-6.09\pm 0.40$	$-5.79\pm 0.36$	$\underset{¯}{-6.86\pm 0.29}$	$-7.33\pm 0.17$

Table 2. (a) Comparison with true score

Metrics	True	Diff-Inst (1)	Diff-Inst (5)	Diff-Inst (10)	VarDiU Gaussian	VarDiU NSF
Log-Density ($\uparrow$)	-6.65	$-10.17\pm 1.77$	$-8.09\pm 0.86$	$-7.86\pm 0.56$	$\underset{¯}{-7.00\pm 0.02}$	$-6.99\pm 0.01$
Log-MMD ($\downarrow$)	/	$-5.59\pm 0.36$	$-6.09\pm 0.40$	$-5.79\pm 0.36$	$\underset{¯}{-6.86\pm 0.29}$	$-7.33\pm 0.17$

Table 3. (b) Comparison with training data

Metrics	True	Diff-Inst (1)	Diff-Inst (5)	Diff-Inst (10)	VarDiU Gaussian	VarDiU NSF
Log-Density ($\uparrow$)	-6.65	$-9.65\pm 0.62$	$-8.38\pm 0.42$	$-7.96\pm 0.62$	$\underset{¯}{-7.09\pm 0.03}$	$-7.01\pm 0.02$
Log-MMD ($\downarrow$)	/	$-5.51\pm 0.26$	$-5.60\pm 0.30$	$-5.63\pm 0.68$	$\underset{¯}{-6.36\pm 0.22}$	$-6.81\pm 0.30$

Table 4. (c) Comparison with learned score

Metrics	True	EDM	Diff-Inst (1)	Diff-Inst (5)	Diff-Inst (10)	VarDiU Gaussian	VarDiU NSF
Log-Density ($\uparrow$)	-6.65	-6.69	$-10.88\pm 1.55$	$-9.21\pm 0.70$	$-8.47\pm 0.56$	$\underset{¯}{-8.13\pm 0.19}$	$-7.89\pm 0.15$
Log-MMD ($\downarrow$)	/	-6.67	$-5.24\pm 0.30$	$-5.57\pm 0.27$	$-5.82\pm 0.27$	$-5.64\pm 0.25$	$\underset{¯}{-5.68\pm 0.30}$

Equations

$$q(x_t|x_0)=\mathcal{N}(x_t|\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I),$$

$$p(x_{t-1}|x_t)=\mathcal{N}\left(x_{t-1}|{\mu }_{t-1}(x_t),{\Sigma }_{t-1}(x_t)\right),$$

$$\mu_{t-1}(x_{t})=\frac{1}{\sqrt{1-\beta_{t}}}\Bigl{(}x_{t}+\beta_{t}\nabla_{x_{t}}\log p_{d}(x_{t})\Bigr{)},$$

$$\mathcal{L}_{\mathrm{DSM}}(\psi)=\int p_{d}(x_{0})\,q(x_{t}|x_{0})\,\bigl{\|}\bm{s}_{\psi}(x_{t};t)-\nabla_{x_{t}}\log q(x_{t}|x_{0})\bigr{\|}_{2}^{2}\,\mathrm{d}x_{t}\,\mathrm{d}x_{0}.$$

$$p_{\theta }(x_0)=\int \delta (x_0-g_{\theta }(z))p(z)dz,$$

$$\mathrm{DiKL}(p_{\theta }||p_d)\equiv {\int }_0^1\omega (t)\mathrm{KL}(p_{\theta }^{(t)}(x_t)\Vert p_d^{(t)}(x_t))\mathrm{d}t,$$

$${\nabla }_{\theta }\mathrm{KL}(p_{\theta }^{(t)}\Vert p_d^{(t)})=\int \left[\left({\nabla }_{x_t}\log p_{\theta }^{(t)}(x_t)-{\nabla }_{x_t}\log p_d^{(t)}(x_t)\right)\frac{\partial x_t}{\partial \theta }\right]\mathrm{d}{x}_t,$$

$$\mathcal{L}_{\mathrm{DSM}}(\psi^{\prime})=\int p_{\theta}(x_{0})\,q(x_{t}|x_{0})\,\bigl{\|}\bm{s}_{\psi^{\prime}}(x_{t};t)-\nabla_{x_{t}}\log q(x_{t}|x_{0})\bigr{\|}_{2}^{2}\,\mathrm{d}x_{t}\,\mathrm{d}x_{0}.$$

$$\mathrm{KL}(p_{\theta }^{(t)}(x)||p_d^{(t)}(x))$$

$$\mathrm{DiU}(\theta ,\varphi )\equiv {\int }_0^{1}w(t){\mathrm{U}}^{(t)}(\theta ,\varphi )\mathrm{d}t\ge {\int }_0^{1}w(t)\mathrm{KL}(p_{\theta }^{(t)}||p_d^{(t)})\mathrm{d}t\equiv \mathrm{DiKL}(p_{\theta }||p_d),$$

$${\mathrm{U}}^{(t)}(\theta ,\varphi )=-\mathbb{H}(p_{\theta }^{(t)}(x_t,z))-\iint p_{\theta }^{(t)}(x_t,z)\left(\log p_d^{(t)}(x_t)\mathrm{d}{x}_t-\log q_{\varphi }^{(t)}(z|x_t)\right)\mathrm{d}z\mathrm{d}{x}_t.$$

$${\nabla }_{\theta }\int p_{\theta }^{(t)}(x_t)\log p_d^{(t)}(x_t)\mathrm{d}{x}_t={\nabla }_{\theta }\int p_{\theta }^{(t)}(x_t)\left(x_t^{\top }[{\nabla }_{x_t}\log p_d^{(t)}(x_t){]}_{\text{sg}}\right)\mathrm{d}{x}_t,$$

$${\mathrm{U}}^{(t)}(\theta ,\varphi )\doteq -\iint p_{\theta }^{(t)}(x_t,z)\left(x_t^{\top }[{\nabla }_{x_t}\log p_d^{(t)}(x_t){]}_{\text{sg}}+\log q_{\varphi }^{(t)}(z|x_t)\right)\mathrm{d}z\mathrm{d}{x}_t,$$

$$q_{\varphi }^{(t)}(z|x_t)=\int \delta \left(z-f_{\varphi }(a,x_t,t)\right)r_{\varphi }(a|x_t;t)\mathrm{d}a,$$

$$r_{\varphi }(a|x_t;t)=\mathcal{N}\left(a;{\mu }_{\varphi }(x_t;t),{\sigma }_t^2\mathrm{diag}({\sigma }_{\varphi }^2(x_t;t))\right).$$

$$\log q_{\varphi }^{(t)}(z=f_{\varphi }(a;t)|x_t;t)=\log r_{\varphi }(a|x_t;t)-\log \left|\det \frac{\partial f_{\varphi }}{\partial a}\right|,$$

$${\mathrm{U}}^{(t)}(\theta ,\varphi )=-\int p_{\theta }^{(t)}(x_t)\log p_d^{(t)}(x_t)\mathrm{d}{x}_t-\iint p_{\theta }^{(t)}(x_t,z)\log q_{\varphi }^{(t)}(z|x_t)\mathrm{d}z\mathrm{d}{x}_t.$$

$$\mathrm{KL}(p_{\theta }^{(t)}\Vert p_d^{(t)})=-\int p_{\theta }^{(t)}(x_t)\log p_d^{(t)}(x_t)\mathrm{d}{x}_t+\int p_{\theta }^{(t)}(x_t)\log p_{\theta }^{(t)}(x_t)\mathrm{d}{x}_t.$$

$$\underset{\theta }{\max }\mathbb{H}(p_{\theta }^{(t)}(x_t))=\underset{\theta }{\min }\mathbb{H}(p_{\theta }^{(t)}(z|x_t))\le \underset{\theta ,\varphi }{\min }-\iint p_{\theta }^{(t)}(x_t,z)\log q_{\varphi }^{(t)}(z|x_t)\mathrm{d}z\mathrm{d}{x}_t,$$

$$\mathrm{KL}\left(p_{\theta }^{(t)}(x_t)\Vert p_d^{(t)}(x_t)\right)$$

$$\mathrm{KL}$$

$$=\iint p_{\theta }^{(t)}(x_t\mid z)p(z)\log \frac{p_{\theta }^{(t)}(x_t\mid z)p(z)}{p_d^{(t)}(x_t)q_{\varphi }^{(t)}(z\mid x_t)}\mathrm{d}{x}_t\mathrm{d}z$$

$$\displaystyle=\iint p^{(t)}_{\theta}(x_{t}\mid z)p(z)\,\big{[}\log p^{(t)}_{\theta}(x_{t}\mid z)-\log p^{(t)}_{d}(x_{t})\big{]}\,\mathrm{d}x_{t}\,\mathrm{d}z$$

$$\displaystyle\qquad\qquad\qquad\qquad\qquad+\iint p^{(t)}_{\theta}(x_{t}\mid z)p(z)\,\big{[}\log p(z)-\log q^{(t)}_{\phi}(z\mid x_{t})\big{]}\,\mathrm{d}x_{t}\,\mathrm{d}z$$

$$=\int p_{\theta }^{(t)}(x_t)\log \frac{p_{\theta }^{(t)}(x_t)}{p_d^{(t)}(x_t)}\mathrm{d}{x}_t+\iint p_{\theta }^{(t)}(x_t\mid z)p(z)\log \frac{p_{\theta }^{(t)}(z\mid x_t)}{q_{\varphi }^{(t)}(z\mid x_t)}\mathrm{d}{x}_t\mathrm{d}z$$

$$=\mathrm{KL}\left(p_{\theta }^{(t)}(x_t)\Vert p_d^{(t)}(x_t)\right)+{\mathbb{E}}_{x_t\sim p_{\theta }^{(t)}(x_t)}\left[\mathrm{KL}\left(p_{\theta }^{(t)}(z\mid x_t)\Vert q_{\varphi }^{(t)}(z\mid x_t)\right)\right].$$

$$\mathrm{KL}\left(p_{\theta }^{(t)}(x_t)\Vert p_d^{(t)}(x_t)\right)$$

$${\nabla }_{\theta }\int p_{\theta }^{(t)}(x_t)\log p_d^{(t)}(x_t)\mathrm{d}{x}_t={\nabla }_{\theta }\int p_{\theta }^{(t)}(x_t)\left(x_t^{\top }[{\nabla }_{x_t}\log p_d^{(t)}(x_t){]}_{\text{sg}}\right)\mathrm{d}{x}_t,$$

$$x_t={\mathcal{F}}_{\theta }(g_{\theta }(z),\mathbf{ϵ},t)=g_{\theta }(z)+\mathbf{ϵ}{\sigma }_t,$$

$$\iint p_{\theta }^{(t)}(x_t|z)p(z)\log p_d^{(t)}(x_t)\mathrm{d}{x}_t\mathrm{d}z$$

$$\displaystyle=\iint p(z)p_{\bm{\epsilon}}(\bm{\epsilon})\log p^{(t)}_{d}\Bigl{(}\mathcal{F}_{\theta}(g_{\theta}(z),\bm{\epsilon},t)\Bigr{)}\mathrm{d}\bm{\epsilon}\mathrm{d}z=\mathbb{E}_{z\sim p(z),\,\bm{\epsilon}\sim p_{\bm{\epsilon}}(\bm{\epsilon})}\Bigl{[}\log p^{(t)}_{d}\Bigl{(}\mathcal{F}_{\theta}(g_{\theta}(z),\bm{\epsilon},t)\Bigr{)}\Bigr{]}$$

$$\nabla_{{\theta}}\iint p^{(t)}_{\theta}(x_{t}|z)p(z)\log p^{(t)}_{d}(x_{t})\mathrm{d}z\mathrm{d}x_{t}=\mathbb{E}_{z\sim p(z),\,\bm{\epsilon}\sim p_{\bm{\epsilon}}(\bm{\epsilon})}\Biggl{[}\nabla_{\theta}\log p^{(t)}_{d}\bigl{(}\mathcal{F}_{\theta}(g_{\theta}(z),\bm{\epsilon},t)\bigr{)}\Biggr{]}$$

$$\displaystyle\nabla_{\theta}\log p^{(t)}_{d}\bigl{(}\mathcal{F}_{\theta}(g_{\theta}(z),\bm{\epsilon},t)\bigr{)}=\nabla_{x_{t}}\log p^{(t)}_{d}(x_{t})\bigg{|}_{x_{t}=\mathcal{F}_{\theta}(g_{\theta}(z),\bm{\epsilon},t)}\cdot\frac{\partial\mathcal{F}_{\theta}(g_{\theta}(z),\bm{\epsilon},t)}{\partial\theta}$$

Full text

1 Introduction

Diffusion models sohl2015deep ; ho2020denoising ; song2019generative ; songscore have achieved remarkable success in generating high-quality samples. However, sampling typically requires a large number of denoising steps, making the process computationally inefficient in practice. To address this limitation, many acceleration methods ho2020denoising ; nichol2021improved ; song2020denoising ; liu2022pseudo ; lu2022dpm ; bao2022analytic ; bao2022estimating ; ou2024improving ; debortoli2025distributionaldiffusionmodelsscoring ; xiao2021tackling ; yu2024hierarchical have been proposed to reduce the number of function evaluations (NFEs) in pre-trained diffusion models, lowering the cost from thousands of steps to only tens. More recently, distillation-based approaches have further compressed the sampling process to as few as a single step while largely preserving generation quality zhou2024score ; salimansprogressive ; berthelot2023tract ; song2023consistency ; heek2024multistep ; kim2023consistency ; li2024bidirectional ; luo2023diff ; salimans2024multistep ; xie2024distillation ; zhang2025towards .

In this paper, we investigate a family of distillation methods based on divergence minimisation luo2023diff ; zhou2024adversarial ; zhang2025towards ; zhou2024score . These methods compress the full denoising process of a pre-trained teacher diffusion model into a one-step latent variable model by minimising a diffusive divergence (e.g. Diffusive Reverse KL Divergence he2024training ; zhang2025towards ; luo2023diff ) between the student and teacher through their respective score estimations. Although these approaches have demonstrated promising results, their gradients are approximated using the student’s score function, which is learned through denoising score matching (DSM). Because DSM training is inherently imperfect, the resulting gradient estimates are inevitably biased, often leading to sub-optimal generation performance.

In this study, we introduce Variational Diffusive Upper Bound Distillation (VarDiU), a method that formulates a variational upper bound of the diffusive divergence whose gradient can be estimated unbiasedly. We demonstrate that this unbiased gradient estimation leads to more stable and efficient training, achieving better performance compared to the original diffusive divergence, which relies on biased gradient approximations.

2 Background on Diffusion Models

Let $p_{d}(x_{0})$ denote the data distribution. Diffusion models define a forward noising process as $q(x_{0:T})=p_{d}(x_{0})\prod_{t=1}^{T}q(x_{t}|x_{t-1})$ where the transition kernel is a fixed Gaussian distribution $q(x_{t}|x_{t-1})=\mathcal{N}(x_{t}|\sqrt{1-\beta_{t}}x_{t-1},\beta_{t}I).$ This process admits the closed-form distribution,

[TABLE]

with $\bar{\alpha}_{t}=\prod_{s=1}^{t}(1-\beta_{s})$. As $T\to\infty$, the final state $x_{T}$ converges to a standard Gaussian distribution, i.e., $q(x_{T})\to\mathcal{N}(0,I)$. The generative process also uses a sequence of Gaussian distributions:

[TABLE]

where the covariance $\Sigma_{t-1}(x_{t})$ may either be learned nichol2021improved ; ou2024improving ; bao2022estimating or fixed bao2022analytic ; ho2020denoising . The mean is estimated via Tweedie’s formula efron2011tweedie ; robbins1992empirical :

[TABLE]

where the score function $\nabla_{x_{t}}\log p_{d}(x_{t})$ is approximated with a network $\bm{s}_{\psi}(x_{t};t)$ which is learned through denoising score matching (DSM) vincent2011connection ; song2019generative :

[TABLE]

The classical diffusion models require thousands of time steps to generate high-quality samples, making them computationally expensive. In the following section, we discuss a family of methods that can distill multi-step diffusion processes into one-step generative models.

2.1 Score-Based Distillation via Diffusive Reverse KL Minimisation

Score-based distillation methods aim to compress a teacher diffusion model, which we also denote as $p_{d}$, into a one-step implicit generative model goodfellow2014generative ; huszar2017variational ; zhang2020spread , defined as

[TABLE]

where $\delta(\cdot)$ is the Dirac delta function, $p(z)$ is a standard Gaussian prior, and $g_{\theta}$ is a deterministic neural network that generates samples in a single step. The training objective is formulated as the Diffusive KL divergence (DiKL) zhang2020spread ; wang2024prolificdreamer ; luo2023diff ; he2024training ; zhang2025towards , which measures the discrepancy between the student and teacher distributions across all intermediate noise levels:

[TABLE]

where $\omega(t)$ is a positive weighting function. Here, $p_{\theta}^{(t)}(x_{t})=\int q(x_{t}|x_{0})\,p_{\theta}(x_{0})\,dx_{0},$ and $p_{d}^{(t)}(x_{t})$ is the marginal distribution induced by the pre-trained teacher diffusion model under the same noise kernel $q(x_{t}|x_{0})$. This divergence is well defined even when the underlying distributions have disjoint support or don’t even allow density functions zhang2020spread . However, direct estimation of this divergence is intractable. One can derive the gradient w.r.t. ${\theta}$ at time $t$ as

[TABLE]

see A.2 in (luo2023diff, ) for a derivation. The teacher score $\nabla_{x_{t}}\log p^{(t)}_{d}(x_{t})$ is available from the pre-trained diffusion model. However, the student score $\nabla_{x_{t}}\log p^{(t)}_{\theta}(x_{t})$ is unknown. Fortunately, since we can sample from the student distribution, it can be estimated using denoising score matching (DSM):

[TABLE]

This method is known as VSD wang2024prolificdreamer or Diff-Instruct luo2023diff in the context of diffusion distillation. In practice, DSM approximations are imperfect due to the limited capacity of the score network $\bm{s}_{\psi^{\prime}}$. As a result, the direct gradient estimation of DiKL is biased, which leads to suboptimal training of the distilled student model. In the next section, we introduce a variational upper bound on the reverse KL divergence that enables unbiased gradient estimation for diffusion distillation.

3 Variational Diffusive Upper Bound Distillation

Instead of relying on optimisation-based estimation of the DiKL objective, we propose a variational upper bound whose gradient can be estimated unbiasedly. Specifically, we propose to extend the RKL upper bound of (zhang2019variational, ) into the diffusion space,

[TABLE]

The inequality follows by Jensen’s inequality; and the equality can be attained when the variational posterior is equal to the true posterior $q_{\phi}(z|x_{t};t)=p_{\theta}(z|x_{t})\propto p_{\theta}(x_{t}|z)p(z)$ (see Appendix˜A for a proof). We are then ready to define the diffusive upper bound as

[TABLE]

where the bound is tight when $q^{(t)}_{\phi}(z|x_{t})=p^{(t)}_{\theta}(z|x_{t})\propto p^{(t)}_{\theta}(x_{t}|z)p(z)$ for all $t$ almost everywhere.

3.1 Tractable Upper Bound Estimation

We then discuss how to estimate this upper bound. Writing the model joint as $p^{(t)}_{\theta}(x_{t},z)=p^{(t)}_{\theta}(x_{t}|z)p(z)$, the upper bound at time step $t$ can be written as

[TABLE]

where $\mathbb{H}(p^{(t)}_{\theta}(x_{t},z))$ denotes the joint entropy, i.e. $\mathbb{H}(p)=-\int p(x)\log p(x)\mathrm{d}x$. In the original reverse KL upper bound in (zhang2019variational, ), the joint entropy is trained with the reparametrisation trick. In our case, when the student model is an implicit model, we present the following surprising fact that the joint entropy is a constant, which removes the requirement of joint entropy estimation.

Proposition 3.1**.**

Let $p_{\theta}(x,z)=p_{\theta}(x|z)p(z)$ with a fixed prior $p(z)$ and p_{\theta}(x|z)=\mathcal{N}\bigl{(}x;\,g_{\theta}(z),\sigma^{2}{I}\bigr{)}. Then the joint entropy $\mathbb{H}(p_{\theta}(x,z))$ is a constant and independent of $\theta$.

See Appendix˜D for a proof. This should be distinguished from the marginal entropy $\mathbb{H}(p^{(t)}_{\theta}(x_{t}))$, which still depends on $\theta$, also see Section 3.3 for a deeper discussion. This observation greatly simplifies training by eliminating the need for entropy estimation.

The second term of Equation 9 cannot be directly optimised via automatic differentiation because the teacher density $p^{(t)}_{d}(x_{t})$ is unknown. However, in typical distillation settings, the score function $\nabla_{x_{t}}\log p^{(t)}_{d}(x_{t})$ is assumed to be available. Leveraging this, we propose an alternative loss whose gradient is equivalent and can be computed via the score function:

[TABLE]

where $[\,\cdot\,]_{\mathrm{sg}}$ denotes the stop-gradient operator. The proof involves using the reparametrisation trick and chain rule; see Appendix˜B for details. The final loss objective at time $t$ becomes

[TABLE]

where $\doteq$ denotes equivalence up to a constant. Therefore, when using a simple Gaussian variational family $q^{(t)}_{\phi}(z|x_{t})$, this bound can be unbiasedly estimated. A loss used practically can be found in Appendix˜C. We also propose a novel noise schedule and variance reduction techniques for training with this objective, see Section˜F.2 for details.

3.2 Improving Variational Inference with Normalising Flows

The upper bound becomes tight, i.e., equal to the DiKL objective, when the variational posterior matches the true posterior: $q^{(t)}_{\phi}(z|x_{t})=p^{(t)}_{\theta}(z|x_{t})$. Simple Gaussian variational distributions often struggle to approximate complex posteriors, particularly in cases where the true posterior is multi-modal or exhibits strong correlations. To address this limitation, we could use normalising flow  tabak2010density ; tabak2013family ; dinh2014nice ; rezende2015variational ; zhai2024normalizing ; durkan2019neural ; kobyzev2020normalizing to increase the flexibility of the posterior  rezende2015variational . Specifically, we define:

[TABLE]

where $f$ is an invertible function and the base distribution is defined as a Gaussian indexed by $t$.

[TABLE]

The time-dependent scaling factor $\sigma_{t}$ modulates the entropy of the base distribution to match the noise level at diffusion step $t$. The corresponding log-density of $q^{(t)}_{\phi}(z|x_{t})$ can be computed using the change-of-variables formula:

[TABLE]

We can then substitute this log-density into Equation (12) to obtain a tractable objective for distillation.

3.3 An Information Maximisation View of the Variational Diffusive Upper Bound

In the previous section, we showed that the joint entropy term in Equation (12) is constant with respect to $\theta$. As a result, the variational upper bound at diffusion step $t$ simplifies to:

[TABLE]

Now consider the original DiKL divergence at time step $t$, defined as:

[TABLE]

We observe that the first term in both objectives (Equations (16) and (17)) is identical. The second term in Equation (17) corresponds to the negative entropy of the marginal distribution, i.e., $-\mathbb{H}(p^{(t)}_{\theta}(x_{t}))$, which depends on $\theta$. We can then obtain an alternative perspective for justifying the proposed upper bound, formalised in the following theorem:

Theorem 3.2**.**

By the chain rule of entropy, we have: $\mathbb{H}(p^{(t)}_{\theta}(x_{t},z))=\mathbb{H}(p^{(t)}_{\theta}(z|x_{t}))+\mathbb{H}(p^{(t)}_{\theta}(x_{t}))$. Since $\mathbb{H}(p^{(t)}_{\theta}(x_{t},z))$ is constant with respect to $\theta$, it follows that:

[TABLE]

where $q^{(t)}_{\phi}(z|x_{t})$ is a variational approximation to the true posterior $p^{(t)}_{\theta}(z|x_{t})$.

See Appendix˜E for a proof. This result mirrors the variational approach in the Information Maximisation (IM) algorithm (barber2004algorithm, ), which maximises mutual information by minimising the conditional entropy using a variational decoder. Hence, this entropy bound can be viewed as a continuous latent-variable adaptation of the IM algorithm, where entropy maximisation serves to promote informative and diverse model outputs.

4 Experiments

To demonstrate the effectiveness of our approach, we consider a toy 2d mixture of 40 Gaussians (MoG-40) proposed in midgleyflow , where the means are uniformly distributed over $[-40,40]^{2}$. We train VarDiU with two types of posteriors: (a) a Gaussian posterior with learnable mean and variance (VarDiU-Gaussian), and (b) a Neural Spline Flow (NSF; durkan2019neural ) with a learnable-parameters Gaussian base distribution (VarDiU-NSF). We compare VarDiU to Diff-Instruct (luo2023diff, ) using varying numbers of student score training steps $(1,5,10)$. Experimental details can be found in Appendix˜F.

Settings We consider three practical settings: (1) access to the true analytical score; (2) access to the training dataset; (3) access to a teacher score pre-trained by a diffusion model. For the given training data case, we use the empirical score, defined by $\nabla_{x_{t}}\log\hat{p}_{d}(x_{t})=\frac{1}{N}\sum_{n=1}^{N}\nabla_{x_{t}}\log q(x_{t}|x_{0}^{(n)})$ with the given dataset $\mathcal{D}=\{x_{0}^{(n)}\}^{N}_{n=1}$, where $q(x_{t}|x_{0})$ matches the diffusion schedule of the VarDiU. This empirical score is a consistent estimator of the true data score song2019generative ; song2020score ; zhang2019variational , see Section˜F.3 for details. The learned score estimator is obtained by training EDM (karras2022edm, ), which is a popular class of diffusion models, on the observed data. For each setting, we report results over 10 independent runs. The noise schedule and weighting strategy are detailed in Section˜F.2.

Metrics We report the log-density of generated samples under the true data distribution, which reflects sample quality but not sample diversity. We additionally compute the Maximum Mean Discrepancy (MMD; gretton2012kernel ) with five kernels which can measure diversity, see for details in Section˜F.1.

Sample Quality We compare samples generated by one-step models trained with VarDiU and Diff-Instruct across three practical settings: using the true score, the training dataset, and a learned score (see Figure˜1 for visualizations). For Diff-Instruct, we train the model with 10 score-matching steps per generator gradient step to obtain the best sample quality we can, whereas the original Diff-Instruct paper luo2023diff uses only a single step. Both VarDiU-Gaussian and VarDiU-NSF produce cleaner and sharper samples than Diff-Instruct when using the true score or the training dataset. With a learned score, however, all methods yield different sub-optimal results due to biases in the provided scores. We also compare log-density and log-MMD across the samples (see Tables˜1(a), 1(c) and 1(b)). We can find VarDiU consistently performs best under reliable evaluations (true score and training data), demonstrating high accuracy and stability with tight confidence intervals, while Diff-Instruct lags despite larger steps. Under learned scores, evaluations are noisier: VarDiU achieves higher log-density but slightly worse MMD, making comparisons unreliable due to bias in the given pre-trained scores. Overall, VarDiU is more accurate and stable across settings.

Training Stability and Efficiency We present plots of log-MMD against both generator gradient iterations and total training time in Figure˜2; additional results are provided in Appendix Figures˜4, 5 and 6. Each curve shows the mean of 10 independent runs, with shaded regions indicating one standard deviation. For VarDiU, we apply a noise scheme with an annealing schedule. For Diff-Instruct, we use both the same annealing schedule (denoted as Diff-Instruct-a) and a non-annealing variant (see Section˜F.2 for details). In the log-MMD vs. generator gradient step plot, we find that under the same generator gradient budget (1M steps), VarDiU methods significantly outperform Diff-Instruct in both MMD values and stability (smaller variance). Moreover, incorporating a normalizing flow further improves performance compared to the Gaussian posterior, highlighting the importance of accurate generator gradient estimation. In the log-MMD vs. training time plot, we observe that under the same training hours, VarDiU-Gaussian clearly outperforms all other methods.

5 Conclusion and Future Work

In this paper, we introduced a novel method for diffusion distillation by proposing a variational diffusive upper bound on the DiKL divergence whose gradient can be estimated without bias. Our bound extends the Reverse KL upper bound from (zhang2019variational, ) to the diffusion setting and further eliminates the need to estimate the joint entropy. We also adapt this bound naturally to scenarios where only the teacher model’s score function is accessible. Additionally, we establish a connection between our method and the information maximisation (IM barber2004algorithm ) framework, see Appendix G for a discussion on related works. Empirically, we demonstrate on a toy problem that our method achieves improved one-step generation quality while significantly reducing training time.

As future work, we are extending this approach to higher-dimensional data, such as image or video generation, which we believe to be a promising direction.

Appendix A Proof of Variational Upper Bound

Proposition A.1** (Reverse KL Variational Upper Bound (zhang2019variational, )).**

For any choice of auxiliary variational distribution $q^{(t)}_{\phi}(z\mid x_{t})$,

[TABLE]

Moreover, equality holds iff $q^{(t)}_{\phi}(z\mid x_{t})=p^{(t)}_{\theta}(z\mid x_{t})$ for $p^{(t)}_{\theta}(x_{t})$-almost every $x_{t}$.

Proof.

We start from the joint KL:

[TABLE]

Since the conditional KL is non-negative, it follows that

[TABLE]

Equality holds iff $q^{(t)}_{\phi}(z\mid x_{t})=p^{(t)}_{\theta}(z\mid x_{t})$ for $p^{(t)}_{\theta}(x_{t})$ almost every $x_{t}$. ∎

Appendix B Proof of VarDiU Gradient Estimator

Proposition B.1**.**

The gradient can be estimated by

[TABLE]

where $[\,\cdot\,]_{\mathrm{sg}}$ denotes the stop-gradient operator.

Proof.

Since $x_{t}$ is generated via a differentiable transformation

[TABLE]

where $\bm{\epsilon}$ is drawn from a noise distribution $p_{\bm{\epsilon}}(\bm{\epsilon})$ that is independent of $\theta$. Then we have

[TABLE]

Since $\mathcal{F}_{\theta}(g_{\theta}(z),\bm{\epsilon},t)$ is differentiable in $\theta$ and under appropriate regularity conditions, we have

[TABLE]

Using the chain rule, the gradient inside the expectation becomes

[TABLE]

Substituting back into our expression for the gradient, we obtain

[TABLE]

as required. ∎

Appendix C A Practically Equivalent Variational Diffusive Upper Bound Loss

Tweedie’s formula efron2011tweedie ; robbins1992empirical .

For the Gaussian corruption kernel $p_{t}(x_{t}\mid x)=\mathcal{N}(x_{t};x,\sigma_{t}^{2}\mathbf{I})$, Tweedie’s identity gives

[TABLE]

We have from Equation˜12

[TABLE]

Proof.

We are given the VarDiU objective:

[TABLE]

With the stop-gradient applied to the score term,

[TABLE]

Write $x_{t}=x+\sigma_{t}\varepsilon$ with $\varepsilon\sim\mathcal{N}(0,\mathbf{I})$, and note that \big{[}\mu(x_{t};t)-x_{t}\big{]}_{\mathrm{sg}} is treated as a constant w.r.t. $(\phi,\theta)$. Taking the expectation over $x_{t}\mid x$ of (36) yields

[TABLE]

The first term on the right is exactly x^{\top}\big{[}\big{(}\mu(x_{t};t)-x\big{)}/\sigma_{t}^{2}\big{]}_{\mathrm{sg}}. The second term decomposes as

[TABLE]

Because $(x_{t}-x)=\sigma_{t}\varepsilon$ with $\varepsilon$ independent of $(\phi,\theta)$ and \big{[}\mu(\cdot;t)\big{]}_{\mathrm{sg}} blocks all gradients through its argument, this entire bracket has zero gradient w.r.t. $(\phi,\theta)$; in particular, \mathbb{E}_{x_{t}\mid x}\!\big{[}\|x_{t}-x\|^{2}\big{]}=d\,\sigma_{t}^{2} (with $d$ the data dimension) and \mathbb{E}_{x_{t}\mid x}\!\big{[}(x_{t}-x)^{\top}[\mu(x_{t};t)]_{\mathrm{sg}}\big{]} does not contribute to parameter gradients under reparametrisation. Hence it can be absorbed into “$\mathrm{const}$”.

Combining, we obtain (up to a parameter-independent constant)

[TABLE]

Substituting $x={g}_{\theta}(z)$ and keeping the $\log q^{(t)}_{\phi}(z\mid x_{t};t)$ term unchanged gives

[TABLE]

as required. ∎

Appendix D Joint Entropy is Constant

Proposition D.1**.**

Let

[TABLE]

with a fixed prior $p(z)$ and conditional \mathcal{N}\bigl{(}x;\,g_{\theta}(z),\sigma^{2}_{t}{I}\bigr{)}. Then the joint entropy $\mathbb{H}(p_{\theta}^{(t)}(x_{t},z))$ is independent of $\theta$.

Proof.

First, write the joint entropy as the sum of marginal and conditional entropies:

[TABLE]

where

[TABLE]

By assumption, the prior $p(z)$ does not depend on $\theta$, so $\mathbb{H}(p(z))$ is a fixed number.

For each $z$, $x_{t}|z\sim\mathcal{N}(g_{\theta}(z),\sigma_{t}^{2}I)$. The entropy of a Gaussian depends only on its covariance:

[TABLE]

where $d$ is independent of $\theta$, Here $\Sigma=\sigma_{t}^{2}I$, so

[TABLE]

which is independent of $g_{\theta}(z)$. Hence

[TABLE]

also constant. Since both $\mathbb{H}(p(z))$ and $\mathbb{H}(p_{\theta}^{(t)}(x_{t}|z))$ are independent of $\theta$, their sum $\mathbb{H}(p_{\theta}^{(t)}(x_{t},z))$ is a constant. ∎

Appendix E Proof of Theorem˜3.2

Theorem E.1** (Restatement of Theorem˜3.2).**

By the chain rule of entropy, we have: $\mathbb{H}(p^{(t)}_{\theta}(x_{t},z))=\mathbb{H}(p^{(t)}_{\theta}(z|x_{t}))+\mathbb{H}(p^{(t)}_{\theta}(x_{t}))$. Since $\mathbb{H}(p^{(t)}_{\theta}(x_{t},z))$ is constant with respect to $\theta$, it follows that:

[TABLE]

where $q^{(t)}_{\phi}(z|x_{t})$ is a variational approximation to the true posterior $p^{(t)}_{\theta}(z|x_{t})$.

Proof.

We have

[TABLE]

Therefore,

[TABLE]

and the bound is tight when $\mathrm{KL}(p_{\theta}(z|x_{t})||q^{(t)}_{\phi}(z|x_{t}))=0$. To minimise $\mathrm{KL}(p_{\theta}(z|x_{t})||q^{(t)}_{\phi}(z|x_{t}))$ wrt $\phi$, we have

[TABLE]

where the last line since Proposition D.1

[TABLE]

∎

Appendix F Experimental Details

In this section, we present the details of the experimental setting.

F.1 Implementation Details

For the generator $g_{\theta}$, we use a 5-layer MLP with latent dimension 2, hidden dimension 400 and SiLU activation. We use Adam to train the generator with learning rate $10^{-4}$, batch size $1024$, and gradient norm clip $10.0$.

We employ the same 5-layer MLP with hidden dimension 400, and SiLU activation as well to predict the mean and log-variance of the Gaussian posterior. We use Adam to train the generator and posterior network with learning rate $10^{-4}$, batch size $1024$, and gradient norm clip $10.0$.

For Diff-instruct luo2023diff , we use the same 5-layer MLP with hidden dimension 400, and SiLU activation to parametrise the student score network. We use Adam to train the student score network with learning rate $5*10^{-5}$, batch size $1024$, and gradient norm clip $10.0$.

Both methods are trained for $1{,}000{,}000$ epochs until convergence. For the learned score, we train an EDM (karras2022edm, ) on $10{,}000$ observed samples, while the empirical score is estimated from $10{,}000$ observed samples. All experiments are conducted on NVIDIA RTX 3090 GPUs.

Maximum Mean Discrepancy Implementation Details

For evaluation, we employ the Maximum Mean Discrepancy (MMD) (gretton2012kernel, ) with five Gaussian kernels of bandwidths $\{2^{-2},2^{-1},2^{0},2^{1},2^{2}\}$ for evaluation across multiple scales. Given two sets of samples $\{x_{i}\}_{i=1}^{n}\sim p_{d}$ and $\{y_{j}\}_{j=1}^{m}\sim p_{\theta}$, the squared MMD is defined as

[TABLE]

where $k(\cdot,\cdot)$ is the positive definite kernel. A smaller MMD indicates closer alignment between the generated and real distributions. For the results in Table 1, we perform 10 independent training runs. Evaluation metrics are recorded every $1{,}000$ epochs, and at each evaluation step, we compute the mean and standard deviation across runs. We use $10,000$ samples for MMD evaluation. To summarise performance, we report the average of these statistics over the last $50$ evaluations.

VarDiU Posterior Parametrisation

For VarDiU, we parametrise the Gaussian posterior as

[TABLE]

The time-dependent scaling factor $\sigma_{t}$ modulates the entropy of the base distribution to match the noise level at diffusion step $t$. This allows the posterior to capture fine-grained detail at small $t$, while maintaining flexibility to denoise highly corrupted samples at larger $t$.

For flow-based posterior, we implement Neural Splines Flows durkan2019neural using an open-source package normflows111https://github.com/VincentStimper/normalizing-flows Stimper2023 , we choose the flow length to be 4. We use Adam to train the generator and flow-based posterior network with learning rate $10^{-4}$, batch size $1024$, and gradient norm clip $10.0$ as well.

F.2 Noise Weighting and Scheduling

The one-step generation procedure is given by

[TABLE]

where the latent prior is sampled as $z\sim\mathcal{N}(\mathbf{0},\sigma^{2}_{\text{init}}\bm{I}).$ We set $\sigma^{2}_{\text{init}}=1$ in all toy data experiments.

For the forward diffusion distribution, we set

[TABLE]

Unlike Variance Preserving (VP;song2019generative ; ho2020denoising ) schedule which must satisfy $\alpha_{t}^{2}+\sigma_{t}^{2}=1$, we employ the schedule used in EDM karras2022edm which is a popular class of diffusion models, where we set $\alpha_{t}=1$ for all $t$ and we sample the time variable $t\sim\mathrm{Unif}[0,1]$ and define the corresponding noise levels as

[TABLE]

A visualisation of the resulting density of $\sigma_{t}$ under different values of $\rho$ is provided in Figure˜3. For Diff-Instruct, we fix $\rho=1.5$ across all three score estimation settings. We found an annealing $\rho$ does not help with the training dynamics of Diff-Instruct, so we use a fixed schedule.

Annealing $\rho$.

In VarDiU, we adopt an annealing strategy in which $\rho$ gradually increases with training iterations. Specifically, $\rho$ is initialised at $\rho_{\text{init}}=0.1$ and incremented by $0.01$ every $1000$ epochs until it reaches $\rho_{\text{end}}$.

The annealed noise schedule plays a crucial role in stabilising training. The central idea is to initially place more emphasis on larger noise levels ($\sigma$), where the learning problem is smoother and less sensitive to estimation errors. As training progresses, the schedule gradually shifts towards smaller $\sigma$ values, allowing the model to refine details and capture higher-frequency structures. This coarse-to-fine strategy not only mitigates instability at the early stage but also improves convergence and sample quality in later stages.

(1) Given The True Score.

We set $\sigma_{\min}=0.1$ and $\sigma_{\max}=20$ for both Diff-Instruct and VarDiU. For VarDiU, we set $\rho_{\text{end}}=5.0$.

(2) Given A Pre-trained DM.

For Diff-Instruct, we set $\sigma_{\min}=1.1$ and $\sigma_{\max}=40$. For VarDiU, we set $\sigma_{\min}=1.5$, $\sigma_{\max}=40$, and $\rho_{\text{end}}=2.0$.

(3) Given A Dataset.

We employ $10{,}000$ samples for empirical score estimation, with details provided in Section˜F.3. For both Diff-Instruct and VarDiU, we set $\sigma_{\min}=0.65$ and $\sigma_{\max}=40$, and for VarDiU, we additionally set $\rho_{\text{end}}=2.0$.

Noise Weighting.

For both Diff-Instruct and VarDiU, we adopt the weighting function

[TABLE]

The noise schedule aligns with zhang2025towards ; luo2023diff . We further divide the weighting by the $\sigma_{\max}$ for not changing the learning rate of training.

F.3 Empirical Score Estimation

Score of a Gaussian-smoothed empirical distribution.

Let $\{x_{0}^{(n)}\}_{n=1}^{N}\subset\mathbb{R}^{d}$ be data points, and define the Gaussian kernel

[TABLE]

The $\sigma$-smoothed empirical distribution is

[TABLE]

Lemma F.1** (Log-sum gradient identity).**

For positive functions $\{a_{n}(x)\}_{n=1}^{N}$,

[TABLE]

Proof.

By the chain and quotient rules,

[TABLE]

Applying the lemma with $a_{n}(x)=\phi_{\sigma}(x-x_{0}^{(n)})$ gives

[TABLE]

with responsibilities

[TABLE]

The Gaussian score is

[TABLE]

Therefore,

[TABLE]

Final form. For any query point $x_{t}$, the score of the Gaussian–KDE is

[TABLE]

We employ Equation˜57 to compute the empirical score using the given dataset practically. The accuracy of empirical score estimation tends to be small when a large number of samples are given, providing useful teacher guidance.

F.4 Symmetric Sampling

To reduce the variance of the stochastic estimator, we employ a symmetric sampling scheme botev2017variance ; kroese2013handbook for both VarDiU and Diff-Instruct.

We first draw a half-batch of latent codes,

[TABLE]

where $B$ is the batch size. We then duplicate the samples to construct a full batch,

[TABLE]

Passing these latents through the generator yields samples

[TABLE]

Next, we simulate forward diffusion process

[TABLE]

and extend them symmetrically as

[TABLE]

Thus, each noise vector is paired with its negative, ensuring that the perturbations are exactly centred. We also sample noise levels from the power distribution

[TABLE]

and replicate them symmetrically:

[TABLE]

Noisy inputs. The perturbed inputs are then constructed as

[TABLE]

where $\odot$ denotes element-wise multiplication.

We found symmetric sampling can reduce the variance during training, further enhance the training stability of VarDiU, and can help with stabilising the training with Diff-Instruct.

Appendix G Related Works

Diffusion acceleration has received significant attention as a means of reducing the number of steps required in the reverse process. At the core of this problem lies the posterior distribution $q(x_{0}\mid x_{t})$, which connects a noisy intermediate $x_{t}$ to the original data $x_{0}$. Accurate modelling of this posterior enables faithful reconstruction of $x_{0}$ in fewer steps, thereby lowering the sampling cost. Existing approaches can be broadly grouped into strategies such as designing alternative samplers lu2022dpm ; lu2025dpm ; song2020denoising ; de2021diffusion , relaxing the Gaussian posterior assumption debortoli2025distributionaldiffusionmodelsscoring ; lipmanflow ; liuflow , or projecting the dynamics onto lower-dimensional subspaces rombach2022high . Whilst these techniques substantially reduce sampling time, they still typically require a non-trivial number of steps to achieve high-quality generation.

Another important direction focuses on compressing the reverse diffusion process into shorter chains. This includes progressive distillation salimansprogressive ; meng2023distillation , consistency models and their improvements song2023consistency ; songimproved ; kim2023consistency , as well as extensions to latent diffusion luo2023latent ; rombach2022high . More recently, Inductive Moment Matching has been proposed as a single-stage framework that trains few-step models from scratch by matching distributional moments zhou2025inductive .

In parallel, several works have investigated single-step diffusion models. Some fine-tune pre-trained diffusion models by minimising divergences such as the Diffusive KL or Fisher divergence luo2023diff ; xie2024distillation ; zhou2024score ; xu2025one , while others incorporate adversarial losses to enhance perceptual fidelity yin2024improved ; zhou2024adversarial . Our work builds on this line by eliminating the need for a separate student score network, thereby avoiding bias from inaccurate score estimation and enabling more stable training with improved generation quality.

Finally, recent work has shown that teacher guidance is not strictly necessary for effective diffusion distillation, highlighting that the multi-scale representations learned by diffusion models are the crucial factor behind the success of one-step generators zhang2025towards .

Appendix H More Convergence and Efficiency Results

We present the log-MMD plots with respect to both generator gradient iterations and total training time, using the pre-trained DM and the given dataset in Figures˜5 and 6, respectively.

When relying on the pre-trained DM, both methods exhibit high variance and unstable curves due to inaccurate score predictions from the pre-trained model. The noticeable jump of the VarDiU variants in the middle of training arises from the annealing schedule, where the density concentrates on a range of $\sigma$ values for which the pre-trained DM suffers from large estimation errors. Nevertheless, the VarDiU variants maintain relatively low variance in the log-density curves, highlighting their robustness.

In contrast, when using the given dataset, we estimate the score empirically. The empirical score becomes increasingly accurate with larger sample sizes. As shown in Figure˜6, the VarDiU variants achieve performance comparable to that obtained with the true analytical score. Furthermore, VarDiU demonstrates substantially improved training stability and efficiency when employing a Gaussian-parameterised posterior.

Overall, these results highlight the superiority of VarDiU, which consistently achieves more stable and efficient convergence across different settings compared to the baseline.