"""Original DDPM arithmetic lesson: 12-step schedule and tiny time-conditioned U-Net. The larger beta schedule (0.02..0.18) is for a short teaching chain, not the paper's 1000-step training schedule. Fixed posterior variance; epsilon prediction. """ import torch import torch.nn as nn import torch.nn.functional as F class ForwardNoise(nn.Module): def __init__(self): super().__init__() beta = torch.cat( [ torch.zeros(1, dtype=torch.float64), torch.linspace(0.02, 0.18, 12, dtype=torch.float64), ] ) self.register_buffer("beta", beta) self.register_buffer("alpha", 1 - beta) self.register_buffer("alpha_bar", (1 - beta).cumprod(0)) def forward(self, clean, noise, timestep): a = self.alpha_bar[timestep].to(clean.dtype).reshape(-1, 1, 1, 1) return a.sqrt() * clean + (1 - a).sqrt() * noise class NoisePredictor(nn.Module): def __init__(self): super().__init__() self.time_embedding = nn.Embedding(13, 8) self.stem = nn.Conv2d(1, 8, 3, padding=1) self.down = nn.Conv2d(8, 8, 3, stride=2, padding=1) self.bottleneck = nn.Conv2d(8, 8, 3, padding=1) self.up = nn.ConvTranspose2d(8, 8, 2, stride=2) self.merge = nn.Conv2d(16, 8, 3, padding=1) self.noise = nn.Conv2d(8, 1, 3, padding=1) def forward(self, noisy, timestep): time = self.time_embedding(timestep)[:, :, None, None] skip = F.silu(self.stem(noisy) + time) coarse = F.silu(self.bottleneck(F.silu(self.down(skip)) + time)) fused = F.silu(self.merge(torch.cat([skip, self.up(coarse)], dim=1))) return self.noise(fused) class TinyDDPM(nn.Module): def __init__(self): super().__init__() self.corrupt = ForwardNoise() self.denoiser = NoisePredictor() def forward(self, clean, noise, timestep): noisy = self.corrupt(clean, noise, timestep) return self.denoiser(noisy, timestep) def reverse_step(schedule, noisy, predicted_noise, t, noise): beta, alpha, abar = schedule.beta[t], schedule.alpha[t], schedule.alpha_bar[t] mean = (noisy - beta / (1 - abar).sqrt() * predicted_noise) / alpha.sqrt() variance = beta * (1 - schedule.alpha_bar[t - 1]) / (1 - abar) return mean + (variance.sqrt() * noise if t > 1 else 0), mean, variance def setup(): torch.manual_seed(37) model = TinyDDPM().double() clean = torch.zeros(1, 1, 8, 8, dtype=torch.float64) clean[:, :, 2:6, 2:6] = 1 clean[:, :, 3:5, 3:5] = -1 noise = torch.randn_like(clean) return model, clean, noise def verify(): model, clean, noise = setup() s = model.corrupt assert torch.all(s.alpha_bar[1:] < s.alpha_bar[:-1]) for t in [1, 6, 12]: a = s.alpha_bar[t] xt = s(clean, noise, torch.tensor([t])) recovered = (xt - (1 - a).sqrt() * noise) / a.sqrt() torch.testing.assert_close(recovered, clean) _, mean, var = reverse_step(s, xt, noise, t, torch.zeros_like(clean)) posterior = ( s.beta[t] * s.alpha_bar[t - 1].sqrt() / (1 - a) * clean + s.alpha[t].sqrt() * (1 - s.alpha_bar[t - 1]) / (1 - a) * xt ) torch.testing.assert_close(mean, posterior, atol=2e-7, rtol=2e-6) assert var >= 0 x1 = s(clean, noise, torch.tensor([1])) a, _, _ = reverse_step(s, x1, noise, 1, torch.zeros_like(clean)) b, _, _ = reverse_step(s, x1, noise, 1, torch.ones_like(clean) * 99) torch.testing.assert_close(a, b, rtol=0, atol=0) torch.testing.assert_close(a, clean, atol=2e-7, rtol=2e-6) t = torch.tensor([6]) predicted = model(clean, noise, t) assert not torch.allclose( model.denoiser(s(clean, noise, t), torch.tensor([1])), predicted ) opt = torch.optim.SGD(model.denoiser.parameters(), lr=0.001) loss = F.mse_loss(predicted, noise) opt.zero_grad() loss.backward() assert model.denoiser.time_embedding.weight.grad[6].norm() > 0 opt.step() assert F.mse_loss(model(clean, noise, t), noise) < loss return [ "Cumulative signal decreases; known noise reconstructs the clean image at early, middle and late steps", "Epsilon-parameterized reverse means match the analytic forward posterior", "The last reverse step adds no stochastic noise and recovers x0 with oracle epsilon", "Time conditioning affects predictions and receives gradients; one fixed-batch SGD step lowers epsilon MSE", ] def experiment(): model, clean, noise = setup() s = model.corrupt cases = [] with torch.no_grad(): for t in [1, 6, 12]: xt = s(clean, noise, torch.tensor([t])) pred = model.denoiser(xt, torch.tensor([t])) a = s.alpha_bar[t] estimate = (xt - (1 - a).sqrt() * pred) / a.sqrt() oracle = (xt - (1 - a).sqrt() * noise) / a.sqrt() cases.append( dict( id=f"noise-{t}", label=f"Forward corruption · t = {t}", target="corrupt", note="The same clean pattern and same Gaussian noise are used at each displayed t to isolate the signal/noise coefficients. These are coupled marginal examples, not consecutive forward-chain samples. The untrained prediction is compared with the actual added noise; oracle reconstruction uses known noise unavailable during generation.", matrices=[ dict(label="Clean x₀", values=clean[0, 0].tolist()), dict(label=f"Noisy x{t}", values=xt[0, 0].tolist()), dict( label="Untrained estimate of x₀", values=estimate[0, 0].tolist(), ), ], vectors=[ dict(label="Signal coefficient √ᾱₜ", values=[a.sqrt().item()]), dict( label="Noise coefficient √(1−ᾱₜ)", values=[(1 - a).sqrt().item()], ), ], metrics=[ dict( label="Noise prediction MSE", value=F.mse_loss(pred, noise).item(), ), dict( label="Oracle x₀ maximum error", value=(oracle - clean).abs().max().item(), ), ], ) ) generator = torch.Generator().manual_seed(41) current = torch.randn(clean.shape, generator=generator, dtype=clean.dtype) start = current.clone() for t in range(12, 0, -1): pred = model.denoiser(current, torch.tensor([t])) z = torch.randn(clean.shape, generator=generator, dtype=clean.dtype) current, mean, variance = reverse_step(s, current, pred, t, z) if t in [9, 5, 1]: cases.append( dict( id=f"reverse-{t-1}", label=f"Reverse chain · reached t = {t-1}", target="denoiser.noise", note=f"Start from an independent Gaussian sample and reuse the same untrained denoiser for all 12 reverse steps. This records the state after step {t}. At t = 1 no new noise is added. The result is an arithmetic demonstration, not a generated image from a trained DDPM.", matrices=[ dict( label="Starting Gaussian x₁₂", values=start[0, 0].tolist(), ), dict( label=f"Current sample x{t-1}", values=current[0, 0].tolist(), ), dict( label="Reverse mean at this step", values=mean[0, 0].tolist(), ), ], vectors=[ dict( label="Completed reverse timesteps", values=list(range(12, t - 1, -1)), ) ], metrics=[ dict( label="Posterior variance at this step", value=variance.item(), ), dict(label="Remaining denoiser calls", value=t - 1), ], ) ) return dict( kind="matrices", title="Corrupt once; denoise repeatedly.", description="Recorded forward marginals and a separate 12-step reverse chain. The denoiser is untrained; the panel teaches DDPM scheduling, epsilon prediction and posterior sampling.", controlLabel="Diffusion checkpoint", cases=cases, )