The DDIM Update Rule

hardFREEDDPM and DDIM SamplingDiffusion ModelsDiffusion Models 1

Problem Statement

Implement the general DDIM update, with the interpolation parameter $\eta$ that sweeps continuously from a deterministic ODE solver to full DDPM ancestral sampling.

Background

DDIM (Song et al., 2021) observes that the DDPM objective constrains only the marginals $q(x_t \mid x_0)$ , so you are free to pick a non-Markovian reverse process with the same marginals. That yields a family of samplers indexed by $\eta \ge 0$ :

$\sigma_t = \eta \sqrt{\frac{1-\bar{\alpha}_{t'}}{1-\bar{\alpha}_t}}\sqrt{1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_{t'}}}$

x_{t'} = \underbrace{\sqrt{\bar{\alpha}_{t'}}\,\hat{x}_0}_{\text{predicted } x_0} + \underbrace{\sqrt{1-\bar{\alpha}_{t'}-\sigma_t^2}\;\hat{\varepsilon}}_{\text{direction pointing to } x_t} + \sigma_t z$$ where $t'$ is the *next* (smaller) timestep, not necessarily $t-1$. Read the middle term carefully: the amount of noise pointed back along $\hat{\varepsilon}$ is reduced by exactly the amount of fresh noise $\sigma_t$ that is about to be injected, so the total variance stays $1-\bar{\alpha}_{t'}$ whatever $\eta$ is. At $\eta = 0$ nothing random enters and the sampler becomes a deterministic map from the initial latent to the image -- which is what makes DDIM invertible and makes latent interpolation meaningful. At $\eta = 1$ with $t' = t-1$ it reduces exactly to DDPM ancestral sampling. ## Your Task Implement: ```python def ddim_step(x_t, eps, alpha_bar_t, alpha_bar_prev, eta=0.0, z=None): ``` Return $x_{t'}$. When **eta == 0** the **z** argument is unused and may be **None**. ## Input Format - **x_t**, **eps**: NumPy arrays of matching shape. - **alpha_bar_t** (float): $\bar{\alpha}$ at the current step, in **(0, 1)**. - **alpha_bar_prev** (float): $\bar{\alpha}$ at the target step, in **(0, 1]**. Use **1.0** for the final step to a clean sample. - **eta** (float): **0.0** is deterministic DDIM, **1.0** is DDPM. - **z**: array shaped like **x_t**, or **None**. ## Output Format A NumPy array shaped like **x_t**. ## Sample ```python x_t = np.array([0.4, -0.6]) eps = np.array([0.1, 0.2]) print(np.round(ddim_step(x_t, eps, 0.5, 0.8, eta=0.0), 4).tolist()) ``` **x0_hat = (x_t - sqrt(0.5)*eps)/sqrt(0.5)**, **sigma = 0** so the direction term is **sqrt(1-0.8)*eps**, and the output is **sqrt(0.8)*x0_hat + sqrt(0.2)*eps**.

Example

Input

x_t = np.array([0.4, -0.6])
eps = np.array([0.1, 0.2])
print(np.round(ddim_step(x_t, eps, 0.5, 0.8, eta=0.0), 4).tolist())

Output

[0.4612, -0.8484]

Reasoning

With eta = 0 the sigma term vanishes entirely. x0_hat = (x_t - sqrt(0.5)*eps)/sqrt(0.5) = [0.4657, -1.0485]. The direction term is sqrt(1 - 0.8 - 0) * eps = sqrt(0.2)*eps. Adding sqrt(0.8)*x0_hat to it gives the result. Note that the schedule moves up here (0.5 to 0.8), which is fine: DDIM only needs the two alpha-bar values.

Constraints

Clamp the argument of the direction-term square root at zero: use max(1 - alpha_bar_prev - sigma**2, 0.0) to absorb floating-point negatives.
Add the sigma * z term only when eta > 0.
alpha_bar_prev = 1.0 is legal and must yield exactly $\hat{x}_0$ (sigma and the direction term both vanish).
Do not draw randomness; z is supplied.
Do not round inside the function.

Learn this in

Reverse Process & DenoisingDiffusion Study Plan Sampling & AccelerationDiffusion 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 generate data by reversing a gradual noising process. The original DDPM (Denoising Diffusion Probabilistic Models) defines a Markovian reverse process where each step $x_{t-1}$ depends only on $x_t$ and a predicted noise $\hat{\varepsilon}$ . However, DDIM (Denoising Diffusion Implicit Models) reveals that the training objective only constrains the marginal distribution $q(x_t \mid x_0)$ , not the conditional transition $q(x_{t-1} \mid x_t)$ . This freedom allows us to define a non-Markovian reverse process that shares the same marginals but behaves differently.

The key innovation in DDIM is the introduction of an interpolation parameter $\eta \ge 0$ . When $\eta = 0$ , the reverse process becomes deterministic, effectively solving an ordinary differential equation (ODE). This determinism makes the process invertible, enabling applications like latent interpolation and image editing. When $\eta = 1$ and steps are unit-sized ( $t' = t-1$ ), the update rule reduces exactly to the standard DDPM ancestral sampling. For $0 < \eta < 1$ , the sampler interpolates between these two extremes, allowing for faster sampling with fewer steps while maintaining high sample quality.

The update rule relies on predicting the clean data $\hat{x}_0$ from the noisy input $x_t$ and the predicted noise $\hat{\varepsilon}$ . The next state $x_{t'}$ is constructed as a linear combination of $\hat{x}_0$ and $\hat{\varepsilon}$ , plus a controlled amount of fresh noise $\sigma_t z$ . The coefficient $\sigma_t$ is carefully designed so that the total variance of the resulting distribution matches the required marginal variance $1 - \bar{\alpha}_{t'}$ , regardless of the value of $\eta$ . This ensures that the generated samples remain consistent with the learned data distribution.

2. Algorithm Approach

The core task is to implement the vectorized update equation provided in the problem statement. The approach involves three main computational stages:

Predict Clean Data: Calculate $\hat{x}_0$ using the current noisy sample $x_t$ , the predicted noise $\hat{\varepsilon}$ , and the current noise schedule value $\bar{\alpha}_t$ .
Calculate Noise Scale: Compute $\sigma_t$ based on $\eta$ , $\bar{\alpha}_t$ , and $\bar{\alpha}_{t'}$ . This term controls how much randomness is injected into the next step.
Construct Next State: Combine the predicted clean data, the direction towards the current noisy sample (scaled by the remaining variance), and the fresh noise term.

The implementation must handle the special case where $\eta = 0$ (deterministic), in which case the noise term $z$ is irrelevant and $\sigma_t$ becomes zero. The function should be robust to None inputs for z when $\eta=0$ .

3. Step-by-Step Strategy

Compute $\hat{x}_0$ : Use the formula $\hat{x}_0 = \frac{x_t - \sqrt{1-\bar{\alpha}_t}\,\hat{\varepsilon}}{\sqrt{\bar{\alpha}_t}}$ .

Calculate $\sqrt{1-\bar{\alpha}_t}$ and $\sqrt{\bar{\alpha}_t}$ .
Perform element-wise operations to get $\hat{x}_0$ .

Compute $\sigma_t$ : Use the formula $\sigma_t = \eta \sqrt{\frac{1-\bar{\alpha}_{t'}}{1-\bar{\alpha}_t}}\sqrt{1 - \frac{\bar{\alpha}_t}{\bar{\alpha}_{t'}}}$ .

Note: If $\eta = 0$ , $\sigma_t$ is simply $0$ .
Be careful with the order of operations and square roots.
Handle the case where $\bar{\alpha}_t = \bar{\alpha}_{t'}$ (though typically $t' < t$ implies $\bar{\alpha}_{t'} > \bar{\alpha}_t$ in standard schedules, the math holds).

Compute the Direction Term: The term $\sqrt{1-\bar{\alpha}_{t'}-\sigma_t^2}\;\hat{\varepsilon}$ represents the direction pointing back towards $x_t$ adjusted for the new noise.

Calculate the coefficient $\sqrt{1-\bar{\alpha}_{t'}-\sigma_t^2}$ .
Multiply this coefficient by $\hat{\varepsilon}$ (element-wise).

Compute the Noise Term:

If $\eta > 0$ , ensure z is provided. Multiply $\sigma_t$ by z.
If $\eta = 0$ , this term is zero.

Combine Terms: Calculate $x_{t'} = \sqrt{\bar{\alpha}_{t'}}\,\hat{x}_0 + \text{Direction Term} + \text{Noise Term}$ .

Multiply $\hat{x}_0$ by $\sqrt{\bar{\alpha}_{t'}}$ .
Add the direction term.
Add the noise term.

Return Result: Return the resulting array $x_{t'}$ .

4. Common Pitfalls

Incorrect Variance Calculation: A common mistake is misinterpreting the term $\sqrt{1-\bar{\alpha}_{t'}-\sigma_t^2}$ . This term ensures that the sum of variances from the predicted mean and the noise equals the target marginal variance $1-\bar{\alpha}_{t'}$ . Do not confuse this with $\sqrt{1-\bar{\alpha}_{t'}}$ .
Handling $\eta = 0$ : When $\eta = 0$ , $\sigma_t$ becomes 0. The code should not crash if z is None. Ensure you check for eta == 0 or handle None gracefully before attempting to multiply z.
Square Root Domains: Ensure that the arguments to np.sqrt are non-negative. Due to floating-point precision, values like $1-\bar{\alpha}_t$ might become slightly negative if $\bar{\alpha}_t$ is very close to 1. Using np.maximum(0, value) inside sqrt can prevent NaNs.
Broadcasting Issues: Ensure that scalar values like $\alpha\_bar\_t$ are broadcast correctly against the array shapes of $x_t$ and $\varepsilon$ . NumPy usually handles this, but explicit scalar multiplication is safer.
Confusing $t$ and $t'$ : Remember that $t'$ is the previous (smaller) timestep in the reverse process, meaning $\bar{\alpha}_{t'} > \bar{\alpha}_t$ (since $\bar{\alpha}$ decreases as $t$ increases). The formula uses $\bar{\alpha}_{t'}$ for the target state's scale.

5. Time & Space Complexity

Time Complexity: $O(N)$ , where $N$ is the total number of elements in the input array $x_t$ . The operations involve element-wise arithmetic (addition, multiplication, square roots) across the entire array. There are no nested loops or complex dependencies between elements.
Space Complexity: $O(N)$ , to store the intermediate arrays $\hat{x}_0$ , the direction term, and the final result $x_{t'}$ . If implemented in-place or with careful memory management, it could be reduced, but typically temporary arrays of size $N$ are created.

solution.py

Python 3.13

Test results