← back

DDPM Noise Schedule (Linear Beta Schedule)

#395 · Deep Learning · Medium

⊣ Solve on deep-ml.com

Problem

Implement the linear beta noise schedule for Denoising Diffusion Probabilistic Models (DDPM). Compute the sequence of beta values, and derive alpha, alpha_bar (cumulative product), and the forward diffusion noise parameters for any timestep.

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
import numpy as np

def ddpm_linear_schedule(num_timesteps: int, beta_start: float = 1e-4, beta_end: float = 0.02) -> dict:
    betas = np.linspace(beta_start, beta_end, num_timesteps)
    alphas = 1.0 - betas
    alpha_bars = np.cumprod(alphas)

    sqrt_alpha_bars = np.sqrt(alpha_bars)
    sqrt_one_minus_alpha_bars = np.sqrt(1.0 - alpha_bars)

    return {
        "betas": betas,
        "alphas": alphas,
        "alpha_bars": alpha_bars,
        "sqrt_alpha_bars": sqrt_alpha_bars,
        "sqrt_one_minus_alpha_bars": sqrt_one_minus_alpha_bars,
    }


def forward_diffusion(x0: np.ndarray, t: int, schedule: dict) -> tuple[np.ndarray, np.ndarray]:
    noise = np.random.randn(*x0.shape)
    sqrt_ab = schedule["sqrt_alpha_bars"][t]
    sqrt_1_ab = schedule["sqrt_one_minus_alpha_bars"][t]
    xt = sqrt_ab * x0 + sqrt_1_ab * noise
    return xt, noise

Explanation

  1. The linear schedule linearly interpolates beta from beta_start to beta_end over T timesteps.
  2. Alpha at each step is 1 - beta. Alpha_bar is the cumulative product of alphas, representing total signal retained.
  3. The forward process adds noise: x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon.
  4. sqrt_alpha_bar controls how much original signal remains; sqrt_one_minus_alpha_bar controls noise magnitude.

Complexity

  • Time: O(T) for schedule computation
  • Space: O(T) for storing schedule arrays
← #394#396 →