← Back to the problem

USA USAAIO 2026 · Correction by IOAI Training

Diffusion Models

This page gives the full approach. Read it once you have searched.

Metric
written derivations (50 points)
Evaluated on
written answers
Our score
all 9 parts derived, checked numerically
Expected level
complete derivations for every part

Expected level reached

We solved the problem ourselves: downloaded the data, ran the experiments and measured the score below. The statement and data belong to their organisers.

The problem in brief

This is the theory half of a denoising diffusion model (DDPM), in one dimension so that every formula fits on a line. A data value x0x_0 is corrupted step by step: at step tt it is shrunk by αt\sqrt{\alpha_t} and fresh Gaussian noise of variance 1−αt1 - \alpha_t is added. You must show that after tt steps xtx_t is still a simple Gaussian around a shrunk x0x_0, find what it converges to, compute the Kullback-Leibler divergence between two Gaussians, derive the exact "reverse" distribution of xt−1x_{t-1} given xtx_t and x0x_0, and finally show that the training objective collapses to a weighted squared error between the true noise and the noise predicted by a network. A last conceptual part asks how such a network should be told which time step it is working on.

How you are scored

50 points over 9 parts (5 or 10 points each), graded on complete derivations in a notebook. There is no dataset. We derived every formula by hand and then checked each one numerically in solution.py (Monte Carlo for the forward process, numerical integration for the KL, brute-force Bayes on a grid for the posterior, random tests for the final identities): all agree to 4 to 6 decimals.

partresult
3.1.1A2=α1α2A_2 = \sqrt{\alpha_1\alpha_2}, B2=1−α1α2B_2 = \sqrt{1 - \alpha_1\alpha_2}
3.1.2xt∣x0∼N(αˉt x0, 1−αˉt)x_t \mid x_0 \sim \mathcal N(\sqrt{\bar\alpha_t}\,x_0,\ 1 - \bar\alpha_t)
3.1.3μt∣0→0\mu_{t\mid 0} \to 0, σt∣02→1\sigma^2_{t\mid 0} \to 1 (when αˉt→0\bar\alpha_t \to 0)
3.2.1ln⁡σpσq+σq2+(μq−μp)22σp2−12\ln\frac{\sigma_p}{\sigma_q} + \frac{\sigma_q^2 + (\mu_q - \mu_p)^2}{2\sigma_p^2} - \frac12
3.2.3β~t=(αt−αˉt)(1−αt)αt(1−αˉt)\tilde\beta_t = \frac{(\alpha_t - \bar\alpha_t)(1-\alpha_t)}{\alpha_t(1-\bar\alpha_t)}, μ~t\tilde\mu_t below
3.2.4μ~t=1αt(xt−1−αt1−αˉtϵˉt)\tilde\mu_t = \frac{1}{\sqrt{\alpha_t}}\big(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\bar\epsilon_t\big)
3.2.5DKL=1−αt2(αt−αˉt) (ϵˉt−ϵθ(xt,t))2D_{KL} = \frac{1-\alpha_t}{2(\alpha_t - \bar\alpha_t)}\,(\bar\epsilon_t - \epsilon_\theta(x_t, t))^2

The key idea

Everything rests on one fact: a linear combination of independent Gaussians is Gaussian, with means and variances that add. If ϵ,ϵ′\epsilon, \epsilon' are independent N(0,1)\mathcal N(0,1), then

a ϵ+b ϵ′∼N(0, a2+b2),soa ϵ+b ϵ′=a2+b2 ϵˉ with ϵˉ∼N(0,1).a\,\epsilon + b\,\epsilon' \sim \mathcal N(0,\ a^2 + b^2), \quad\text{so}\quad a\,\epsilon + b\,\epsilon' = \sqrt{a^2 + b^2}\ \bar\epsilon \text{ with } \bar\epsilon \sim \mathcal N(0,1).

That turns tt noisy steps into a single jump from x0x_0 to xtx_t (so training can sample any tt directly), makes every distribution in the model Gaussian (so the KL is a closed formula), and lets the reverse mean be rewritten in terms of the noise (so the network only has to guess "which noise was added"). The schedule is designed so that each step preserves variance: Var⁡[xt]=αtVar⁡[xt−1]+(1−αt)\operatorname{Var}[x_t] = \alpha_t \operatorname{Var}[x_{t-1}] + (1 - \alpha_t) equals 1 whenever Var⁡[xt−1]=1\operatorname{Var}[x_{t-1}] = 1.

Histograms of x_t for a data distribution made of two narrow clusters at -2 and +2. At t = 10 the clusters are barely blurred; by t = 60 they are wide and overlapping; at t = 100 the density is a broad flat bump; at t = 200 it matches the standard normal density almost perfectly.

Part 3.1: the forward process

3.1.1 Two steps

Write the two steps and substitute the first into the second:

x2=α2 x1+1−α2 ϵ2=α2(α1 x0+1−α1 ϵ1)+1−α2 ϵ2=α1α2 x0+α2(1−α1) ϵ1+1−α2 ϵ2⏟noise.x_2 = \sqrt{\alpha_2}\,x_1 + \sqrt{1-\alpha_2}\,\epsilon_2 = \sqrt{\alpha_2}\big(\sqrt{\alpha_1}\,x_0 + \sqrt{1-\alpha_1}\,\epsilon_1\big) + \sqrt{1-\alpha_2}\,\epsilon_2 = \sqrt{\alpha_1\alpha_2}\,x_0 + \underbrace{\sqrt{\alpha_2(1-\alpha_1)}\,\epsilon_1 + \sqrt{1-\alpha_2}\,\epsilon_2}_{\text{noise}} .

The noise term is a combination of the independent ϵ1,ϵ2∼N(0,1)\epsilon_1, \epsilon_2 \sim \mathcal N(0,1), so it is Gaussian with mean 0 and variance

α2(1−α1)+(1−α2)=1−α1α2.\alpha_2(1-\alpha_1) + (1-\alpha_2) = 1 - \alpha_1\alpha_2 .

Hence it equals 1−α1α2 ϵˉ2\sqrt{1-\alpha_1\alpha_2}\ \bar\epsilon_2 with ϵˉ2∼N(0,1)\bar\epsilon_2 \sim \mathcal N(0,1), and

x2=A2x0+B2ϵˉ2,A2=α1α2=αˉ2,B2=1−α1α2=1−αˉ2.x_2 = A_2 x_0 + B_2\bar\epsilon_2, \qquad A_2 = \sqrt{\alpha_1\alpha_2} = \sqrt{\bar\alpha_2},\quad B_2 = \sqrt{1-\alpha_1\alpha_2} = \sqrt{1-\bar\alpha_2}.

Conditionally on x0x_0, the first term is a constant and the second is N(0,B22)\mathcal N(0, B_2^2), and the ϵ\epsilon's are independent of x0x_0. So x2∣x0∼N(α1α2 x0, 1−α1α2)x_2 \mid x_0 \sim \mathcal N(\sqrt{\alpha_1\alpha_2}\,x_0,\ 1 - \alpha_1\alpha_2).

3.1.2 Any number of steps

Claim: xt=αˉt x0+1−αˉt ϵˉtx_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\bar\epsilon_t with ϵˉt∼N(0,1)\bar\epsilon_t \sim \mathcal N(0,1) independent of x0x_0 (a combination of ϵ1,…,ϵt\epsilon_1, \dots, \epsilon_t).

Base case t=1t = 1: αˉ1=α1\bar\alpha_1 = \alpha_1 and ϵˉ1=ϵ1\bar\epsilon_1 = \epsilon_1; this is the definition.

Induction step: assume the claim at t−1t-1. Then

xt=αt(αˉt−1 x0+1−αˉt−1 ϵˉt−1)+1−αt ϵt=αˉt x0+αt(1−αˉt−1) ϵˉt−1+1−αt ϵt,x_t = \sqrt{\alpha_t}\big(\sqrt{\bar\alpha_{t-1}}\,x_0 + \sqrt{1-\bar\alpha_{t-1}}\,\bar\epsilon_{t-1}\big) + \sqrt{1-\alpha_t}\,\epsilon_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{\alpha_t(1-\bar\alpha_{t-1})}\,\bar\epsilon_{t-1} + \sqrt{1-\alpha_t}\,\epsilon_t ,

using αtαˉt−1=αˉt\alpha_t\bar\alpha_{t-1} = \bar\alpha_t. ϵˉt−1\bar\epsilon_{t-1} depends only on ϵ1,…,ϵt−1\epsilon_1,\dots,\epsilon_{t-1}, so it is independent of ϵt\epsilon_t; the noise is Gaussian with variance αt−αtαˉt−1+1−αt=1−αˉt\alpha_t - \alpha_t\bar\alpha_{t-1} + 1 - \alpha_t = 1 - \bar\alpha_t. This proves the claim with

At=αˉt,Bt=1−αˉt,xt∣x0∼N(αˉt x0, 1−αˉt).A_t = \sqrt{\bar\alpha_t}, \qquad B_t = \sqrt{1-\bar\alpha_t}, \qquad x_t \mid x_0 \sim \mathcal N\big(\sqrt{\bar\alpha_t}\,x_0,\ 1-\bar\alpha_t\big).

Monte Carlo check (400 000 chains, x0=1.7x_0 = 1.7, linear schedule): at t=50t = 50 the empirical mean and variance are 1.4538 and 0.2690 against 1.4534 and 0.2691 predicted.

3.1.3 The limit

μt∣0=αˉt x0\mu_{t\mid0} = \sqrt{\bar\alpha_t}\,x_0 and σt∣02=1−αˉt\sigma^2_{t\mid0} = 1 - \bar\alpha_t, so everything depends on lim⁡αˉt\lim \bar\alpha_t. The sequence αˉt\bar\alpha_t is positive and decreasing (each factor is in (0,1)(0,1)), so it converges. For any practical schedule, where the αt\alpha_t stay below some αmax⁡<1\alpha_{\max} < 1, we have αˉt≤αmax⁡t→0\bar\alpha_t \le \alpha_{\max}^t \to 0, hence

lim⁡t→∞μt∣0=0,lim⁡t→∞σt∣02=1:\lim_{t\to\infty}\mu_{t\mid0} = 0, \qquad \lim_{t\to\infty}\sigma^2_{t\mid0} = 1 :

the chain forgets x0x_0 entirely and xtx_t tends to N(0,1)\mathcal N(0,1), whatever the data. That is what makes generation possible: we can start the reverse process from pure noise.

A rigorous footnote worth knowing: 0<αt<10 < \alpha_t < 1 alone is not quite enough. Since ln⁡αˉt=∑sln⁡αs\ln\bar\alpha_t = \sum_s \ln \alpha_s and ln⁡αs≈−(1−αs)\ln \alpha_s \approx -(1 - \alpha_s), we get αˉt→0\bar\alpha_t \to 0 exactly when ∑t(1−αt)=∞\sum_t (1 - \alpha_t) = \infty. With αt=1−1/(t+1)2\alpha_t = 1 - 1/(t+1)^2 the noise is summable and αˉt→1/2\bar\alpha_t \to 1/2: the limit is then N(x0/2, 1/2)\mathcal N(x_0/\sqrt2,\ 1/2) and the data is never fully erased.

Part 3.2: the training objective

3.2.1 KL divergence between two Gaussians

By definition DKL(q ∥ p)=Ex∼q[ln⁡q(x)−ln⁡p(x)]D_{KL}(q\,\|\,p) = \mathbb E_{x\sim q}[\ln q(x) - \ln p(x)]. With the Gaussian density,

ln⁡q(x)−ln⁡p(x)=ln⁡σpσq−(x−μq)22σq2+(x−μp)22σp2.\ln q(x) - \ln p(x) = \ln\frac{\sigma_p}{\sigma_q} - \frac{(x-\mu_q)^2}{2\sigma_q^2} + \frac{(x-\mu_p)^2}{2\sigma_p^2}.

Take the expectation under qq: Eq[(x−μq)2]=σq2\mathbb E_q[(x-\mu_q)^2] = \sigma_q^2, and writing x−μp=(x−μq)+(μq−μp)x - \mu_p = (x - \mu_q) + (\mu_q - \mu_p), the cross term has zero mean, so Eq[(x−μp)2]=σq2+(μq−μp)2\mathbb E_q[(x-\mu_p)^2] = \sigma_q^2 + (\mu_q - \mu_p)^2. Therefore

DKL(q ∥ p)=ln⁡σpσq+σq2+(μq−μp)22σp2−12=12ln⁡σp2σq2+σq2+(μq−μp)22σp2−12.\boxed{D_{KL}(q\,\|\,p) = \ln\frac{\sigma_p}{\sigma_q} + \frac{\sigma_q^2 + (\mu_q - \mu_p)^2}{2\sigma_p^2} - \frac12 = \frac12\ln\frac{\sigma_p^2}{\sigma_q^2} + \frac{\sigma_q^2 + (\mu_q-\mu_p)^2}{2\sigma_p^2} - \frac12 .}

Sanity checks: it is 0 when q=pq = p, and numerical integration on a fine grid matches it to 6 decimals on three random pairs (e.g. 0.740647 both ways).

3.2.2 Applied to the diffusion step

Substitute μq=μ~t(xt,x0)\mu_q = \tilde\mu_t(x_t, x_0), σq2=β~t\sigma_q^2 = \tilde\beta_t, μp=μθ(xt,t)\mu_p = \mu_\theta(x_t, t), σp2=σt2\sigma_p^2 = \sigma_t^2:

DKL(q(xt−1∣xt,x0) ∥ pθ(xt−1∣xt))=12ln⁡σt2β~t+β~t+(μ~t(xt,x0)−μθ(xt,t))22σt2−12.D_{KL}\big(q(x_{t-1}\mid x_t, x_0)\,\|\,p_\theta(x_{t-1}\mid x_t)\big) = \frac12\ln\frac{\sigma_t^2}{\tilde\beta_t} + \frac{\tilde\beta_t + \big(\tilde\mu_t(x_t,x_0) - \mu_\theta(x_t,t)\big)^2}{2\sigma_t^2} - \frac12 .

3.2.3 The true reverse step

By Bayes' rule, and because the chain is Markov (xtx_t depends on the past only through xt−1x_{t-1}),

q(xt−1∣xt,x0)∝q(xt∣xt−1)  q(xt−1∣x0)∝exp⁡ ⁣(−(xt−αt xt−1)22(1−αt)−(xt−1−αˉt−1 x0)22(1−αˉt−1)),q(x_{t-1}\mid x_t, x_0) \propto q(x_t \mid x_{t-1})\; q(x_{t-1}\mid x_0) \propto \exp\!\Big(-\frac{(x_t - \sqrt{\alpha_t}\,x_{t-1})^2}{2(1-\alpha_t)} - \frac{(x_{t-1} - \sqrt{\bar\alpha_{t-1}}\,x_0)^2}{2(1-\bar\alpha_{t-1})}\Big),

using 3.1.2 at time t−1t-1. The exponent is a quadratic in xt−1x_{t-1}, so the posterior is Gaussian. Matching it to −(xt−1−μ~)22β~+const-\frac{(x_{t-1} - \tilde\mu)^2}{2\tilde\beta} + \text{const}:

  • coefficient of −12xt−12-\tfrac12 x_{t-1}^2 (the precision): 1β~t=αt1−αt+11−αˉt−1=αt−αˉt+1−αt(1−αt)(1−αˉt−1)=1−αˉt(1−αt)(1−αˉt−1)\dfrac{1}{\tilde\beta_t} = \dfrac{\alpha_t}{1-\alpha_t} + \dfrac{1}{1-\bar\alpha_{t-1}} = \dfrac{\alpha_t - \bar\alpha_t + 1 - \alpha_t}{(1-\alpha_t)(1-\bar\alpha_{t-1})} = \dfrac{1-\bar\alpha_t}{(1-\alpha_t)(1-\bar\alpha_{t-1})};
  • coefficient of xt−1x_{t-1}: μ~tβ~t=αt xt1−αt+αˉt−1 x01−αˉt−1\dfrac{\tilde\mu_t}{\tilde\beta_t} = \dfrac{\sqrt{\alpha_t}\,x_t}{1-\alpha_t} + \dfrac{\sqrt{\bar\alpha_{t-1}}\,x_0}{1-\bar\alpha_{t-1}}.

Hence

β~t=(1−αˉt−1)(1−αt)1−αˉt,μ~t(xt,x0)=αt (1−αˉt−1)1−αˉt xt+αˉt−1 (1−αt)1−αˉt x0.\tilde\beta_t = \frac{(1-\bar\alpha_{t-1})(1-\alpha_t)}{1-\bar\alpha_t}, \qquad \tilde\mu_t(x_t,x_0) = \frac{\sqrt{\alpha_t}\,(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}\,x_t + \frac{\sqrt{\bar\alpha_{t-1}}\,(1-\alpha_t)}{1-\bar\alpha_t}\,x_0 .

The statement wants only αt\alpha_t and αˉt\bar\alpha_t, so replace αˉt−1=αˉt/αt\bar\alpha_{t-1} = \bar\alpha_t/\alpha_t (hence 1−αˉt−1=αt−αˉtαt1 - \bar\alpha_{t-1} = \frac{\alpha_t - \bar\alpha_t}{\alpha_t} and αˉt−1=αˉt/αt\sqrt{\bar\alpha_{t-1}} = \sqrt{\bar\alpha_t/\alpha_t}):

β~t=(αt−αˉt)(1−αt)αt (1−αˉt),μ~t(xt,x0)=(αt−αˉt) xt+αˉt (1−αt) x0αt (1−αˉt).\boxed{\tilde\beta_t = \frac{(\alpha_t - \bar\alpha_t)(1-\alpha_t)}{\alpha_t\,(1-\bar\alpha_t)}, \qquad \tilde\mu_t(x_t,x_0) = \frac{(\alpha_t - \bar\alpha_t)\,x_t + \sqrt{\bar\alpha_t}\,(1-\alpha_t)\,x_0}{\sqrt{\alpha_t}\,(1-\bar\alpha_t)} .}

Interpretation: the mean is a weighted average of "where we are" (xtx_t, rescaled) and "where we came from" (x0x_0), weighted by how much each one tells us; the variance is smaller than both 1−αt1-\alpha_t and 1−αˉt−11 - \bar\alpha_{t-1}, because two sources of information are combined. Brute-force check: normalising the product of the two densities on a grid of 400 001 points (for t=30t = 30, xt=0.4x_t = 0.4, x0=−1.2x_0 = -1.2) gives mean 0.294652 and variance 0.006911, exactly the formulas. (For t=1t = 1, β~1=0\tilde\beta_1 = 0: given x0x_0, x0x_0 is known; the formulas are meant for t≥2t \ge 2.)

3.2.4 In terms of the noise

From 3.1.2, x0=xt−1−αˉt ϵˉtαˉtx_0 = \dfrac{x_t - \sqrt{1-\bar\alpha_t}\,\bar\epsilon_t}{\sqrt{\bar\alpha_t}}. Substituting into μ~t\tilde\mu_t, the αˉt\sqrt{\bar\alpha_t} cancels in the second term:

μ~t=(αt−αˉt) xt+(1−αt) xt−(1−αt)1−αˉt ϵˉtαt (1−αˉt).\tilde\mu_t = \frac{(\alpha_t - \bar\alpha_t)\,x_t + (1-\alpha_t)\,x_t - (1-\alpha_t)\sqrt{1-\bar\alpha_t}\,\bar\epsilon_t}{\sqrt{\alpha_t}\,(1-\bar\alpha_t)} .

The xtx_t coefficients add to 1−αˉt1 - \bar\alpha_t, which cancels the denominator, and 1−αˉt1−αˉt=11−αˉt\frac{\sqrt{1-\bar\alpha_t}}{1-\bar\alpha_t} = \frac{1}{\sqrt{1-\bar\alpha_t}}:

μ~t=1αt(xt−1−αt1−αˉt ϵˉt).\boxed{\tilde\mu_t = \frac{1}{\sqrt{\alpha_t}}\Big(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\,\bar\epsilon_t\Big).}

To go one step back, remove a well-chosen fraction of the noise that was put in, then undo the shrinking by αt\sqrt{\alpha_t}.

3.2.5 The simplified objective

The network sees only xtx_t and tt, not x0x_0, so it cannot compute ϵˉt\bar\epsilon_t; it predicts it, and we use the same formula: μθ(xt,t)=1αt(xt−1−αt1−αˉt ϵθ(xt,t))\mu_\theta(x_t,t) = \frac{1}{\sqrt{\alpha_t}}\big(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon_\theta(x_t,t)\big). With σt2=β~t\sigma_t^2 = \tilde\beta_t, the log term and the β~t2σt2−12\frac{\tilde\beta_t}{2\sigma_t^2} - \frac12 terms of 3.2.2 vanish, leaving

DKL=(μ~t−μθ)22β~t,μ~t−μθ=−1−αtαt1−αˉt(ϵˉt−ϵθ(xt,t)).D_{KL} = \frac{(\tilde\mu_t - \mu_\theta)^2}{2\tilde\beta_t}, \qquad \tilde\mu_t - \mu_\theta = -\frac{1-\alpha_t}{\sqrt{\alpha_t}\sqrt{1-\bar\alpha_t}}\big(\bar\epsilon_t - \epsilon_\theta(x_t,t)\big).

So

DKL=(1−αt)22 αt(1−αˉt) β~t (ϵˉt−ϵθ)2=(1−αt)22 αt(1−αˉt)⋅αt(1−αˉt)(αt−αˉt)(1−αt) (ϵˉt−ϵθ)2,D_{KL} = \frac{(1-\alpha_t)^2}{2\,\alpha_t(1-\bar\alpha_t)\,\tilde\beta_t}\,\big(\bar\epsilon_t - \epsilon_\theta\big)^2 = \frac{(1-\alpha_t)^2}{2\,\alpha_t(1-\bar\alpha_t)}\cdot\frac{\alpha_t(1-\bar\alpha_t)}{(\alpha_t - \bar\alpha_t)(1-\alpha_t)}\,\big(\bar\epsilon_t - \epsilon_\theta\big)^2 , DKL(q(xt−1∣xt,x0) ∥ pθ(xt−1∣xt))=1−αt2 (αt−αˉt) (ϵˉt−ϵθ(xt,t))2.\boxed{D_{KL}\big(q(x_{t-1}\mid x_t,x_0)\,\|\,p_\theta(x_{t-1}\mid x_t)\big) = \frac{1-\alpha_t}{2\,(\alpha_t - \bar\alpha_t)}\,\big(\bar\epsilon_t - \epsilon_\theta(x_t,t)\big)^2 .}

(Equivalently 1−αt2αt(1−αˉt−1)(ϵˉt−ϵθ)2\frac{1-\alpha_t}{2\alpha_t(1-\bar\alpha_{t-1})}(\bar\epsilon_t - \epsilon_\theta)^2.) Training a diffusion model is therefore regression on the noise: sample x0x_0, tt and ϵˉt\bar\epsilon_t, form xt=αˉtx0+1−αˉtϵˉtx_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\bar\epsilon_t, and penalise the squared error of the network's guess. The weight in front depends only on tt:

The weight (1 - alpha_t) / (2 (alpha_t - alpha-bar_t)) on a log scale over t = 2 to 200 for a linear schedule: about 1.75 at t = 2, dropping below 0.1 by t = 12, then flat around 0.02 to 0.03.

The exact KL puts most of the weight on the first few steps (fine details). Ho et al. found that simply dropping it, i.e. minimising E ∥ϵˉt−ϵθ(xt,t)∥2\mathbb E\,\lVert\bar\epsilon_t - \epsilon_\theta(x_t,t)\rVert^2 with tt uniform, gives better samples: the "simple loss" every DDPM implementation uses.

Part 3.3: telling the network the time step

The same network must denoise slightly noisy inputs (small tt) and almost pure noise (large tt), so it needs tt as an input. Since ϵθ\epsilon_\theta must keep the shape (B,C,L0,…,LM−1)(B, C, L_0, \dots, L_{M-1}), the standard design (as in DDPM's U-Net) is:

  1. Represent tt by a vector, not a raw integer. Use a sinusoidal embedding, as for positions in transformers: for a dimension dd and frequencies ωk=10000−2k/d\omega_k = 10000^{-2k/d}, emb(t)=(sin⁡(ω0t),cos⁡(ω0t),…,sin⁡(ωd/2−1t),cos⁡(ωd/2−1t))∈Rd.\text{emb}(t) = \big(\sin(\omega_0 t), \cos(\omega_0 t), \dots, \sin(\omega_{d/2-1} t), \cos(\omega_{d/2-1} t)\big) \in \mathbb R^{d}. It is smooth in tt, bounded, distinguishes all steps, and lets the network read both coarse and fine time differences. A learned lookup table nn.Embedding(T, d) also works when TT is fixed. The embedding is then passed through a small MLP (for example Linear, SiLU, Linear) to get a learned time feature of shape (B,d′)(B, d').
  2. Inject it into every block, broadcast over space. In each residual block with C′C' feature channels, a block-specific linear layer maps the time feature to (B,C′)(B, C'), which is reshaped to (B,C′,1,…,1)(B, C', 1, \dots, 1) and added to the feature map (a per-channel bias that depends on tt), or used as a per-channel scale and shift h↦(1+γ(t)) h+β(t)h \mapsto (1 + \gamma(t))\,h + \beta(t) (FiLM / adaptive normalisation). Broadcasting makes this work for any number MM of spatial dimensions and keeps the output shape equal to the input shape.
  3. Simpler alternatives, less effective: append t/Tt/T as an extra constant input channel (shape (B,1,L0,… )(B, 1, L_0, \dots)), or feed emb(t)\text{emb}(t) as an extra token in a transformer backbone (as in DiT).

Injecting tt at every block, rather than only at the input, matters because each layer's job (how much to smooth, how much detail to keep) depends on the noise level.

Results

check (in solution.py)formulanumerical
E[x50∣x0=1.7]\mathbb E[x_{50}\mid x_0 = 1.7], Var⁡[x50∣x0]\operatorname{Var}[x_{50}\mid x_0] (Monte Carlo, 400 000 chains)1.4534, 0.26911.4538, 0.2690
x200x_{200} (linear schedule, αˉ200=0.006\bar\alpha_{200} = 0.006)≈N(0,1)\approx\mathcal N(0,1)variance 0.994
summable schedule αt=1−1/(t+1)2\alpha_t = 1 - 1/(t+1)^2αˉ∞=1/2\bar\alpha_\infty = 1/20.5000
Gaussian KL, three random pairsclosed formmatches integration to 10−610^{-6}
posterior mean and variance (t=30t = 30)0.294652, 0.006911same on a 400 001-point grid
3.2.4 and 3.2.5 on random (t,x0,ϵ,ϵθ)(t, x_0, \epsilon, \epsilon_\theta)identitieshold to machine precision

What to remember

  • Sums of independent Gaussians stay Gaussian and their variances add: that single fact gives the closed-form jump xt=αˉtx0+1−αˉtϵˉtx_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\bar\epsilon_t.
  • Gaussian posteriors are found by completing the square: precisions add, and the mean is a precision-weighted average of the information sources.
  • KL between Gaussians is a closed formula; with equal variances it is just the squared mean difference over twice the variance.
  • Reparameterising the mean by the noise turns a generative model into a regression problem: predict ϵˉt\bar\epsilon_t with a squared loss.
  • Conditioning a network on a scalar (time, noise level, class) is done with an embedding plus per-channel broadcasting into every block.
Full code (solution.py)
"""USAAIO 2026 Round 2 Day 2, Problem 3: 1-D diffusion models.

Numerical checks of every derived formula:
  3.1  forward marginals x_t | x_0 ~ N(sqrt(abar_t) x_0, 1 - abar_t)  (Monte Carlo)
  3.2.1 Gaussian KL in closed form vs numerical integration
  3.2.3 posterior q(x_{t-1} | x_t, x_0) vs brute-force Bayes on a grid
  3.2.4/3.2.5 the epsilon parameterisation and the simplified KL
plus the two figures of the correction.
Run:  python solution.py
"""
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

rng = np.random.default_rng(0)
T = 200
alpha = 1 - np.linspace(1e-4, 0.05, T)          # a typical (linear) schedule
abar = np.cumprod(alpha)                          # abar[t-1] = prod_{s<=t} alpha_s

def gauss(x, m, v):
    return np.exp(-(x - m) ** 2 / (2 * v)) / np.sqrt(2 * np.pi * v)

# ------------------------------------------------ 3.1 forward process (MC)
x0 = 1.7
n = 400_000
x = np.full(n, x0)
for t in range(1, T + 1):
    x = np.sqrt(alpha[t - 1]) * x + np.sqrt(1 - alpha[t - 1]) * rng.standard_normal(n)
    if t in (2, 50, 200):
        print(f"t={t:3d}: MC mean {x.mean():+.4f} vs {np.sqrt(abar[t-1]) * x0:+.4f} | "
              f"MC var {x.var():.4f} vs {1 - abar[t-1]:.4f}")
print("abar_T =", abar[-1].round(6), "-> x_T | x_0 is close to N(0, 1)")

# a schedule whose noise is summable: abar does NOT go to 0
k = np.arange(1, 100_001)
print("alpha_t = 1 - 1/(t+1)^2 : abar_inf ~", np.prod(1 - 1 / (k + 1.0) ** 2).round(4), "(= 1/2)")

# ------------------------------------------------ 3.2.1 KL of two Gaussians
def kl_gauss(mq, vq, mp, vp):
    return 0.5 * np.log(vp / vq) + (vq + (mq - mp) ** 2) / (2 * vp) - 0.5

grid = np.linspace(-30, 30, 600_001)
dx = grid[1] - grid[0]
for (mq, vq, mp, vp) in [(0, 1, 0, 1), (1.0, 0.5, -0.3, 2.0), (2.0, 3.0, 2.5, 0.7)]:
    q, p = gauss(grid, mq, vq), gauss(grid, mp, vp)
    num = np.sum(q * (np.log(q + 1e-300) - np.log(p + 1e-300))) * dx
    print(f"KL closed form {kl_gauss(mq, vq, mp, vp):.6f}  numerical {num:.6f}")

# ------------------------------------------------ 3.2.3 posterior by Bayes
def posterior(xt, x0, t):
    a, ab = alpha[t - 1], abar[t - 1]
    mu = ((a - ab) / np.sqrt(a) * xt + np.sqrt(ab) * (1 - a) / np.sqrt(a) * x0) / (1 - ab)
    beta = (a - ab) * (1 - a) / (a * (1 - ab))
    return mu, beta

t, xt, x0b = 30, 0.4, -1.2
a, ab_prev = alpha[t - 1], abar[t - 2]
g = np.linspace(-8, 8, 400_001)
unnorm = gauss(xt, np.sqrt(a) * g, 1 - a) * gauss(g, np.sqrt(ab_prev) * x0b, 1 - ab_prev)
dens = unnorm / (unnorm.sum() * (g[1] - g[0]))
m_num = (g * dens).sum() * (g[1] - g[0])
v_num = ((g - m_num) ** 2 * dens).sum() * (g[1] - g[0])
mu, beta = posterior(xt, x0b, t)
print(f"posterior mean {mu:.6f} (grid {m_num:.6f}), var {beta:.6f} (grid {v_num:.6f})")

# ------------------------------------------------ 3.2.4 / 3.2.5
for _ in range(5):
    t = rng.integers(2, T + 1)
    a, ab = alpha[t - 1], abar[t - 1]
    x0r, eps, eps_th = rng.normal(size=3)
    xtr = np.sqrt(ab) * x0r + np.sqrt(1 - ab) * eps
    mu_eps = (xtr - (1 - a) / np.sqrt(1 - ab) * eps) / np.sqrt(a)
    assert np.isclose(mu_eps, posterior(xtr, x0r, t)[0])
    mu_th = (xtr - (1 - a) / np.sqrt(1 - ab) * eps_th) / np.sqrt(a)
    beta = posterior(xtr, x0r, t)[1]
    kl = kl_gauss(posterior(xtr, x0r, t)[0], beta, mu_th, beta)
    assert np.isclose(kl, (1 - a) / (2 * (a - ab)) * (eps - eps_th) ** 2)
print("3.2.4 and 3.2.5 identities verified")

# ------------------------------------------------ figures
plt.rcParams.update({"font.size": 12})
# (1) forward process on a bimodal data distribution
n = 200_000
x = np.where(rng.random(n) < 0.5, -2.0, 2.0) + 0.25 * rng.standard_normal(n)
fig, ax = plt.subplots(figsize=(10, 5))
xs = np.linspace(-4.5, 4.5, 400)
cm = plt.get_cmap("viridis")
shown = [0, 10, 30, 60, 100, 200]
for t in range(0, T + 1):
    if t > 0:
        x = np.sqrt(alpha[t - 1]) * x + np.sqrt(1 - alpha[t - 1]) * rng.standard_normal(n)
    if t in shown:
        h, e = np.histogram(x, bins=160, range=(-4.5, 4.5), density=True)
        lbl = f"t = {t}" + ("" if t == 0 else f",  $\\bar\\alpha_t$ = {abar[t-1]:.3f}")
        ax.plot((e[:-1] + e[1:]) / 2, h, color=cm(shown.index(t) / (len(shown) - 1)), lw=2, label=lbl)
ax.plot(xs, gauss(xs, 0, 1), "k--", lw=1.5, label="N(0, 1)")
ax.set_xlabel("$x_t$"); ax.set_ylabel("density")
ax.set_title("Forward diffusion: two clusters dissolve into a standard Gaussian")
ax.legend(fontsize=10)
fig.tight_layout(); fig.savefig("figures/forward_process.png", dpi=150)

# (2) the weight in front of ||eps - eps_theta||^2 across t
tt = np.arange(2, T + 1)
w = (1 - alpha[tt - 1]) / (2 * (alpha[tt - 1] - abar[tt - 1]))
fig, ax = plt.subplots(figsize=(9, 4.5))
ax.plot(tt, w, lw=2)
ax.set_yscale("log"); ax.grid(alpha=0.3, which="both")
ax.set_xlabel("time step $t$")
ax.set_ylabel(r"$\frac{1-\alpha_t}{2(\alpha_t-\bar\alpha_t)}$")
ax.set_title(r"KL = weight$_t$ $\times\ (\bar\epsilon_t-\epsilon_\theta)^2$: the exact weight")
fig.tight_layout(); fig.savefig("figures/kl_weight.png", dpi=150)
print("weight at t=2: %.3g, t=10: %.3g, t=200: %.3g" % (w[0], w[8], w[-1]))