DDIM (Denoising Diffusion Implicit Models): Non-Markovian Sampling and Probability-Flow ODEs

The Sampling Bottleneck in Standard Markovian Diffusion

Standard Denoising Diffusion Probabilistic Models (DDPM) (Ho et al., NeurIPS 2020) define the forward diffusion chain as a strictly Markovian Gaussian process:

$q({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1})=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I}\right)$

While the closed-form marginal allows direct training at arbitrary timesteps $t$:

$q({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I}\right),{\bar{\alpha }}_t={\prod }_{s=1}^t(1-{\beta }_s)$

the reverse generation process is tightly bound to sequential step-by-step Markovian ancestral transitions:

$p_{\theta }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t)=\mathcal{N}\left({\mathbf{x}}_{t-1};\frac{1}{\sqrt{{\alpha }_t}}\left({\mathbf{x}}_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)\right),{\tilde{\beta }}_t\mathbf{I}\right)$

Generating a single high-resolution image requires evaluating the neural noise predictor ${\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)$ across all $T=1000$ sequential timesteps. This constraint renders inference computationally prohibitive ($\ge 10\text{–}30\text{ seconds}$ per sample on modern GPUs).

Denoising Diffusion Implicit Models (DDIM) (Song et al., ICLR 2021) eliminates this limitation by generalizing the forward trajectory to a non-Markovian family of distributions that strictly preserves the identical marginal distributions $q({\mathbf{x}}_t\mid {\mathbf{x}}_0)$. This unlocks:

1. Sub-Sequence Step Skipping: Sampling converges in $20\text{–}50$ discretization steps without retraining the base ${\mathbf{ϵ}}_{\theta }$ model.
2. Deterministic Invertibility ($\eta =0$): The reverse process becomes a deterministic discretization of the continuous-time Probability-Flow Ordinary Differential Equation (ODE), establishing an exact bidirectional mapping between latent noise ${\mathbf{x}}_T$ and synthesized data ${\mathbf{x}}_0$.

Architectural & Sampler Comparison

Mathematical Foundations

Figure 1: Mathematical unification of DDIM showing marginal preservation, non-Markovian parameterization, and deterministic ODE discretization.

1. Non-Markovian Forward Distribution Formulation

DDIM constructs a family of non-Markovian inference distributions $q_{\sigma }({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)$ that maintain identical marginals $q_{\sigma }({\mathbf{x}}_t\mid {\mathbf{x}}_0)=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,(1-{\bar{\alpha }}_t)\mathbf{I}\right)$:

$q_{\sigma }({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0)=q_{\sigma }({\mathbf{x}}_T\mid {\mathbf{x}}_0){\prod }_{t=2}^T{q}_{\sigma }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)$

The conditional forward step $q_{\sigma }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)$ is parameterized as a Gaussian:

$q_{\sigma }({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}\left({\mathbf{x}}_{t-1};\sqrt{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}\frac{{\mathbf{x}}_t-\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0}{\sqrt{1-{\bar{\alpha }}_t}},{\sigma }_t^2\mathbf{I}\right)$

where the variance ${\sigma }_t^2$ is controlled by a hyperparameter $\eta \ge 0$:

${\sigma }_t=\eta \cdot \sqrt{\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\left(1-\frac{{\bar{\alpha }}_t}{{\bar{\alpha }}_{t-1}}\right)}=\eta \cdot {\tilde{\beta }}_t^{1/2}$

2. General Reverse Sampling Derivation

During reverse sampling, the clean data point ${\mathbf{x}}_0$ is unknown and approximated via Tweedie’s formula using the trained noise predictor ${\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)$:

${\hat{\mathbf{x}}}_0({\mathbf{x}}_t,t)=\frac{{\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)}{\sqrt{{\bar{\alpha }}_t}}$

Substituting ${\hat{\mathbf{x}}}_0$ into the forward conditional yields the Universal DDIM Update Rule:

${\mathbf{x}}_{t-1}=\underset{\text{Predicted Clean Sample}}{\underbrace{\sqrt{{\bar{\alpha }}_{t-1}}{\hat{\mathbf{x}}}_0({\mathbf{x}}_t,t)}}+\underset{\text{Direction Pointing to }{\mathbf{x}}_t}{\underbrace{\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)}}+\underset{\text{Random Noise}}{\underbrace{{\sigma }_t\mathbf{ϵ}}},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I})$

Special Regimes of the Stochasticity Parameter $\eta$:

1. $\eta =1$ (DDPM Ancestral Sampling): ${\sigma }_t={\tilde{\beta }}_t^{1/2}⟹\sqrt{1-{\bar{\alpha }}_{t-1}-{\sigma }_t^2}=\sqrt{\frac{{\alpha }_t(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}}$ The formulation reverts exactly to the stochastic Markovian DDPM transition kernel.

2. $\eta =0$ (Deterministic DDIM): ${\sigma }_t=0⟹{\mathbf{x}}_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}{\hat{\mathbf{x}}}_0({\mathbf{x}}_t,t)+\sqrt{1-{\bar{\alpha }}_{t-1}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)$ The reverse trajectory becomes completely deterministic, removing all injected random noise at each step.

3. The Probability-Flow ODE Derivation

To establish the continuous-time connection, consider an infinitesimal timestep transition $\Delta t\to 0$. Rewriting the deterministic DDIM update ($\eta =0$) in terms of continuous time variables $t\in [0,1]$:

${\mathbf{x}}_{t-\Delta t}=\sqrt{{\bar{\alpha }}_{t-\Delta t}}\left(\frac{{\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)}{\sqrt{{\bar{\alpha }}_t}}\right)+\sqrt{1-{\bar{\alpha }}_{t-\Delta t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)$

Rearranging into discrete difference quotient form:

$\frac{{\mathbf{x}}_{t-\Delta t}-{\mathbf{x}}_t}{\Delta t}=\left(\frac{\sqrt{{\bar{\alpha }}_{t-\Delta t}}-\sqrt{{\bar{\alpha }}_t}}{\Delta t\sqrt{{\bar{\alpha }}_t}}\right){\mathbf{x}}_t+\left(\frac{\sqrt{1-{\bar{\alpha }}_{t-\Delta t}}-\sqrt{\frac{{\bar{\alpha }}_{t-\Delta t}}{{\bar{\alpha }}_t}}\sqrt{1-{\bar{\alpha }}_t}}{\Delta t}\right){\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)$

Taking the continuous limit $\Delta t\to 0$ yields the continuous ODE:

$\frac{\mathrm{d}{\mathbf{x}}_t}{\mathrm{d}t}=\frac{1}{2}\frac{\mathrm{d}\log {\bar{\alpha }}_t}{\mathrm{d}t}{\mathbf{x}}_t-\frac{1}{2}\frac{\mathrm{d}{\bar{\alpha }}_t}{\mathrm{d}t}\frac{{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)}{\sqrt{1-{\bar{\alpha }}_t}}$

Using the score function identity ${\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)=-\frac{{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)}{\sqrt{1-{\bar{\alpha }}_t}}$, this simplifies to the Probability-Flow ODE:

$\mathrm{d}{\mathbf{x}}_t=\left[f(t){\mathbf{x}}_t-\frac{1}{2}{g}^2(t){\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)\right]\mathrm{d}t$

where $f(t)=\frac{1}{2}\frac{\mathrm{d}\log {\bar{\alpha }}_t}{\mathrm{d}t}$ is the drift coefficient and $g^2(t)=-\frac{\mathrm{d}{\bar{\alpha }}_t}{\mathrm{d}t}$ is the diffusion coefficient. This proves that deterministic DDIM is an Euler-Maruyama discretization of the exact Probability-Flow ODE.

4. Spherical Latent Interpolation ($SLERP$)

Because deterministic DDIM ($\eta =0$) forms a continuous bijective mapping $\Psi :{\mathbf{x}}_T\leftrightarrow {\mathbf{x}}_0$, interpolating between two generated images ${\mathbf{x}}_0^{(A)}$ and ${\mathbf{x}}_0^{(B)}$ is executed via Spherical Linear Interpolation ($SLERP$) on the initial standard Gaussian sphere:

${\mathbf{x}}_T(\gamma )=\frac{\sin ((1-\gamma )\theta )}{\sin \theta }{\mathbf{x}}_T^{(A)}+\frac{\sin (\gamma \theta )}{\sin \theta }{\mathbf{x}}_T^{(B)},\theta =\arccos \left(\frac{⟨{\mathbf{x}}_T^{(A)},{\mathbf{x}}_T^{(B)}⟩}{\Vert {\mathbf{x}}_T^{(A)}{\Vert }_2\Vert {\mathbf{x}}_T^{(B)}{\Vert }_2}\right)$

where $\gamma \in [0,1]$. Sampling ${\mathbf{x}}_0(\gamma )=\Psi ({\mathbf{x}}_T(\gamma ))$ generates semantically smooth interpolations without semantic collapse or structural artifacts.

Production PyTorch Implementation

Below is a complete, vectorized PyTorch implementation of the DDIM Scheduler, supporting both forward deterministic ODE inversion and reverse accelerated sampling with sub-sequence stride scheduling.

Step 1: Vectorized DDIM Sampler & Inversion Engine

from __future__ import annotations

import math
import torch
import torch.nn as nn


class DDIMScheduler:
    def __init__(
        self,
        num_train_timesteps: int = 1000,
        beta_start: float = 0.0001,
        beta_end: float = 0.02,
        beta_schedule: str = "linear",
    ) -> None:
        self.num_train_timesteps = num_train_timesteps
        
        if beta_schedule == "linear":
            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
        elif beta_schedule == "cosine":
            steps = num_train_timesteps + 1
            s = 0.008
            x = torch.linspace(0, num_train_timesteps, steps, dtype=torch.float32)
            alphas_cumprod = torch.cos(((x / num_train_timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
            alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
            self.betas = torch.clip(1.0 - alphas_cumprod[1:] / alphas_cumprod[:-1], 0.0001, 0.9999)
        else:
            raise NotImplementedError(f"Schedule {beta_schedule} is not supported.")

        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)

    def set_timesteps(self, num_inference_steps: int, device: torch.device) -> torch.Tensor:
        """Constructs an evenly spaced sub-sequence of timesteps."""
        self.num_inference_steps = num_inference_steps
        step_ratio = self.num_train_timesteps // self.num_inference_steps
        timesteps = (torch.arange(0, num_inference_steps) * step_ratio).round().flip(0).long().to(device)
        self.timesteps = timesteps
        return self.timesteps

    def step(
        self,
        model_output: torch.Tensor,
        timestep: int,
        prev_timestep: int,
        sample: torch.Tensor,
        eta: float = 0.0,
    ) -> torch.Tensor:
        """Executes a single DDIM step from timestep t to prev_timestep."""
        alpha_prod_t = self.alphas_cumprod[timestep]
        alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else torch.tensor(1.0)

        # 1. Predict clean x_0 from current noise estimate
        pred_original_sample = (sample - torch.sqrt(1.0 - alpha_prod_t) * model_output) / torch.sqrt(alpha_prod_t)

        # 2. Compute variance sigma_t
        variance = ((1.0 - alpha_prod_t_prev) / (1.0 - alpha_prod_t)) * (1.0 - alpha_prod_t / alpha_prod_t_prev)
        sigma_t = eta * torch.sqrt(variance.clamp(min=0.0))

        # 3. Compute direction pointing to x_t
        pred_sample_direction = torch.sqrt((1.0 - alpha_prod_t_prev - sigma_t**2).clamp(min=0.0)) * model_output

        # 4. Compute x_{t-1}
        prev_sample = torch.sqrt(alpha_prod_t_prev) * pred_original_sample + pred_sample_direction

        if eta > 0.0:
            noise = torch.randn_like(sample)
            prev_sample = prev_sample + sigma_t * noise

        return prev_sample

    def invert_step(
        self,
        model_output: torch.Tensor,
        timestep: int,
        next_timestep: int,
        sample: torch.Tensor,
    ) -> torch.Tensor:
        """Deterministic Probability-Flow ODE Inversion: steps forward from t to next_timestep."""
        alpha_prod_t = self.alphas_cumprod[timestep] if timestep >= 0 else torch.tensor(1.0)
        alpha_prod_t_next = self.alphas_cumprod[next_timestep]

        # 1. Estimate x_0
        pred_original_sample = (sample - torch.sqrt(1.0 - alpha_prod_t) * model_output) / torch.sqrt(alpha_prod_t)

        # 2. Step forward along deterministic trajectory
        pred_sample_direction = torch.sqrt(1.0 - alpha_prod_t_next) * model_output
        next_sample = torch.sqrt(alpha_prod_t_next) * pred_original_sample + pred_sample_direction

        return next_sample

Step 2: Full Inference Loop with Classifier-Free Guidance (CFG)

from __future__ import annotations

import torch
import torch.nn as nn
from ddim_scheduler import DDIMScheduler


@torch.no_grad()
def ddim_generate(
    unet: nn.Module,
    scheduler: DDIMScheduler,
    shape: tuple[int, int, int, int],
    text_embeddings: torch.Tensor,
    uncond_embeddings: torch.Tensor,
    guidance_scale: float = 7.5,
    num_inference_steps: int = 50,
    eta: float = 0.0,
    device: str = "cuda",
) -> torch.Tensor:
    dev = torch.device(device)
    timesteps = scheduler.set_timesteps(num_inference_steps, dev)
    
    # 1. Initialize random Gaussian latent
    latents = torch.randn(shape, device=dev, dtype=torch.float32)

    # 2. Reverse Diffusion Loop
    for i, t in enumerate(timesteps):
        prev_t = timesteps[i + 1] if i + 1 < len(timesteps) else -1
        
        # Dual-batch forward pass for Classifier-Free Guidance
        latent_model_input = torch.cat([latents] * 2)
        encoder_hidden_states = torch.cat([uncond_embeddings, text_embeddings])
        
        # Predict noise
        t_input = torch.tensor([t, t], device=dev)
        noise_pred = unet(latent_model_input, t_input, encoder_hidden_states=encoder_hidden_states)
        
        # CFG Extrapolation
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred_guided = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        # DDIM Step
        latents = scheduler.step(
            model_output=noise_pred_guided,
            timestep=t.item(),
            prev_timestep=prev_t.item() if isinstance(prev_t, torch.Tensor) else prev_t,
            sample=latents,
            eta=eta,
        )

    return latents

Empirical Benchmark Evaluation

Quantitative evaluation measuring Fréchet Inception Distance (FID) vs. Number of Function Evaluations (NFE) across unconditional and text-to-image datasets:

Troubleshooting Common Sampling Faults

1. High-Frequency Noise Leakage at Low Step Counts ($NFE<20$)

• Symptom: Generated images retain noticeable grain, washed-out textures, or high-frequency checkerboard artifacts.
• Root Cause: First-order Euler discretization introduces cumulative truncation errors when step size $\Delta t$ is large.
• Remedy: Increase num_inference_steps to $\ge 30$, or switch to a higher-order multi-step solver (e.g., DPMSolverMultistepScheduler with order 2).

2. Inversion Divergence During Image Editing

• Symptom: Inverting a real image ${\mathbf{x}}_0\to {\mathbf{x}}_T$ via invert_step and subsequently running step fails to reconstruct the original content.
• Root Cause: Accumulation of local linearization errors in Euler ODE steps under high classifier-free guidance scales.
• Remedy: Perform inversion with guidance_scale = 1.0 (unconditioned), increase inversion steps to $NFE=100$, and utilize Null-Text Optimization (Mokady et al., 2023).

3. Numerical Instability in Variance Computation

• Symptom: NaN tensors appearing during the final sampling steps ($t\to 0$).
• Root Cause: Dividing by $(1-{\bar{\alpha }}_t)$ or taking the square root of negative values caused by floating-point precision clamping.
• Remedy: Clamp variance terms strictly using torch.clamp(variance, min=0.0) and ensure ${\bar{\alpha }}_0$ explicitly evaluates to $1.0$.

References

1. Song, J., Meng, C., & Ermon, S. (2021). Denoising Diffusion Implicit Models. International Conference on Learning Representations (ICLR 2021 - Oral).
2. Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. Advances in Neural Information Processing Systems (NeurIPS 2020).
3. Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., & Poole, B. (2021). Score-Based Generative Modeling through Stochastic Differential Equations. ICLR 2021.
4. Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., & Zhu, J. (2022). DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling. NeurIPS 2022.
5. Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the Design Space of Diffusion-Based Generative Models. NeurIPS 2022.