← back

Forward & Backward Diffusion Process

#304 · Deep Learning · Medium

⊣ Solve on deep-ml.com

Problem

Implement both the forward and backward (reverse) diffusion processes. The forward process adds noise, and the backward process iteratively denoises using a trained model's noise predictions.

Solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import numpy as np

def linear_beta_schedule(T: int, beta_start: float = 1e-4,
                          beta_end: float = 0.02) -> np.ndarray:
    return np.linspace(beta_start, beta_end, T)

def forward_diffusion(x_0: np.ndarray, t: int, betas: np.ndarray,
                      noise: np.ndarray = None) -> dict:
    alphas = 1.0 - betas
    alpha_bars = np.cumprod(alphas)
    alpha_bar_t = alpha_bars[t]

    if noise is None:
        noise = np.random.randn(*x_0.shape)

    x_t = np.sqrt(alpha_bar_t) * x_0 + np.sqrt(1.0 - alpha_bar_t) * noise
    return {"x_t": x_t, "noise": noise}

def reverse_step(x_t: np.ndarray, t: int, predicted_noise: np.ndarray,
                 betas: np.ndarray) -> np.ndarray:
    alphas = 1.0 - betas
    alpha_bars = np.cumprod(alphas)

    alpha_t = alphas[t]
    alpha_bar_t = alpha_bars[t]

    # Predicted x_0 component
    coeff1 = 1.0 / np.sqrt(alpha_t)
    coeff2 = betas[t] / np.sqrt(1.0 - alpha_bar_t)

    mean = coeff1 * (x_t - coeff2 * predicted_noise)

    if t > 0:
        sigma = np.sqrt(betas[t])
        z = np.random.randn(*x_t.shape)
        x_prev = mean + sigma * z
    else:
        x_prev = mean

    return x_prev

def backward_diffusion(x_T: np.ndarray, noise_predictor, betas: np.ndarray) -> np.ndarray:
    T = len(betas)
    x_t = x_T.copy()

    for t in range(T - 1, -1, -1):
        predicted_noise = noise_predictor(x_t, t)
        x_t = reverse_step(x_t, t, predicted_noise, betas)

    return x_t

Explanation

  1. Forward process: same as before, adds noise using closed-form expression with alpha_bar.
  2. Reverse step: uses the model's noise prediction to estimate the mean of p(x_{t-1} | x_t):
  3. - mean = (1/sqrt(alpha_t)) (x_t - (beta_t / sqrt(1 - alpha_bar_t)) epsilon_predicted)
  4. - For t > 0, add stochastic noise scaled by sqrt(beta_t). At t=0, return the mean directly.
  5. Full backward diffusion: start from pure noise x_T and iteratively denoise for T steps.

Complexity

  • Time: O(T * (d + M)) where T is timesteps, d is data dimension, and M is model inference cost
  • Space: O(T + d) for schedule and current sample
← #303#305 →