Variational Diffusion Models Overview

• Variational Diffusion Models (VDMs) are generative models that combine variational inference with diffusion-based methods, using a fixed noising process and a learnable denoising process.
• They leverage a likelihood-based training objective via the ELBO, unifying score-matching and variational autoencoder perspectives through optimized noise schedules.
• Practical implementations of VDMs have advanced applications in image synthesis, density estimation, and inverse problems, achieving state-of-the-art performance.

Variational Diffusion Models (VDMs) define a class of generative models at the intersection of variational inference and diffusion-based probabilistic modeling. They operationalize deep hierarchical variational autoencoders using a fixed, tractable forward "noising" process and a learned reverse "denoising" process, forming the foundation of most modern score-based and denoising diffusion models. VDMs are distinguished by their capacity for likelihood-based learning, efficient noise-schedule optimization, and unified theoretical characterization, enabling both state-of-the-art synthesis and tractable maximum likelihood estimation across diverse data domains (Kingma et al., 2021, Luo, 2022, Ribeiro et al., 2024).

1. Probabilistic Formulation and Graphical Model Structure

VDMs are formulated as Markovian hierarchical latent-variable models. The forward (inference) chain is fixed as a Gaussian diffusion:

$$q(z_{1:T}|x) = q(z_T | x) \prod_{t=2}^T q(z_{t-1}|z_t, x)$$

where $q(z_T|x) = \mathcal{N}(0, I)$ and $q(z_{t-1}|z_t, x)$ are linear-Gaussian with parameters derived from the preceding variables and the data $x$. The generative model reverses this process:

$$p_\theta(x, z_{1:T}) = p(z_T) \prod_{t=2}^T p_\theta(z_{t-1}|z_t) \cdot p_\theta(x|z_1)$$

In the VDM paradigm, only the reverse kernels $p_\theta(z_{t-1}|z_t)$ are learned, while the forward kernels remain fixed. This structure can be viewed as a hierarchical variational autoencoder (HVAE) with infinitely many stochastic layers in the continuous-time limit (Kingma et al., 2021, Ribeiro et al., 2024, Luo, 2022).

2. Variational Objective and Loss Derivation

The central training objective is the evidence lower bound (ELBO):

$$\log p_\theta(x) \geq \mathbb{E}_{q(z_{1:T}|x)} [ \log p_\theta(x, z_{1:T}) - \log q(z_{1:T}|x) ]$$

For the diffusion hierarchy, this expands as:

$$\begin{aligned} \log p_\theta(x) \geq&\, \mathbb{E}_{q(z_1|x)}[\log p_\theta(x|z_1)] - \mathrm{KL}[q(z_T|x) \| p(z_T)] \ &- \sum_{t=2}^T \mathbb{E}_{q(z_t|x)} \left[ \mathrm{KL}[q(z_{t-1}|z_t,x)\|p_\theta(z_{t-1}|z_t)] \right] \end{aligned}$$

The per-step KL simplifies, via Gaussian conjugacy, to a quadratic loss between the posterior mean $\mu_Q$ and the predicted mean $\mu_\theta$, typically reducing to a "denoising" or "noise prediction" loss weighted by the change in signal-to-noise ratio (SNR):

$$D_{\mathrm{KL}}[q \,\|\, p_\theta] = \frac{1}{2\sigma_Q^2} \|\mu_Q - \mu_\theta\|^2 = \frac{1}{2} (SNR(s)-SNR(t)) \|x - \hat x_\theta(z_t, t)\|_2^2$$

where $SNR(t) = \alpha_t^2/\sigma_t^2$, and $$z_t = \alpha_t x + \sigma_t \epsilon,\, \epsilon\sim\mathcal{N}(0,I)$$ (Kingma et al., 2021, Ribeiro et al., 2024).

This loss admits three equivalent parameterizations: predicting $x$, predicting noise $\epsilon$, or predicting the score $\nabla_{z_t} \log p(z_t)$, unified via Tweedie's formula (Luo, 2022).

3. Continuous-Time Limit, Schedule Invariance, and Training

By taking the $T\to\infty$ limit, the discrete diffusion process converges to a continuous stochastic differential equation (SDE) and the loss becomes an integral over time:

$$L_\infty(x) = -\frac{1}{2} \mathbb{E}_{t\sim U[0,1],\,\epsilon} [ SNR'(t) \|x - \hat x_\theta(z_t,t)\|^2 ] + \text{const}$$

Equivalently, when the model is parameterized to predict the noise:

$$L_{\text{simple}} = \mathbb{E}_{t\sim U[0,1],\,\epsilon} [ \|\epsilon - \hat\epsilon_\theta(z_t,t)\|^2 ]$$

A crucial property is that the continuous-time ELBO is invariant to the specific shape of the noise schedule $\gamma(t)$ (or equivalently SNR$(t)$), up to its endpoint values. This underpins the possibility of learning or tuning the schedule during training without affecting the maximum-likelihood objective (Kingma et al., 2021, Ribeiro et al., 2024).

Joint optimization of the denoising network and the schedule improves convergence, as gradient signal is concentrated at challenging SNR regimes (Kingma et al., 2021).

4. Extensions: Conditional, Expressive Posteriors, and Schrödinger Diffusion

Conditional VDMs (CVDMs) generalize the framework to model conditional distributions $p(y|x)$ via a variance-preserving, data-conditioned forward SDE:

$$q(z_t|y,x) = \mathcal{N}(\sqrt{\gamma(t,x)}\,y, (1-\gamma(t,x))I)$$

Here, the schedule $\gamma(t, x)$ is itself learned, with explicit mechanisms (e.g., monotonic neural nets and regularization) to ensure smoothness and task-adaptive noise injection. Conditioning is incorporated both in the forward process and in the reverse denoiser, enabling high-quality solutions for inverse problems without schedule fine-tuning (Maggiora et al., 2023).

Expressive variational posteriors are realized in the denoising diffusion variational inference (DDVI) approach, replacing conventional VAEs' one-shot encoders with iterative diffusion-based posteriors. Training alternates between "wake" steps fitting a regularized ELBO and "sleep" steps enforcing mode-covering via forward KL, yielding tighter bounds and improved latent inference (Piriyakulkij et al., 2024).

For efficient distributional transport, variational Schrödinger Diffusion Models (VSDM) use a multivariate Ornstein–Uhlenbeck forward SDE, parameterized by time-dependent matrices $A_t$, enabling closed-form transitions and simulation-free training of the backward score via explicit score-matching. Stochastic approximation optimizes $A_t$, combining transport-theoretic efficiency with tractable variational losses (Deng et al., 2024).

5. Theoretical Insights and Unified Perspective

The ELBO in VDMs subsumes classical denoising and score-matching objectives. Weighted diffusion losses and alternative training objectives (e.g., "simple" noise-prediction loss, weighted score-matching) correspond to alternative forms of the same ELBO. This unifies the VAE, score-based, and diffusion-modeling viewpoints: the loss functional, model class, and sample generation are determined—modulo trivial linear rescalings—by the SNR endpoints and the variational denoiser (Luo, 2022, Kingma et al., 2021, Ribeiro et al., 2024).

Importantly, VDMs demonstrate that state-of-the-art sample quality and likelihood-based objectives are not inherently opposed: maximum-likelihood training yields both (Ribeiro et al., 2024). Schedule invariance justifies learned scheduling and data-dependent augmentation.

In conditional or learned-schedule variants, regularization (smoothness, physics-informed constraints) controls the convergence rate of the discrete to continuous ELBO, preventing the emergence of pathological noise schedules and ensuring robust optimization in practice (Maggiora et al., 2023).

6. Practical Implementation and Empirical Performance

VDMs have been deployed for density estimation, image synthesis, inverse problems, and bits-back compression. Architectural innovations include monotonic schedule parameterization, Fourier feature augmentation, deep U-Nets with attention at bottlenecks, and robust optimization protocols (Adam, EMA, dropout, low-discrepancy sampling for $t$) (Kingma et al., 2021).

Empirical results have established new likelihood benchmarks on CIFAR-10, ImageNet, and others, with performance (in bits/dim) surpassing regular autoregressive and normalizing flow models. Notably, VDMs achieve near-theoretical lossless compression rates using standard bits-back schemes (Kingma et al., 2021).

Conditional VDMs demonstrate state-of-the-art or competitive results on scientific inverse problems (super-resolution microscopy, Quantitative Phase Imaging, natural image super-resolution) without manual schedule tuning (Maggiora et al., 2023).

In the variational Schrödinger and DDVI settings, the methodology achieves sample efficiency, mode-coverage, and improved transport path straightness or latent inference, outperforming standard flows and adversarial approaches (Piriyakulkij et al., 2024, Deng et al., 2024).

7. Outlook and Open Directions

Key future research avenues for VDMs include:

• Designing non-Gaussian and discrete noising processes to widen the range of generative models (Ribeiro et al., 2024).
• Coupling diffusion models with structured or disentangled latent representations for enhanced interpretability and causal discovery (Ribeiro et al., 2024).
• Integrating optimal-transport and Schrödinger-bridge formulations for increased sample efficiency and principled path interpolation (Deng et al., 2024, Ribeiro et al., 2024).
• Developing data-driven and adaptive scheduling, uncertainty quantification, and new architectures exploiting SNR-invariant training (Maggiora et al., 2023, Kingma et al., 2021).
• Further bridging variational, optimal-transport, and score-based learning paradigms for unified generative modeling.

These directions are expected to expand the applicability, interpretability, and theoretical grounding of VDMs in both foundation-model and application-driven regimes.