Importance-Sampled Timesteps for Loss Estimation

mediumPROSignal-to-Noise Ratio and Loss WeightingDiffusion ModelsDiffusion Models 1

Problem Statement

Improved-DDPM reduces variance of the loss estimate by sampling timesteps proportional to their recent per-step loss, then reweighting each sample by the inverse probability. Compute those sampling probabilities and importance weights.

Background

Given the running mean loss L_t at each timestep, the importance-sampling distribution is

$p_t = \frac{\sqrt{\mathbb{E}[L_t^2]}}{\sum_s \sqrt{\mathbb{E}[L_s^2]}}$

Here you are given sq_loss[t] = E[L_t^2] directly. The unbiased loss estimator then weights a sample at t by 1/(T * p_t). Return both the distribution and the per-timestep importance weight.

Your Task

Implement:

def importance_sampling(sq_loss):

Return a dict with "probs" (the p_t) and "weights" (the 1/(T*p_t)), each rounded to 4 decimals. T = len(sq_loss).

Input Format

sq_loss: list of T non-negative values E[L_t^2], not all zero.

Output Format

A dict with two lists.

Sample

print(importance_sampling([4.0, 1.0]))

Output:

{'probs': [0.6667, 0.3333], 'weights': [0.75, 1.5]}

Example

Input

print(importance_sampling([4.0, 1.0]))

Output

{'probs': [0.6667, 0.3333], 'weights': [0.75, 1.5]}

Reasoning

Compute the square root of each input value to determine the unnormalized importance scores, as the distribution is proportional to $\sqrt{\mathbb{E}[L_t^2]}$ : $\sqrt{4.0} = 2.0$ and $\sqrt{1.0} = 1.0$ .
Sum these scores to find the normalization constant required for a probability distribution: $2.0 + 1.0 = 3.0$ .
Calculate the sampling probabilities $p_t$ by dividing each score by the total sum: $p_0 = 2.0 / 3.0 \approx 0.6667$ and $p_1 = 1.0 / 3.0 \approx 0.3333$ .
Determine the importance weights using the formula $1/(T \cdot p_t)$ , where $T=2$ is the number of timesteps: for the first timestep, $1 / (2 \cdot 0.6667) = 0.75$ , and for the second, $1 / (2 \cdot 0.3333) = 1.5$ .
The final output is {'probs': [0.6667, 0.3333], 'weights': [0.75, 1.5]}.

Constraints

1 <= T <= 100000, all sq_loss[t] >= 0, not all zero.
p_t is proportional to sqrt(sq_loss[t]), normalized.
weight_t = 1 / (T * p_t); round both to 4 decimals.

Learn this in

Probability & StatisticsFoundations

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 primerPREVIEW

1. Background Knowledge

In denoising diffusion probabilistic models (DDPM), the training objective is a sum of per-timestep losses $L_t$ . Because the variance of $L_t$ varies dramatically across the diffusion schedule (early timesteps often have much larger gradients than late ones), a uniform random draw of $t$ yields a high-variance Monte Carlo estimate of the total loss. Importance sampling fixes this: if we sample $t$ from a distribution $p_t$ that is proportional to the magnitude of the per-step loss, the resulting estimator has lower variance while remaining unbiased.

The Improved-DDPM paper proposes choosing $p_t$ proportional to $\sqrt{\mathbb{E}[L_t^2]}$ , i.e., the root-mean-square of the loss at step $t$ . This is a heuristic that balances bias and variance: sampling more from high-loss steps reduces the spread of the estimator, while the square-root (rather than the raw second moment) prevents a single dominant timestep from monopolizing the distribution. The key identity is that if $X \sim p$ , then $\mathbb{E}_p[f(X)/p(X)] = \mathbb{E}_q[f(X)]$ for any reference distribution $q$ ; here $q$ is the uniform distribution over $\{1,\dots,T\}$ , so the unbiased weight for a sample at $t$ is $1/(T \cdot p_t)$ .

Because the true $\mathbb{E}[L_t^2]$ is unknown during training, Improved-DDPM maintains an exponential moving average of the squared loss at each timestep. The problem hands you these running means directly as sq_loss[t], so you only need to normalize and invert.

2. Algorithm Approach

This is a normalization-then-inversion pattern:

Transform each element of sq_loss by taking its square root.
Normalize the transformed vector so it sums to 1 — this gives probs.
Compute weights as the element-wise reciprocal of T * probs.

No iterative or dynamic-programming structure is needed; it is a single pass over the input plus a few scalar reductions.

3. Step-by-Step Strategy

PRO

Continue the full explanation

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