§3.4, Eq. 14 — Simplified training objective L_simple. Computes MSE between true noise and predicted noise: L = ||ε − ε_θ(x_t, t)||² This module handles only the loss computation. The caller (training loop) is responsible for sampling t, computing x_t from x_0, and calling the
| 28 | |
| 29 | |
| 30 | class DDPMLoss(nn.Module): |
| 31 | """§3.4, Eq. 14 — Simplified training objective L_simple. |
| 32 | |
| 33 | Computes MSE between true noise and predicted noise: |
| 34 | L = ||ε − ε_θ(x_t, t)||² |
| 35 | |
| 36 | This module handles only the loss computation. The caller (training loop) |
| 37 | is responsible for sampling t, computing x_t from x_0, and calling the model. |
| 38 | """ |
| 39 | |
| 40 | def __init__(self): |
| 41 | super().__init__() |
| 42 | |
| 43 | def forward( |
| 44 | self, |
| 45 | noise_pred: torch.Tensor, |
| 46 | noise_true: torch.Tensor, |
| 47 | ) -> torch.Tensor: |
| 48 | """ |
| 49 | §3.4, Eq. 14 — L_simple = E[||ε − ε_θ(x_t, t)||²] |
| 50 | |
| 51 | Args: |
| 52 | noise_pred: (batch, C, H, W) — predicted noise ε_θ(x_t, t) |
| 53 | noise_true: (batch, C, H, W) — true noise ε ~ N(0, I) |
| 54 | |
| 55 | Returns: |
| 56 | scalar — mean squared error loss |
| 57 | """ |
| 58 | # §3.4 — Simple MSE between predicted and true noise |
| 59 | # "equivalent to (a re-weighted variant of) the ELBO" |
| 60 | return nn.functional.mse_loss(noise_pred, noise_true) |