Cosine Noise Schedule

mediumFREEForward Process and Noise SchedulesDiffusion ModelsDiffusion Models 1

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.

Example

Input

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

Output

[0.847, 0.4938, 0.1443, 0.0]
[0.153, 0.417, 0.7079, 0.999]

Reasoning

With T = 4 and s = 0.008 the offset grid is t/T = 0, 0.25, 0.5, 0.75, 1. Each maps through cos^2(((t/T + s)/(1+s)) * pi/2), and dividing by the t = 0 value normalises alpha_bar_0 to 1. Dropping that first entry leaves the four returned alpha-bars, which decay smoothly to ~0. Each beta is one minus the ratio of consecutive alpha-bars; the last ratio is essentially 0, so the raw beta would be ~1 and the 0.999 clip takes effect.

Constraints

Evaluate $f$ on the T+1 grid points $t = 0, 1, \dots, T$ , then divide by $f(0)$ .
Return only the T entries for $t \geq 1$ .
$\beta_t$ uses the ratio to the previous $\bar{\alpha}$ , with $\bar{\alpha}_0 = 1$ .
Clip betas to [0, 0.999]. Do not clip alpha_bars.
No rounding inside the function.

Learn this in

Forward Diffusion ProcessDiffusion Study Plan

More from Diffusion Models 1

Classifier-Free Guidance: Uncond/Cond Mixingeasy DDPM Posterior Mean and Variance from x0medium Adaptive CFG Scheduling over the Trajectoryhard Classifier Gradient Guidance Stepmedium Negative Prompt as the Unconditional Branchmedium

Concept primer

1. Background Knowledge

Diffusion models operate by gradually adding Gaussian noise to data over a fixed number of steps $T$ . The core of this process is the noise schedule, which dictates how much noise is added at each step. The parameter $\beta_t$ represents the variance of the noise added at step $t$ , while $\bar{\alpha}_t$ represents the cumulative product of the signal retention factors up to step $t$ . Specifically, $\bar{\alpha}_t = \prod_{i=1}^t (1 - \beta_i)$ . In the forward process, the noisy data $x_t$ is related to the original data $x_0$ by $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$ , where $\epsilon$ is standard Gaussian noise.

The linear noise schedule simply sets $\beta_t$ to increase linearly from a small value to a larger one. However, this causes information to be lost too quickly in the early stages, leaving the later stages with little signal to recover. The cosine schedule, introduced in Improved Denoising Diffusion Probabilistic Models (Nichol & Dhariwal, 2021), addresses this by defining $\bar{\alpha}_t$ directly using a cosine function. This ensures that the signal remains strong for a longer portion of the trajectory, allowing the model to learn more effectively from the intermediate steps. The schedule is defined as $\bar{\alpha}_t = \frac{f(t)}{f(0)}$ , where $f(t) = \cos^2\left(\frac{t/T + s}{1+s} \cdot \frac{\pi}{2}\right)$ and $s$ is a small offset to prevent $\beta_1$ from being too small.

Once $\bar{\alpha}_t$ is computed for all $t$ , the individual noise levels $\beta_t$ are derived from the ratio of consecutive cumulative alphas: $\beta_t = 1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}$ . This relationship stems from the definition $\bar{\alpha}_t = \bar{\alpha}_{t-1} (1 - \beta_t)$ . Since $\bar{\alpha}_t$ approaches 0 as $t$ approaches $T$ , the ratio $\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}$ can become very small, causing $\beta_t$ to approach 1. To maintain numerical stability during sampling, $\beta_t$ is clipped to a maximum value of 0.999.

2. Algorithm Approach

The problem requires implementing a specific mathematical function to generate two arrays: alpha_bars and betas. The approach involves vectorized computation using NumPy to efficiently calculate the cosine schedule for all time steps simultaneously.

Generate Time Steps: Create an array of time indices $t$ from $1$ to $T$ .
Compute $\bar{\alpha}_t$ : Apply the cosine formula directly to these time steps to get the cumulative signal retention factors. Ensure normalization by $f(0)$ .
Compute $\beta_t$ : Use the relationship between consecutive $\bar{\alpha}$ values to derive $\beta_t$ . This involves shifting the alpha_bars array to align $\bar{\alpha}_{t-1}$ with $\bar{\alpha}_t$ .
Clip $\beta_t$ : Apply the clipping constraint to ensure no $\beta_t$ exceeds 0.999.
Return Results: Return the unclipped alpha_bars and the clipped betas.

The key insight is that $\bar{\alpha}_t$ is computed independently for each $t$ , while $\beta_t$ depends on the ratio of adjacent $\bar{\alpha}$ values. Vectorization allows this to be done without explicit loops, leveraging NumPy's broadcasting and slicing capabilities.

3. Step-by-Step Strategy

Define Time Array: Create a NumPy array t containing integers from $1$ to $T$ . This represents the diffusion steps.
Calculate $f(t)$ : Compute the cosine squared function for each $t$ . The argument to the cosine is $\frac{t/T + s}{1+s} \cdot \frac{\pi}{2}$ . Use np.cos and square the result.
Normalize to get $\bar{\alpha}_t$ : Calculate $f(0)$ by substituting $t=0$ into the formula for $f(t)$ . Divide the array $f(t)$ by $f(0)$ to obtain alpha_bars. This ensures $\bar{\alpha}_0 = 1$ (though we only return $t=1 \dots T$ ).
Prepare for $\beta_t$ Calculation: To compute $\beta_t = 1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}$ , you need access to $\bar{\alpha}_{t-1}$ . Create a shifted version of alpha_bars where the first element is $\bar{\alpha}_0 = 1$ and the rest are alpha_bars[:-1]. Let's call this alpha_bars_prev.
Compute Raw Betas: Calculate betas_raw = 1 - alpha_bars / alpha_bars_prev. This gives the unclipped noise levels.
Clip Betas: Use np.clip(betas_raw, 0, 0.999) to ensure all values are within the valid range. This prevents numerical instability in the reverse process.
Return: Return the tuple (alpha_bars, betas_clipped).

4. Common Pitfalls

Indexing Errors: The problem specifies that the output arrays should cover $t = 1, \dots, T$ . Be careful with 0-based indexing in NumPy. When computing $\beta_t$ , you need $\bar{\alpha}_{t-1}$ . For $t=1$ , this is $\bar{\alpha}_0$ , which is not in the alpha_bars array (since it starts at $t=1$ ). You must explicitly handle $\bar{\alpha}_0 = 1$ .
Division by Zero: Although $\bar{\alpha}_0 = 1$ and $\bar{\alpha}_t > 0$ for $t < T$ , ensure that your implementation of $f(0)$ is correct. If $s$ is not handled correctly, $f(0)$ might not be 1, leading to incorrect normalization.
Clipping Order: The problem states that alpha_bars must be unclipped, while betas must be clipped. Do not clip alpha_bars. Also, ensure that the clipping is applied to betas after they are computed from the ratio, not before.
Floating Point Precision: When $t$ is close to $T$ , $\bar{\alpha}_t$ becomes very small. The ratio $\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}$ might suffer from precision issues, but the clipping to 0.999 mitigates the impact on the final result. Ensure you use float64 precision if necessary, though float32 is often sufficient.
Formula Misinterpretation: The formula for $f(t)$ involves $\cos^2$ , which means you should square the result of np.cos(), not pass a squared argument to np.cos().

5. Time & Space Complexity

Time Complexity: The algorithm involves vectorized operations on arrays of size $T$ . Computing $f(t)$ , normalizing, shifting, and clipping all take $O(T)$ time. Thus, the overall time complexity is $O(T)$ .
Space Complexity: The algorithm stores several arrays of size $T$ : t, f(t), alpha_bars, alpha_bars_prev, betas_raw, and betas_clipped. Each of these requires $O(T)$ space. Therefore, the overall space complexity is $O(T)$ .

This efficiency makes the cosine schedule suitable for large $T$ values, such as $T=1000$ or $T=2000$ , commonly used in diffusion models.

solution.py

Python 3.13

Test results