Cosine Noise Schedule

Problem Statement

Implement the cosine noise schedule from Improved Denoising Diffusion Probabilistic Models (Nichol & Dhariwal, 2021), which is defined on $\bar{\alpha}$ directly and only afterwards converted back into betas.

Background

The linear schedule destroys information too fast: by the middle of the trajectory the image is already almost pure noise, so the later steps contribute little. The cosine schedule instead defines

$f(t) = \cos^2\!\left(\frac{t/T + s}{1+s} \cdot \frac{\pi}{2}\right), \qquad \bar{\alpha}_t = \frac{f(t)}{f(0)}$

for $t = 0, 1, \dots, T$, where the small offset $s = 0.008$ keeps $\beta_1$ from becoming vanishingly small. By construction $\bar{\alpha}_0 = 1$ (no noise) and $\bar{\alpha}_T \approx 0$ (pure noise).

The betas are then recovered from the ratio of consecutive cumulative products:

$\beta_t = 1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}$

Near $t = T$ this ratio goes to zero and $\beta_t$ approaches 1, which makes the sampler unstable, so the paper clips every beta to at most 0.999.

Your Task

Implement:

def cosine_schedule(T, s=0.008):

Return (alpha_bars, betas), both 1-D NumPy arrays of length T covering $t = 1, \dots, T$ (0-based index i is $t = i+1$). alpha_bars must be the unclipped $\bar{\alpha}_t$ values; betas must be clipped to the range [0, 0.999].

Input Format

• T (int): number of diffusion steps, T >= 1.
• s (float): the small offset, default 0.008.

Output Format

A tuple (alpha_bars, betas) of NumPy arrays of length T.

Sample

ab, b = cosine_schedule(4)
print(np.round(ab, 4).tolist())
print(np.round(b, 4).tolist())

With T = 4 the cosine curve drops from just below 1 down to essentially 0, and the betas grow monotonically, with the final one hitting the 0.999 clip.