📘

Forward-Diffuse a Batch to Timestep t

MediumDiffusion Models 📚 Diffusion Models 1

Problem Statement

Apply the closed-form forward process to a batch of clean samples at possibly different timesteps, the exact operation inside a training step. Each sample gets noised with its own t and its own noise vector.

Background

The forward marginal is q(x_t | x_0) = N(sqrt(alpha_bar_t) x_0, (1 - alpha_bar_t) I), so with a standard-normal noise eps:

$x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1 - \bar{\alpha}_t}\, \varepsilon$

In a training batch each row i has its own timestep t_i (drawn uniformly) and its own noise, so the per-row scale factors are gathered from the alpha_bar table.

Your Task

Implement:

def forward_diffuse(x0, eps, alpha_bar, t):

x0, eps: B x D nested lists (clean samples and noise).
alpha_bar: the full schedule (list).
t: list of B integer timesteps, one per row.

Return the B x D noised batch as a nested list rounded to 4 decimals.

Input Format

x0, eps: B x D.
alpha_bar: list of cumulative products.
t: list of B indices into alpha_bar.

Output Format

A B x D nested list rounded to 4 decimals.

Sample

x0 = [[1.0, 1.0]]
eps = [[1.0, -1.0]]
print(forward_diffuse(x0, eps, [0.99, 0.5], [1]))

Output:

[[1.4142, 0.0]]

Example:

Input:

x0 = [[1.0, 1.0]]
eps = [[1.0, -1.0]]
print(forward_diffuse(x0, eps, [0.99, 0.5], [1]))

Output:

[[1.4142, 0.0]]

Reasoning:

Look up the noise schedule value for the given timestep: with $t = [1]$ , we index into alpha_bar at position 1 to get $\bar{\alpha} = 0.5$ .
Compute the signal and noise scale factors for this timestep: the signal weight is $a = \sqrt{0.5} \approx 0.7071$ and the noise weight is $b = \sqrt{1 - 0.5} = \sqrt{0.5} \approx 0.7071$ .
Apply the forward diffusion formula $x_t = a \cdot x_0 + b \cdot \epsilon$ $x_{t} = a \cdot x_{0} + b \cdot ϵ$ element-wise to the single sample:
- For the first dimension: $0.7071 \cdot 1.0 + 0.7071 \cdot 1.0 = 1.4142$ .
- For the second dimension: $0.7071 \cdot 1.0 + 0.7071 \cdot (-1.0) = 0.0$ .
The final output is [[1.4142, 0.0]]

Constraints:

x0 and eps share shape B x D; len(t) == B.
Gather alpha_bar[t_i] per row; scales are sqrt(alpha_bar) and sqrt(1 - alpha_bar).
Round to 4 decimals; avoid -0.0.

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

1. Background Knowledge

In denoising diffusion probabilistic models (DDPMs), the forward process gradually adds Gaussian noise to a clean data sample $x_0$ over $T$ discrete timesteps. The key insight is that although the process is defined as a Markov chain, the marginal distribution $q(x_t \mid x_0)$ has a closed-form solution: it is a Gaussian with mean $\sqrt{\bar{\alpha}_t}\, x_0$ and covariance $(1 - \bar{\alpha}_t)\, I$ . This means you can jump directly from $x_0$ to any $x_t$ in one step, without simulating all intermediate timesteps.

The quantity $\bar{\alpha}_t$ is the cumulative product of the per-step signal retention factors $\alpha_i = 1 - \beta_i$ , where $\beta_i$ is the noise variance at step $i$ . As $t$ increases, $\bar{\alpha}_t$ decreases toward zero, meaning the signal is progressively drowned out by noise. At $t=0$ , $\bar{\alpha}_0 = 1$ (pure signal); at large $t$ , $\bar{\alpha}_t \approx 0$ (pure noise).

In a training batch, each sample is noised at a different randomly sampled timestep $t_i$ . This is essential for the model to learn to denoise at all noise levels simultaneously. The per-sample scale factors $\sqrt{\bar{\alpha}_{t_i}}$ and $\sqrt{1 - \bar{\alpha}_{t_i}}$ are gathered from the precomputed schedule and applied row-wise.

2. Algorithm Approach

This is a row-wise element-wise operation with per-row scalar coefficients. The pattern is:

For each row $i$ in the batch, look up $\bar{\alpha}_{t_i}$ from the schedule.
Compute the two scalar coefficients: $s_i = \sqrt{\bar{\alpha}_{t_i}}$ and $n_i = \sqrt{1 - \bar{\alpha}_{t_i}}$ .
Apply the linear combination $x_t[i] = s_i \cdot x_0[i] + n_i \cdot \varepsilon[i]$ element-wise across the feature dimension $D$ .

No loops over timesteps are needed—just a single pass over the batch.

3. Step-by-Step Strategy

👑

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.