Perception Prioritized Training of Diffusion Models

TL;DR

This work addresses the suboptimality of standard fixed-weight training objectives in diffusion models by proposing Perception Prioritized (P2) weighting. P2 adjusts the loss weights across diffusion steps based on the current SNR, emphasizing noise levels that foster learning of perceptually rich content and de-emphasizing corollary learning of imperceptible details. Empirical results across diverse datasets and sampling configurations show consistent improvements in FID and KID, with state-of-the-art performance on Oxford Flowers and CelebA-HQ, and strong results on FFHQ. The approach offers a simple, broadly applicable modification to training objectives that enhances perceptual quality without added computational burden, suggesting wide applicability in diffusion-based generation tasks.

Abstract

Diffusion models learn to restore noisy data, which is corrupted with different levels of noise, by optimizing the weighted sum of the corresponding loss terms, i.e., denoising score matching loss. In this paper, we show that restoring data corrupted with certain noise levels offers a proper pretext task for the model to learn rich visual concepts. We propose to prioritize such noise levels over other levels during training, by redesigning the weighting scheme of the objective function. We show that our simple redesign of the weighting scheme significantly improves the performance of diffusion models regardless of the datasets, architectures, and sampling strategies.

Figures (11)

• Figure 1: Information removal of a diffusion process. (Left) Perceptual distance of corrupted images as a function of signal-to-noise ratio (SNR). Distances are measured between two noisy images either corrupted from the same image (blue) or different images (orange). We averaged distances measured with 200 random triplets from CelebA-HQ. Perceptually recognizable contents are removed when SNR magnitude is between $10^{-2}$ and $10^0$. (Right) Illustration of the diffusion process.
• Figure 2: Stochastic reconstruction. (Left) Illustration of reconstruction, where sample are obtained from full sampling chain. (Right) Reconstructions $\hat{x}_0$ with input images $x_0$ on the rightmost column and SNR of $x_t$ on the bottom. Samples in the 1st, 2nd columns share only the coarse attributes (e.g., global color structure) with the input. The 3rd, 4th columns share perceptually discriminative contents with the input. 5th column are almost identical to the input, suggesting that the model learns imperceptible details when SNRs are large.
• Figure 3: Weighting schemes. (Left) Signal-to-noise ratio (SNR) of linear and cosine noise schedules for reference. (Middle) Weights of our P2 weighting and the baseline with a cosine schedule. (Right) Weights of P2 weighting and the baseline with a linear schedule. Compared to the baseline, P2 weighting suppresses weights for large SNRs, where the model learns imperceptible details.
• Figure 4: Comparison of FID-10k through training on FFHQ. P2 weighting consistently improves performance for both linear and cosine schedules. Training progress refers to the number of images seen by the model. Samples are generated with 250 steps.
• Figure 5: Samples generated by our models trained on several datasets (FFHQ, CelebA-HQ, MetFaces, AFHQ-Dogs, Oxford Flowers, CUB Bird) at 256$\times$256 resolution. See appendix for more samples.