← back

Diffusion Cosine Noise Schedule

#403 · Deep Learning · Medium

⊣ Solve on deep-ml.com

Problem

Implement the cosine noise schedule for diffusion models, as proposed by Nichol & Dhariwal. The cosine schedule provides a smoother noise progression than the linear schedule, with more steps at low noise levels where details matter most.

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

def cosine_noise_schedule(num_timesteps: int, s: float = 0.008, max_beta: float = 0.999) -> dict:
    steps = np.arange(num_timesteps + 1)
    f_t = np.cos(((steps / num_timesteps) + s) / (1 + s) * (np.pi / 2)) ** 2

    alpha_bars = f_t / f_t[0]
    # Clip to avoid numerical issues
    alpha_bars = np.clip(alpha_bars, 1e-8, 1.0)

    # Compute betas from alpha_bars
    betas = np.zeros(num_timesteps)
    for t in range(num_timesteps):
        betas[t] = 1 - alpha_bars[t + 1] / alpha_bars[t]
    betas = np.clip(betas, 0.0, max_beta)

    alphas = 1.0 - betas
    alpha_bars_final = np.cumprod(alphas)

    return {
        "betas": betas,
        "alphas": alphas,
        "alpha_bars": alpha_bars_final,
        "sqrt_alpha_bars": np.sqrt(alpha_bars_final),
        "sqrt_one_minus_alpha_bars": np.sqrt(1.0 - alpha_bars_final),
    }

Explanation

  1. Define f(t) = cos((t/T + s) / (1 + s) * pi/2)^2 where s is a small offset preventing beta from being too small near t=0.
  2. Compute alpha_bar_t = f(t) / f(0), giving a smooth curve from 1 (no noise) to near 0 (full noise).
  3. Derive betas from consecutive alpha_bars: beta_t = 1 - alpha_bar_t / alpha_bar_{t-1}.
  4. Clip betas to max_beta for numerical stability.
  5. The cosine schedule spends more timesteps at low noise levels, improving detail generation compared to the linear schedule.

Complexity

  • Time: O(T)
  • Space: O(T)
← #402#404 →