"""DDPM, the reverse-time equation written on a finite grid. A U-Net reads a corrupted image and the step it sits at, and regresses the Gaussian draw that corrupted it. The reverse transition rebuilds the previous state from that one prediction, with a mean that comes from completing the square and the posterior variance as its spread. The sampler starts at a standard Gaussian and applies that transition K times. The three pieces match the derivation one to one, and we keep the chapter's names. `forward_diffusion.py` covers the schedule and the forward corruption, and we rebuild both here so that this file stands alone. In Chapter 3 step k = 0 holds a clean image and corruption grows with k, up to pure noise at k = K. Chapter 3, Sections 3.19, 3.20 and 3.24 through 3.28 of the notes. Trained on 28-by-28 images, the offline toy set by default and MNIST with --dataset mnist. Prints L_simple over the run, the reverse posterior variance and the ELBO weight formed from it, the chapter's two-step reverse transition, and the range of the generated images. python lecture_3/ddpm.py # 300 steps on the offline images python lecture_3/ddpm.py --steps 20 # a few seconds, for a first look python lecture_3/ddpm.py --dataset mnist # the real digits, downloaded on first use """ # %% 01. Imports and configuration import math import torch import torch.nn as nn import torch.nn.functional as F from cis6270.nets import count_parameters from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train NUM_STEPS = 1000 # Diffusion steps K. BETA_START = 1e-4 # Noise variance added at step 1. BETA_END = 0.02 # Noise variance added at step K. # %% 02. The schedule tensors and the forward corruption # # alpha_k = 1 - beta_k is the one-step retention, abar_k its running product, and beta-tilde_k the # variance of the reverse posterior q(x_{k-1} | x_k, x_0). We prefix a zero to the betas, so # Python index k holds beta_k and the cumulative product starts at abar_0 = 1, the clean state. # One Gaussian draw then reaches any level, x_k = a_k x_0 + sigma_k eps with a_k = sqrt(abar_k) and # sigma_k = sqrt(1 - abar_k). betas = torch.cat([ torch.zeros(1), # Placeholder so that index k holds beta_k. torch.linspace(BETA_START, BETA_END, NUM_STEPS), # beta_1 ... beta_K. ]) alphas = 1.0 - betas # One-step signal retention alpha_k. alphas_cumprod = torch.cumprod(alphas, dim=0) # Cumulative retention. alphas_cumprod_prev = torch.cat([ torch.ones(1), # abar_0 = 1. alphas_cumprod[:-1], # Index k now holds abar_{k-1}. ]) posterior_variance = ( betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod) ) # Reverse variance beta-tilde_k. def extract(schedule, k, shape): # One schedule value per item, reshaped to broadcast over C, H, W. out = schedule.to(k.device).gather(0, k) return out.reshape(k.shape[0], *((1,) * (len(shape) - 1))) def q_sample(x_0, k, noise=None): if noise is None: noise = torch.randn_like(x_0) # Mean 0, covariance I, shape of x_0. abar = extract(alphas_cumprod, k, x_0.shape) a_k = abar.sqrt() # Signal coefficient a_k. sigma_k = (1.0 - abar).sqrt() # Noise standard deviation sigma_k. return a_k * x_0 + sigma_k * noise, noise # x_k and the exact noise. # %% 03. Sinusoidal step embedding # # A raw index is a poor input to a linear layer, so we spread the step over many frequencies # first. def timestep_embedding(k, dim): half = dim // 2 # Half the coordinates sine, half cosine. freqs = torch.exp( -math.log(10000.0) * torch.arange(half, device=k.device).to(torch.get_default_dtype()) / half ) args = k.to(torch.get_default_dtype())[:, None] * freqs[None, :] # [B, half]. return torch.cat([torch.sin(args), torch.cos(args)], dim=1) # [B, dim]. def num_groups(channels): for g in (8, 4, 2, 1): # Largest valid GroupNorm group count. if channels % g == 0: return g # %% 04. The time-conditioned residual block # # The step embedding enters as a per-channel offset between the two convolutions, so every layer # has the current noise level. class ResidualBlock(nn.Module): def __init__(self, in_channels, out_channels, time_embed_dim): super().__init__() self.norm1 = nn.GroupNorm(num_groups(in_channels), in_channels) self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) self.time_proj = nn.Linear(time_embed_dim, out_channels) self.norm2 = nn.GroupNorm(num_groups(out_channels), out_channels) self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) self.skip = ( nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, kernel_size=1) ) self.activation = nn.SiLU() def forward(self, x, t_emb): h = self.conv1(self.activation(self.norm1(x))) # GN, SiLU, Conv_1. h = h + self.time_proj(t_emb)[:, :, None, None] # Offset per channel. h = self.conv2(self.activation(self.norm2(h))) return h + self.skip(x) # Residual connection. # %% 05. The U-Net # # Three scales, 28 by 28 down to 7 by 7 and back, with a skip at each resolution. The output has # the shape of the input because it predicts one noise coordinate per pixel. class UNet(nn.Module): def __init__(self, in_channels=1, base_channels=64, time_embed_dim=128): super().__init__() self.time_embed_dim = time_embed_dim self.time_mlp = nn.Sequential( # Learned representation of step k. nn.Linear(time_embed_dim, time_embed_dim), nn.SiLU(), nn.Linear(time_embed_dim, time_embed_dim), ) c1 = base_channels # 64 channels at 28 x 28. c2 = base_channels * 2 # 128 channels at 14 x 14. c3 = base_channels * 4 # 256 channels at 7 x 7. self.in_conv = nn.Conv2d(in_channels, c1, kernel_size=3, padding=1) self.down1 = ResidualBlock(c1, c1, time_embed_dim) self.pool1 = nn.Conv2d(c1, c1, kernel_size=3, stride=2, padding=1) self.down2 = ResidualBlock(c1, c2, time_embed_dim) self.pool2 = nn.Conv2d(c2, c2, kernel_size=3, stride=2, padding=1) self.mid = ResidualBlock(c2, c3, time_embed_dim) self.up1 = nn.ConvTranspose2d(c3, c2, kernel_size=4, stride=2, padding=1) self.dec1 = ResidualBlock(c2 + c2, c2, time_embed_dim) self.up2 = nn.ConvTranspose2d(c2, c1, kernel_size=4, stride=2, padding=1) self.dec2 = ResidualBlock(c1 + c1, c1, time_embed_dim) self.out_conv = nn.Conv2d(c1, in_channels, kernel_size=1) def forward(self, x, k): t_emb = self.time_mlp( timestep_embedding(k, self.time_embed_dim) ) # Learned representation of diffusion step k. h = self.in_conv(x) h1 = self.down1(h, t_emb) # 28 x 28 encoder features. h = self.pool1(h1) # 28 -> 14. h2 = self.down2(h, t_emb) # 14 x 14 encoder features. h = self.pool2(h2) # 14 -> 7. h = self.mid(h, t_emb) # Compressed representation. h = self.up1(h) # 7 -> 14. h = self.dec1(torch.cat([h, h2], dim=1), t_emb) # Join 14 x 14 skip. h = self.up2(h) # 14 -> 28. h = self.dec2(torch.cat([h, h1], dim=1), t_emb) # Join 28 x 28 skip. return self.out_conv(h) # [B, 1, 28, 28] predicted noise. # %% 06. One reverse step and the sampler @torch.no_grad() def p_sample(model, x_k, k): eps_theta = model(x_k, k) # Predicted noise at step k. beta_k = extract(betas, k, x_k.shape) alpha_k = extract(alphas, k, x_k.shape) abar_k = extract(alphas_cumprod, k, x_k.shape) mean = ( x_k - beta_k / (1.0 - abar_k).sqrt() * eps_theta ) / alpha_k.sqrt() # Reverse mean mu_theta. var = extract(posterior_variance, k, x_k.shape) # beta-tilde_k. z = torch.randn_like(x_k) # Fresh Gaussian draw. keep = (k > 1).to(torch.get_default_dtype()).reshape(-1, *((1,) * (x_k.dim() - 1))) return mean + keep * var.sqrt() * z # No noise at the final step. @torch.no_grad() def sample(model, num_images, device): x = torch.randn(num_images, 1, 28, 28, device=device) # X_K ~ N(0, I). for step in range(NUM_STEPS, 0, -1): # k = K down to k = 1. k = torch.full((num_images,), step, device=device, dtype=torch.long) x = p_sample(model, x, k) return x # Generated images in the training range [-1, 1]. # %% 07. The dropped weight and one complete reverse transition # # One interior term of the variational bound is beta_k^2 / (2 beta-tilde_k alpha_k (1 - abar_k)) # multiplied by the squared noise error, and L_simple sets every such factor to one. On the # two-step schedule beta_1 = beta_2 = 0.2 the whole transition is small enough to check by hand. # The mean there has two equal forms, one built from the noise prediction and one from the clean # estimate it implies. def elbo_weight(beta_k, posterior, alpha_k, abar_k) -> float: """The schedule constant in front of the squared noise error at one step.""" return beta_k ** 2 / (2.0 * posterior * alpha_k * (1.0 - abar_k)) def two_step_transition(beta: float, x_k: float, eps_theta: float, z: float) -> dict: """The reverse step at k = 2 of the two-step schedule, through both forms of the mean.""" alpha = 1.0 - beta # alpha_2. abar_1, abar_2 = alpha, alpha ** 2 # 0.8 and 0.64 at beta = 0.2. sigma = math.sqrt(1.0 - abar_2) # 0.6, the corruption standard deviation at k = 2. mean = (x_k - beta / sigma * eps_theta) / math.sqrt(alpha) # The noise form of mu_theta. x_0_hat = (x_k - sigma * eps_theta) / math.sqrt(abar_2) # The clean estimate it implies. weights = (math.sqrt(abar_1) * beta / (1.0 - abar_2), # Weight on the clean estimate. math.sqrt(alpha) * (1.0 - abar_1) / (1.0 - abar_2)) # Weight on x_k. variance = beta * (1.0 - abar_1) / (1.0 - abar_2) # beta-tilde_2 = 1/9. return { "mean_from_noise": mean, "mean_from_clean": weights[0] * x_0_hat + weights[1] * x_k, "x_0_hat": x_0_hat, "weights": weights, "variance": variance, "x_{k-1}": mean + math.sqrt(variance) * z, "elbo_weight": elbo_weight(beta, variance, alpha, abar_2), } # %% 08. The report def main() -> None: parser = common_parser(__doc__.split("\n")[0]) parser.add_argument("--dataset", choices=["toy", "mnist"], default="toy") parser.add_argument("--base-channels", type=int, default=8, help="U-Net base width; the notes uses 64") parser.set_defaults(steps=300, batch_size=64, lr=2e-3, samples=4) args = parser.parse_args() set_seed(args.seed) device = pick_device(args.device) if args.dataset == "toy": from cis6270.data import toy_images images = toy_images(512, 28, args.seed).to(device) # [512, 1, 28, 28] offline images. draw = lambda n: images[torch.randint(0, len(images), (n,), device=device)] # noqa: E731 else: from cis6270.data import mnist_loader loader = mnist_loader(batch_size=args.batch_size) stream = iter(loader) def draw(n): # One MNIST batch per call, restarting the loader at the end of an epoch. nonlocal stream batch = next(stream, None) if batch is None: stream = iter(loader) batch = next(stream) return batch[0].to(device) model = UNet(in_channels=1, base_channels=args.base_channels, time_embed_dim=128).to(device) if not args.quiet: print(f"device={device} dataset={args.dataset} parameters={count_parameters(model):,}") def loss_fn(step: int) -> torch.Tensor: x_0 = draw(args.batch_size) # [B, 1, 28, 28] clean images in [-1, 1]. k = torch.randint(1, NUM_STEPS + 1, (x_0.shape[0],), device=device) # One level per item. x_k, noise = q_sample(x_0, k) # Corrupt directly to step k. return F.mse_loss(model(x_k, k), noise) # L_simple on this batch. # The notes trains at 2e-4 for 20 epochs; 2e-3 for a few hundred steps is the # classroom version of the same run and reaches a usable sampler in under a minute. losses = train(model, loss_fn, args.steps, args.lr, quiet=args.quiet) images_out = sample(model, args.samples, device) # [S, 1, 28, 28] after K reverse steps. window = min(100, args.steps) # A short run averages over everything it has. reverse = two_step_transition(0.2, 7.4, -1.0, 0.3) weight_500 = elbo_weight(betas[500], posterior_variance[500], alphas[500], alphas_cumprod[500]) # The dropped factor at k = 500. report = { "parameters": count_parameters(model), f"loss, first {window} steps": sum(losses[:window]) / window, f"loss, last {window} steps": sum(losses[-window:]) / window, "beta-tilde at k=500": f"{float(posterior_variance[500]):.5f}", # 0.01003. # The weight L_simple drops, on the thousand-step grid and on the two-step hand schedule. "ELBO weight, k=500 and two-step": [round(float(weight_500), 6), round(reverse["elbo_weight"], 4)], # 0.005508, 0.625. "two-step reverse mean, two forms": [round(reverse["mean_from_noise"], 4), round(reverse["mean_from_clean"], 4)], # 8.646. "two-step x_0 estimate and weights": [round(reverse["x_0_hat"], 4)] + [round(w, 4) for w in reverse["weights"]], "two-step reverse variance and x_1": [round(reverse["variance"], 4), round(reverse["x_{k-1}"], 4)], # 1/9 and 8.746. "generated range": [round(float(images_out.min()), 3), round(float(images_out.max()), 3)], "generated standard deviation": float(images_out.std()), } print_report(f"DDPM on {args.dataset} images", report) saved = save_run(args.out, model=model, config=vars(args), report=report, losses=losses) if saved is not None: try: # Without torchvision the run still trains, samples and saves its weights. from torchvision.utils import save_image display = ((images_out.cpu() + 1.0) / 2.0).clamp(0.0, 1.0) # [-1, 1] -> [0, 1]. save_image(display, saved / "samples.png", nrow=4) # A grid of the generated images. except ImportError: pass # %% 09. Run the script if __name__ == "__main__": main()