key = jax.random.PRNGKey(0)
mix = d2l.GaussianMixture(means=[[-2.5, -1.5], [2.5, -1.5], [0.0, 2.5]],
weights=[0.5, 0.3, 0.2], std=0.5)
key, k1, k2 = jax.random.split(key, 3)
data, test = mix.sample(k1, 2000), mix.sample(k2, 2000)
T = 200
beta = jnp.linspace(1e-4, 0.1, T)
alpha = 1 - beta
alpha_bar = jnp.cumprod(alpha)
print(f'alpha_bar_T = {alpha_bar[-1]:.1e}: after T steps the data are scaled '
f'by {jnp.sqrt(alpha_bar[-1]):.3f}')
x = data
for t in range(T): # the chain, step by step
key, subkey = jax.random.split(key)
x = jnp.sqrt(alpha[t]) * x + jnp.sqrt(beta[t]) * jax.random.normal(subkey, x.shape)
if t + 1 in (10, 50, 200):
key, subkey = jax.random.split(key)
one_shot = (jnp.sqrt(alpha_bar[t]) * data
+ jnp.sqrt(1 - alpha_bar[t]) * jax.random.normal(subkey, data.shape))
print(f'step {t + 1:3d}: chain: mean {[round(float(m), 3) for m in x.mean(0)]}, '
f'var {x.var(0).mean():.3f} | one shot: mean '
f'{[round(float(m), 3) for m in one_shot.mean(0)]}, '
f'var {one_shot.var(0).mean():.3f}')Denoising Diffusion Probabilistic Models
Dive into Deep Learning · §17.7
Denoising diffusion probabilistic models
a fixed forward chain · a learned reverse chain · a variational bound that becomes noise prediction
The Forward Process Turns Data into Noise in T Steps
q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathcal{N}\big(\sqrt{1 - \beta_t}\, \mathbf{x}_{t-1},\ \beta_t I\big), \qquad q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathcal{N}\big(\sqrt{\bar{\alpha}_t}\, \mathbf{x}_0,\ (1 - \bar{\alpha}_t) I\big)
Variance preserving: the chain ends close to a standard Gaussian, whatever the data scale, once \sqrt{\bar{\alpha}_T}\, \mathbf{x}_0 is negligible.
The Closed-Form Marginal, Checked by Simulation
alpha_bar_T = 3.2e-05: after T steps the data are scaled by 0.006
step 10: chain: mean [-0.546, -0.751], var 3.761 | one shot: mean [-0.549, -0.758], var 3.775
step 50: chain: mean [-0.409, -0.601], var 2.489 | one shot: mean [-0.396, -0.531], var 2.500
step 200: chain: mean [-0.005, -0.014], var 1.000 | one shot: mean [-0.001, -0.029], var 0.994Any step is one Gaussian draw from the data; training never simulates the chain.
The Reverse Process Is a Latent-Variable Model
p_{\boldsymbol{\theta}}(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \prod_{t=1}^{T} p_{\boldsymbol{\theta}}(\mathbf{x}_{t-1} \mid \mathbf{x}_t), \qquad p_{\boldsymbol{\theta}}(\mathbf{x}_{t-1} \mid \mathbf{x}_t) = \mathcal{N}\big(\boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{x}_t, t),\ \sigma_t^2 I\big)
- small steps: the reversal of a Gaussian step is nearly Gaussian;
- \mathbf{x}_{1:T} are latents; the likelihood of \mathbf{x}_0 is an intractable integral;
- so bound it, with the forward process as the auxiliary distribution.
The Bound Splits into One Term per Step
L_{\textrm{VLB}} = \underbrace{\mathrm{KL}\big(q(\mathbf{x}_T \mid \mathbf{x}_0)\, \|\, p(\mathbf{x}_T)\big)}_{L_T} + \sum_{t=2}^{T} \underbrace{\mathbb{E}_q\,\mathrm{KL}\big(q(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0)\, \|\, p_{\boldsymbol{\theta}}(\mathbf{x}_{t-1} \mid \mathbf{x}_t)\big)}_{L_{t-1}} + \underbrace{\mathbb{E}_q[-\log p_{\boldsymbol{\theta}}(\mathbf{x}_0 \mid \mathbf{x}_1)]}_{L_0}
- L_T: no parameters, near zero by construction;
- L_{t-1}: match the reversal of step t given the clean point, a Gaussian with known mean and variance;
- L_0: the last denoising step.
Noise Prediction
\tilde{\boldsymbol{\mu}}_t = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\, \boldsymbol{\epsilon} \Big) \quad\Rightarrow\quad \boldsymbol{\mu}_{\boldsymbol{\theta}} = \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\, \boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_t, t) \Big)
L_{\textrm{simple}} = \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\epsilon}} \big\| \boldsymbol{\epsilon} - \boldsymbol{\epsilon}_{\boldsymbol{\theta}}\big(\sqrt{\bar{\alpha}_t}\, \mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\, \boldsymbol{\epsilon},\ t\big) \big\|^2
Draw \mathbf{x}_0, t, \boldsymbol{\epsilon}; noise the point in one shot; regress. The bound’s per-step weights are dropped.
It Is Denoising Score Matching, and the Sampler Is Langevin-Like
\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t, t) = -\frac{\boldsymbol{\epsilon}_{\boldsymbol{\theta}}(\mathbf{x}_t, t)}{\sqrt{1 - \bar{\alpha}_t}}, \qquad \mathbf{x}_{t-1} = \frac{1}{\sqrt{\alpha_t}} \big( \mathbf{x}_t + \beta_t\, \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t, t) \big) + \sigma_t \mathbf{z}
A score step of size \beta_t, noise \sigma_t, and a rescaling: a Langevin-like step whose drift is doubled and whose schedule the forward process fixes.
The Learned Predictor against the Optimum
def eps_star(x_t, t): # the optimal noise predictor at step index t
a = alpha_bar[t]
return -jnp.sqrt(1 - a) * mix.score(x_t, scale=jnp.sqrt(a),
sigma=jnp.sqrt(1 - a))
print(f'{"step":>4} {"E||eps_theta - eps*||^2":>24} {"E||eps - eps*||^2":>18} '
f'{"E||eps - eps_theta||^2":>23}')
for t in (9, 49, 99, 199):
key, subkey = jax.random.split(key)
eps = jax.random.normal(subkey, test.shape)
x_t = jnp.sqrt(alpha_bar[t]) * test + jnp.sqrt(1 - alpha_bar[t]) * eps
pred, best = net(x_t, jnp.full((len(test),), t)), eps_star(x_t, t)
print(f'{t + 1:>4} {((pred - best) ** 2).sum(1).mean():>24.3f} '
f'{((eps - best) ** 2).sum(1).mean():>18.3f} '
f'{((eps - pred) ** 2).sum(1).mean():>23.3f}')step E||eps_theta - eps*||^2 E||eps - eps*||^2 E||eps - eps_theta||^2
10 0.058 1.779 1.838
50 0.049 0.784 0.831
100 0.010 0.480 0.492
200 0.006 0.000 0.007A mean squared gap to \boldsymbol{\epsilon}^\star of at most a few hundredths at each step checked; the loss exceeds the irreducible noise variance by less than a tenth.
Ancestral Sampling Recovers the Data
Two hundred steps from a Gaussian cloud: no structure at half the schedule, the mixture’s layout by t = 50, detail at the end. The sampler adds no parameters beyond the forward schedule and \sigma_t^2 = \beta_t.
Recap
- Forward: a fixed variance-preserving chain with closed-form marginals.
- Reverse: Gaussian transitions with learned means; the bound splits into per-step Kullback–Leibler terms, and training uses its unweighted simplification.
- Noise prediction turns every term into a regression; the simple loss is reweighted denoising score matching.
- Sampling is Langevin-like with a fixed schedule. Next: the same model on images, with a U-Net and class conditioning.