DDPM Reverse Step

Take one denoising step of DDPM sampling, from x_t to x_{t-1}.

At sampling time a diffusion model starts from pure noise and removes a little of it at each step. The network predicts the noise eps_pred that is in x_t. From it, DDPM computes the mean of x_{t-1}, then adds fresh noise:

alpha_t     = 1 - beta_t
alpha_bar_t = prod_{s=0}^{t} alpha_s
mean        = (x_t - (beta_t / sqrt(1 - alpha_bar_t)) * eps_pred) / sqrt(alpha_t)
x_{t-1}     = mean + sqrt(beta_t) * noise      if t > 0
            = mean                             if t == 0

The last step adds no noise, so the final sample is the mean.

Function signature

def ddpm_step(x_t, t, eps_pred, betas, noise):

Example

Input

ddpm_step(np.array([1.0]), t=2, eps_pred=np.array([0.5]),
          betas=np.array([0.1, 0.2, 0.3]), noise=np.array([0.2]))

Output

array([1.05020653])

Reasoning

alpha_2 = 0.7 and alpha_bar_2 = 0.9 * 0.8 * 0.7 = 0.504. The noise coefficient is 0.3 / sqrt(0.496) = 0.426, so the mean is (1 - 0.426 * 0.5) / sqrt(0.7) = 0.9407. Adding sqrt(0.3) * 0.2 = 0.1095 gives 1.0502.

Constraints and notes

• t is a 0-based step index.
• x_t, eps_pred and noise have the same shape.
• At t = 0, ignore noise.