Mean squared error (MSE) averages the squared differences between predicted and reference numerical values. In diffusion training, it compares noise predicted by a denoising U-Net with the noise actually injected into images. section 6.3 section 6.3.2
A DDPM training step uses MSE in a short sequence:
- Sample a time step and add noise to a batch of clean images. section 6.3.2
- Give the noisy images, labels, and time information to the denoising U-Net. section 6.3
- Compute MSE between injected noise and predicted noise, then adjust model parameters to reduce the loss. section 6.3.3
Fashion-MNIST diffusion training
A conditional DDPM trains for 20 epochs on Fashion-MNIST images. Each batch supplies clean clothing images and labels; the model predicts the added noise, and MSE supplies the loss used during optimization. section 6.3.3
Listing 6.5 Creating the DDPM
device="cuda" if torch.cuda.is_available() else "cpu" class DDPM(nn.Module): def __init__(self, model, n_T, device=device, drop_prob=0.1): super().__init__() self.model = model.to(device) for k, v in noise_scheduler(n_T).items(): self.register_buffer(k, v) self.n_T = n_T self.device = device self.drop_prob = drop_prob self.loss_mse = nn.MSELoss() def forward(self, x, c): _ts=torch.randint(1,self.n_T+1, (x.shape[0],)).to(self.device) #1 noise = torch.randn_like(x) x_t = (self.sqrtab[_ts, None, None, None] * x + self.sqrtmab[_ts, None, None, None] * noise) #2 context_mask = torch.bernoulli(torch.zeros_like(c)+ self.drop_prob).to(self.device) return self.loss_mse(noise, self.model(x_t, c, _ts / self.n_T, context_mask)) #3
Inspect the DDPM implementation to see MSE instantiated as nn.MSELoss and applied to predicted and actual noise. section 6.3.2
