📘

Closed-Form Forward Sample q(x_t | x_0)

EasyDiffusion Models 📚 Diffusion Models 1

Problem Statement

Jump straight from a clean sample $x_0$ to any noise level $t$ in a single operation, without simulating the intermediate steps.

Background

Applying the forward kernel $t$ times and repeatedly using the fact that the sum of two independent Gaussians is Gaussian collapses the whole chain into one closed form:

$q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t) I\right)$

Via the reparameterization trick this becomes an explicit, differentiable formula:

$x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)$

This is the single most-used line of code in diffusion training: it lets you sample a random timestep for every element of a batch and produce its noisy version in O(1), instead of running a t-step loop.

Note the coefficients are square roots of $\bar{\alpha}_t$ and $1-\bar{\alpha}_t$ , and that they satisfy $(\sqrt{\bar{\alpha}_t})^2 + (\sqrt{1-\bar{\alpha}_t})^2 = 1$ : the forward process is variance preserving.

Your Task

Implement:

def q_sample(x0, t, alpha_bars, noise):

Return the noisy sample $x_t$ as a NumPy array with the same shape as x0.

Input Format

x0: NumPy array of any shape, the clean data.
t (int): 0-based index into alpha_bars.
alpha_bars: 1-D NumPy array of cumulative alphas.
noise: NumPy array with the same shape as x0, the standard-normal sample $\varepsilon$ .

Output Format

A NumPy array with the same shape as x0.

Sample

x0 = np.array([1.0, -1.0, 0.5])
noise = np.array([0.2, 0.4, -0.6])
ab = np.array([0.9, 0.5, 0.1])
print(np.round(q_sample(x0, 1, ab, noise), 4).tolist())

With $\bar{\alpha}_1 = 0.5$ both coefficients equal $\sqrt{0.5} \approx 0.7071$ , so the output is 0.7071 * x0 + 0.7071 * noise.

Example:

Input:

x0 = np.array([1.0, -1.0, 0.5])
noise = np.array([0.2, 0.4, -0.6])
ab = np.array([0.9, 0.5, 0.1])
print(np.round(q_sample(x0, 1, ab, noise), 4).tolist())

Output:

[0.8485, -0.4243, -0.0707]

Reasoning:

At t = 1 the schedule gives alpha_bar = 0.5, so sqrt(0.5) = 0.7071 multiplies both terms. Element-wise: 0.7071*1.0 + 0.7071*0.2 = 0.8485, 0.7071*(-1.0) + 0.7071*0.4 = -0.4243, 0.7071*0.5 + 0.7071*(-0.6) = -0.0707. Scaling by alpha_bar instead of its square root would break the variance-preserving property.

Constraints:

0 <= t < len(alpha_bars).
Use the square roots of $\bar{\alpha}_t$ and $1-\bar{\alpha}_t$ , not the raw values.
Do not draw any randomness: noise is supplied.
Do not round inside the function.

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

Diffusion models are generative models that learn to reverse a gradual noising process. The forward process systematically adds Gaussian noise to data over $T$ timesteps, transforming a complex data distribution into a simple isotropic Gaussian distribution. While the theoretical definition describes this as a Markov chain where $x_t$ depends only on $x_{t-1}$ , simulating each step sequentially is computationally expensive and unnecessary for training.

The key insight is that the composition of Gaussian distributions remains Gaussian. By leveraging the properties of variance and mean propagation, we can derive a closed-form solution that allows us to jump directly from the clean data $x_0$ to any noisy state $x_t$ in a single operation. This is known as the reparameterization trick. It expresses $x_t$ as a deterministic function of $x_0$ , the timestep $t$ , and a random noise vector $\varepsilon$ . This enables efficient sampling during training by allowing us to pick random timesteps for each sample in a batch and compute their noisy counterparts in $O(1)$ time, rather than iterating through all previous steps.

The formula relies on cumulative alphas ( $\bar{\alpha}_t$ ), which represent the product of individual noise schedule coefficients up to time $t$ . The signal-to-noise ratio is controlled by $\sqrt{\bar{\alpha}_t}$ , while the noise magnitude is controlled by $\sqrt{1-\bar{\alpha}_t}$ . These coefficients are square roots because they scale the standard deviations of the Gaussian distributions, ensuring the total variance is preserved correctly.

2. Algorithm Approach

The problem requires implementing the direct sampling formula derived from the forward diffusion process. The approach is purely arithmetic and vectorized, relying on NumPy's broadcasting capabilities.

Retrieve Coefficients: Access the specific cumulative alpha value $\bar{\alpha}_t$ corresponding to the given timestep $t$ from the alpha_bars array.
Compute Scaling Factors: Calculate the two scaling coefficients:

Signal coefficient: $\sqrt{\bar{\alpha}_t}$
Noise coefficient: $\sqrt{1 - \bar{\alpha}_t}$

Apply Linear Combination: Compute the weighted sum of the clean data x0 and the provided noise using these coefficients.
Return Result: Return the resulting array, which represents $x_t$ .

This approach avoids any loops over timesteps. It treats the operation as a single vectorized linear transformation, which is optimal for performance in deep learning frameworks.

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.