AncestralSolver

class deepinv.sampling.AncestralSolver(timesteps=None, t_start=None, t_end=None, num_steps=None, eta=None, variance='small', rng=None)

Bases: BaseSDESolver

Ancestral solver for reverse-time diffusion SDEs, generalizing the DDPM and DDIM samplers.

It solves the reverse-time SDE (see deepinv.sampling.EDMDiffusionSDE), from \(t = T\) to \(t = 0\):

\[d x_t = \left(\frac{s'(t)}{s(t)} x_t - (1 + \alpha(t)) s(t)^2 \sigma(t) \sigma'(t) \nabla \log p_t(x_t) \right) dt + s(t) \sqrt{2 \alpha(t) \sigma(t) \sigma'(t)} d w_t.\]

On a step from \(t\) to \(t + dt\) (with \(dt < 0\) for reverse-time sampling), the solver computes the next state \(x_{t+dt}\) as:

\[x_{t+dt} = \frac{s(t+dt)}{s(t)} x_t + s(t) s(t+dt) \sigma(t)^2 \left(1 - r^{1 + \alpha}\right) \nabla \log p_t(x_t) + s(t+dt) \sigma(t+dt) \sqrt{1 - r^{2 \alpha}} \, z, \quad z \sim \mathcal{N}(0, \mathrm{Id}),\]

with \(r = \sigma(t+dt) / \sigma(t)\) and \(\alpha = \alpha(t)\). The noise level of the next state is exactly \(\sigma(t+dt)\).

Compared to a Euler-Maruyama step of deepinv.sampling.EulerSolver, it integrates the linear part and the noise exactly, and freezes the non-linear term. It is thus more accurate when discretizing with few steps.

The ancestral DDPM sampler [1] is obtained for \(\alpha = 1\), and the deterministic DDIM sampler [2] for \(\alpha = 0\). See deepinv.sampling.DDPMSolver and deepinv.sampling.DDIMSolver.

If the parameter \(\eta\) of DDIM [2] is given, the alpha of the SDE is ignored, and replaced on each step by

\[\alpha_\eta = \frac{\log\left(1 - \eta^2 (1 - r^2)\right)}{2 \log r},\]

for which the step is exactly the DDIM step with parameter \(\eta\).

With variance="large", the noise \(s(t+dt) \sigma(t+dt) \sqrt{1 - r^{2 \alpha}}\) is replaced by \(s(t) \sigma(t) \sqrt{1 - r^{2 \alpha}}\). For \(\alpha = 1\), this replaces the posterior variance \(\tilde{\beta}_t\) of DDPM by the variance \(\beta_t\) of the forward transition, see Section 3.2 of Ho et al.[1] and deepinv.sampling.DDPMSolver.

Note

The solver requires sde.sigma_t, sde.scale_t, sde.score and, if eta is None, sde.alpha, provided by deepinv.sampling.EDMDiffusionSDE (and its subclasses) and by deepinv.sampling.PosteriorDiffusion.

Parameters:

• timesteps (torch.Tensor, numpy.ndarray, list) – time steps at which the SDE will be discretized.

• t_start (float) – the starting time of the SDE, optional. If not provided, it will be inferred from the timesteps argument.

• t_end (float) – the ending time of the SDE, optional. If not provided, it will be inferred from the timesteps argument.

• num_steps (int) – the number of time steps for the SDE, optional. If not provided, it will be inferred from the timesteps argument.

• eta (float) – the stochasticity parameter \(\eta \in [0, 1]\) of DDIM, optional. If given, it replaces the alpha of the SDE by \(\alpha_\eta\). If None (default), the alpha of the SDE is used.

• variance (str) – the variance of the noise added on each step, either "small" (default) for the posterior variance (\(\tilde{\beta}_t\) for DDPM), or "large" for the variance of the forward transition (\(\beta_t\) for DDPM).

• rng (torch.Generator) – A random number generator for reproducibility.

Note

You can either provide the timesteps argument directly, or specify t_start, t_end, and num_steps to generate the time steps automatically (linearly with constant stepsize). If both are provided, the timesteps argument will take precedence.

---

References:

[1] (1,2)

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020.

[2] (1,2)

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations. 2020.

step(sde, t0, t1, x0, *args, **kwargs)

Perform a single ancestral step from time t0 to time t1, with current state x0, solving the reverse-time SDE.

Parameters:

• sde (deepinv.sampling.EDMDiffusionSDE, deepinv.sampling.PosteriorDiffusion) – the SDE to solve, which must provide sigma_t, scale_t, score and, if eta is None, alpha.

• t0 (float or torch.Tensor) – Time at the start of the step, of size (,).

• t1 (float or torch.Tensor) – Time at the end of the step, of size (,).

• x0 (torch.Tensor) – Current state of the system, of size (batch_size, d).

• *args – additional arguments for the score of the SDE.

• **kwargs – additional keyword arguments for the score of the SDE.

Return torch.Tensor, int:

Updated state of the system after the step and number of function evaluations (NFE) performed during the step (here 1).

Return type:

tuple[torch.Tensor, int]

Examples using AncestralSolver:

Using state-of-the-art diffusion models from HuggingFace Diffusers with DeepInverse

Using state-of-the-art diffusion models from HuggingFace Diffusers with DeepInverse

Building your diffusion posterior sampling method using SDEs

Building your diffusion posterior sampling method using SDEs

DPS – Posterior Sampling with Diffusion Models

DPS -- Posterior Sampling with Diffusion Models