Source-linked AI summary
Perception Prioritized Training of Diffusion Models
Jooyoung Choi, Jungbeom Lee, Chaehun Shin, Sungwon Kim, Hyunwoo Kim, Sungroh Yoon
TL;DR
The paper addresses the unclear design of loss weighting across diffusion noise levels and asks which levels best support learning useful visual concepts. It introduces perception prioritized weighting, which emphasizes those levels, and reports consistent performance improvements across datasets, model configurations, and sampling steps.
Problem
The standard weighted diffusion objective is widely used, but whether its weighting is optimal and what each noise level teaches the model remain unknown.
Method
The paper investigates visual concepts learned at different noise levels and redesigns the objective weighting to prioritize levels associated with perceptually rich content.
Results
P2 weighting consistently improves diffusion models across datasets, model configurations, and sampling steps, achieving state-of-the-art performance on CelebA-HQ and Oxford-flowers.
Takeaways & Limitations
A simple perception prioritized weighting scheme improves diffusion-model performance across the evaluated datasets, configurations, and sampling strategies.
Takeaways & Limitations
The proposed weighting and noise schedule are correlated but not equivalent, because the noise schedule affects both weights and MSE terms.
Abstract
from arXiv · showhide
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.
1. Introduction
Diffusion models train across noise levels, but the optimal weighting of these denoising losses remains unclear. The paper prioritizes levels that encourage learning rich visual concepts and reports consistent improvements across datasets, configurations, and sampling steps.
- Diffusion models progressively remove noise by learning to recover clean images from corrupted inputs at multiple noise levels.
- The standard weighted objective is widely used, but its effectiveness and optimality for sample quality remain unexplained.
- The paper investigates what visual concepts diffusion models learn at each noise level, distinguishing imperceptible details from perceptually rich content.
- The proposed weighting assigns higher loss weights to levels where models learn perceptually rich content and minimal weights to levels focused on imperceptible details.
- P2 weighting consistently outperforms the standard objective across datasets, model configurations, and sampling steps, including state-of-the-art results on CelebA-HQ and Oxford-flowers.
2. Background
Diffusion models transform data into noise and learn a reverse denoising process, with training objectives that weight losses across noise levels. The background also introduces SNR-based weighting and evaluation with FID and KID.
- Diffusion process: Diffusion models gradually corrupt data with predefined noise scales, then learn a reverse process that recovers images from noise.The forward process uses Gaussian transitions, while generation starts from x_T sampled from N(0, I) and iteratively denoises it.
- Diffusion process: Noisy data can be sampled directly as x_t = sqrt(α_t)x_0 + sqrt(1 − α_t)ϵ, with α_t determined by the cumulative noise schedule.The schedule is designed so α_T is near zero, making the terminal distribution approximately N(0, I).
- Signal-to-noise ratio: SNR(t) = α_t/(1 − α_t) characterizes noise level, decreases through diffusion, and increases through denoising.The variance of x_t can therefore be expressed in terms of SNR(t).
- Training objectives: The variational lower bound is a sum of denoising score matching losses, while the simplified objective trains a noise predictor with a uniform sum of squared errors.For each step, the denoising loss is rewritten as a distance involving the predicted and actual noise.
- Training objectives: The standard weighting can be written as λ_t = (1 − β_t)(1 − α_t)/β_t or, continuously, as λ_t = −SNR(t)/SNR′(t).The paper inherits a hybrid objective that adds cL_vlb to L_simple while modifying L_simple for performance.
- Evaluation metrics: FID and KID evaluate generated samples against the entire training set, using 50k samples for final scores and 10k samples for ablations.FID is associated with human perception, while KID is commonly used for small datasets.
3. Method
The method interprets each diffusion noise level as a signal-recovery pretext task and introduces P2 weighting to prioritize levels that learn perceptually rich content over imperceptible details. P2 suppresses clean-up-stage weights and improves FID across both linear and cosine schedules during FFHQ training.
- Learning to Recover Signals from Noisy Data: Diffusion models learn visual concepts by recovering signals from corrupted images at each noise level.The model predicts noise for a noisy image, while the corresponding recovery task explains how different levels contribute to learned content.
- Learning to Recover Signals from Noisy Data: Large SNR levels preserve perceptually rich signals, so recovery can focus on imperceptible details without holistic context.At small SNRs, noisy images lack recognizable high-level content and recovery requires prior knowledge, encouraging perceptually rich contents.
- Learning to Recover Signals from Noisy Data: The paper groups noise levels into coarse, content, and clean-up stages based on whether recovery learns global structure, perceptually rich content, or remaining noise.The proposed grouping assigns coarse features to SNR 0–10^-2, perceptually rich contents to 10^-2–100, and clean-up to 100–104.
- Perception Prioritized Weighting: P2 weighting prioritizes important noise levels by assigning minimal weights to clean-up and relatively higher weights to coarse and content stages.Its weighting is λ'_t = λ_t/(k+SNR(t))^γ, where γ controls down-weighting strength and k prevents exploding weights at extremely small SNRs.
- Perception Prioritized Weighting: P2 is a generalization of the baseline weighting scheme and reduces to the baseline when γ = 0.The scheme replaces the standard weighting in existing diffusion models without requiring a different model architecture.
- Effectiveness of P2 Weighting: P2 weighting beats the baseline in FID throughout FFHQ training for both linear and cosine schedules.The comparison uses FID-10k, with samples generated using 250 steps; the weighting and noise-schedule designs are correlated but not equivalent because schedules affect weights and MSE terms.
4. Experiment
Experiments show that P2 weighting improves diffusion models over the baseline across datasets, model configurations, and sampling procedures, while also improving qualitative coherence.
- Comparison to the Baseline: P2 weighting consistently improves FID and KID over the baseline across FFHQ, AFHQ-Dog, MetFaces, and CUB.The advantage is especially large on MetFaces, which contains approximately 1k images.
- Qualitative comparison: Baseline-trained models exhibit color-shift artifacts and underemphasize global coherence by focusing unnecessarily on imperceptible details.The authors associate this behavior with poor learning of global color schemes.
- Comparison to the Prior Literature: P2 weighting achieves state-of-the-art FIDs on Oxford Flowers and CelebA-HQ and outperforms most compared models on FFHQ except StyleGAN2.With reduced sampling, it reaches state-of-the-art using 250 steps on Oxford Flowers and 500 steps on CelebA-HQ.
- Model configuration: P2 remains superior to the baseline across varied model configurations, including settings without self-attention.The authors report especially strong effectiveness when self-attention is removed, indicating encouragement of global dependency learning.
- Sampling steps and schedules: P2 consistently outperforms the baseline across sampling steps and achieves better performance with half the baseline’s required sampling steps.Changing sampling schedules slightly improves FID and KID but does not match the improvement from P2 weighting.
5. Related Work
Related work situates diffusion and score-based models among generative approaches, emphasizing their quality, stability, adaptability, and objective-design trade-offs.
- Diffusion and score-based models: Diffusion and score-based models generate data through learned denoising processes and can be expressed as SDEs with different noise schedules.Their sample quality has been improved by recent work, although some approaches rely on heavy architectures and long training or sampling.
- Sampling schedules: Sampling-step schedules can slightly improve FID and KID but do not exceed the improvement obtained by the paper’s training weighting.This comparison motivates improving training rather than relying only on sampling schedules.
- Generative-model comparisons: Diffusion models offer superior sample quality to likelihood-based methods and are easier to scale and apply across domains than GANs.The comparison attributes the latter advantage to stable rather than adversarial training.
- Downstream adaptation: Pre-trained diffusion models can adapt to image translation and editing without task-specific training or loss functions, while supporting stochastic one-to-many generation.These capabilities are contrasted with GAN-based methods.
- Training objectives: New likelihood-oriented objectives can improve likelihood but may degrade sample quality or training stability and impede global consistency and high-level concept understanding.Prior methods address these issues through importance sampling or sophisticated parameterization.
6. Conclusion
The paper concludes that perception-prioritized weighting improves diffusion training across datasets, model configurations, and sampling steps, while leaving more sophisticated weighting designs for future work.
- Conclusion: Perception-prioritized weighting is introduced as a new diffusion-model training-objective weighting scheme.The scheme is motivated by investigating which visual concepts models learn at different noise levels.
- Conclusion: The simplest P2 choice improves diffusion models across datasets, model configurations, and sampling steps.The authors present this consistency as evidence for the scheme’s effectiveness.
- Future work: More sophisticated weighting schemes may further improve performance, but designing them is left as future work.The conclusion frames this as an open direction rather than a demonstrated result.
A.2. Derivations
The appendix derives the baseline weighting λ_t as a function of signal-to-noise ratio and shows its continuous-time approximation as a differential of log-SNR(t).
- Baseline weighting: The derivation rewrites the baseline weighting λ_t using the signal-to-noise ratio associated with diffusion step t.The appendix presents the discrete expression and its transformation into SNR-based terms.
- Continuous-time form: In the continuous-time limit, the baseline weighting is approximated by −SNR(t)/SNR′(t), the differential of log-SNR(t).The approximation is stated for T → ∞.
B.1. Limitations
Diffusion models remain computationally demanding at sampling time despite the proposed method’s promising performance. The method can improve FID while reducing the number of required sampling steps.
- B.1. Limitations: Diffusion models still require multiple sampling steps, limiting their suitability for real-time applications.The paper reports at least 25 feed-forwards with the DDIM sampler.
- B.1. Limitations: The proposed method achieves better FID with half the sampling steps required by the baseline.
- B.1. Limitations: Optimizing sampling schedules or distilling DDIM sampling into a single-step model are suggested directions for faster generation.
B.2. Broader Impacts
The method enables high-fidelity image generation, which may support creative applications but also introduces risks of deceptive use. The paper points to detection and watermarking as possible mitigations.
- B.2. Broader Impacts: The proposed method enables high-fidelity image generation with diffusion-based generative models.
- B.2. Broader Impacts: Improved generative models can support multiple creative applications, while their outputs may also be exploited for deception.
- B.2. Broader Impacts: Deepfake detection, watermarking, and invisible frequency-artifact analysis are identified as approaches that may help detect fake images.
C. Implementation Details
The implementation uses an ADM-style U-Net to predict image-resolution noise and variance, with compact default models and dataset-specific training hyperparameters.
- C. Implementation Details: The model predicts image-resolution noise and variance using an ADM-style U-Net with three input and six output channels.The architecture includes large channel dimensions, BigGAN residual blocks, multi-resolution attention, and multi-head attention.
- C. Implementation Details: The default model has 94M parameters, substantially smaller than recent models exceeding 500M parameters.The implementation uses one residual block per resolution by default, compared with two or four in recent works.
- C. Implementation Details: The training setting uses γ = 0.5 for FFHQ and CelebA-HQ because it yields slightly better FIDs than γ = 1.0.
D. Additional Results
Additional experiments extend the paper’s analyses across datasets and resolutions and provide qualitative samples. The supplementary figures examine weighting behavior, stochastic reconstruction, and generated images.
- D. Additional Results: Quantitative reproductions across datasets and resolutions support the investigation of perceptual corruption and learned concepts across settings.The results include CelebA-HQ, LSUN-Church, and CUB at 256^2 and 64^2 resolutions.
- D. Additional Results: Figure A plots normalized and unnormalized weights against diffusion steps or signal-to-noise ratio.Large t and small SNR correspond to noisy inputs, while small t and large SNR correspond to nearly clean inputs.
- D. Additional Results: The supplementary material includes a hyperparameter table for the reported experiments.
- D. Additional Results: Figure C measures perceptual distance between input and reconstructed FFHQ images as a function of signal-to-noise ratio.The measurement uses 200 randomly selected images.
- D. Additional Results: Figure D provides additional generated samples from models trained on the datasets discussed in the paper.