Perception Prioritized Training of Diffusion Models

In plain words

Diffusion models make images by learning to undo noise. During training, the model sees pictures covered in noise of many different strengths, and it is scored on how well it removes the noise at each strength. The standard recipe gives these strengths a fixed weighting that works well in practice, but nobody had a good explanation for why it works.

The authors look at what the model learns at each noise strength. When the noise is very light, almost the whole picture is still visible, so the model only learns fine detail that it could have guessed from its neighbours. When the noise is heavy, the picture is hidden, so the model has to learn the overall colours, shapes and recognisable content. Their change is to move training weight away from the very light end and toward the middle, where the model learns content. It is a single extra factor in the loss, and it works for the datasets, model settings and sampling setups they tried.

With that change, the same network makes images that score better on the authors' image-quality metric, and the authors say it does not need extra models or heavy extra computation.

TL;DR

• Diffusion models are trained by summing a denoising loss over noise levels. The paper replaces the standard weight lambda_t with lambda_t / (k + SNR(t))^gamma, which down-weights the low-noise (high-SNR) steps §3.2 eq. 8.
• In the released code the change is one factor on the simple MSE loss: 1 / (k + SNR)^gamma, with k = 1 and gamma set to 0.5 or 1 §3.4 gaussian_diffusion.py.
• Author-reported: FID-50k on FFHQ 7.86 → 6.92 at 1000 steps; CUB 9.60 → 6.95; MetFaces 44.34 → 36.80 at 250 steps tab. 1. The paper also reports state-of-the-art FID on CelebA-HQ and Oxford Flowers among the models it compares with §4.2 tab. 2.
• The authors say the method needs no additional models or heavy computation §5.1. It changes only the training loss. In the sampling sweep, P2 at 250 steps beats the baseline at 500 steps tab. 4.
• Official MIT-licensed PyTorch code, built on guided-diffusion, is linked. It differs from the paper's stated setup in several places (see "Paper vs code").

Problem

Standard diffusion training uses a weighted sum of denoising losses over noise levels. Ho et al.'s weighting worked better for sample quality than the plain variational bound, but the reason was unclear, and no one had tried to design a better weighting for sample quality §1.

Two things make the design hard. There are around a thousand noise levels, so an exhaustive grid search is impractical. And it is unclear what the model learns at each level, so the priority of each level is hard to set §1.

The problem is stated for image generation at 256×256 on single-model unconditional settings, where other generative families (GANs, VAEs, flows) were the main competitors §1.

Method

The authors first ask what the model learns at each noise level. They measure LPIPS distances between noisy images as a function of SNR §3.1 fig. 1. When SNR is large, noisy versions of the same image stay perceptually similar, so the model learns only imperceptible details. When SNR is small, the noisy images converge and the model has to learn content §3.1.

They verify this with reconstructions from a trained model. Reconstructions from low-SNR steps share only coarse colour structure with the input, medium-SNR reconstructions keep perceptually discriminative content, and high-SNR reconstructions are almost identical to the input §3.1 fig. 2. From this they group steps into three stages: coarse (SNR below 1e-2), content (SNR between 1e-2 and 1) and clean-up (SNR above 1) §3.1.

The P2 weighting then down-weights the clean-up stage. The weighting is:

\lambda'_t = \frac{\lambda_t}{(k+\text{SNR}(t))^{\gamma}} eq. 8

Here lambda_t is the baseline weight, which the paper defines in terms of the SNR and the schedule in eq. 7 (as -SNR(t)/SNR'(t), an approximation that holds as T goes to infinity app. A.2). gamma controls how strongly the clean-up stage is suppressed, and k prevents exploding weights at very small SNR and sets how sharp the weighting is §3.2. With gamma = 0 the weighting reduces to the baseline lambda_t §3.2.

The steps for one training batch, following the paper's description and the code:

1. Draw a step per image. Sample one diffusion step t per image. The paper does not state the rule; the code samples uniformly resample.py.
2. Noise the image. Use the closed form of the forward process, x_t = sqrt(alpha_t) x_0 + sqrt(1 - alpha_t) eps eq. 2 §2.1.
3. Predict the noise. The network eps_theta(x_t, t) predicts eps §2.2.
4. Compute the weight. The SNR of step t is alpha_t / (1 - alpha_t) eq. 4. The P2 factor is 1 / (k + SNR)^gamma eq. 8 gaussian_diffusion.py gaussian_diffusion.py.
5. Weight the simple loss. The per-image loss is the simple MSE ||eps - eps_theta||^2 eq. 6. The lambda_t of the baseline is already contained in this simple loss eq. 7 §2.2, so P2 multiplies it by the P2 factor only, not by lambda_t again.
6. Average over the batch.

# P2 training loss for one batch  [§3.2] [eq. 8]
# Input:  clean images x0 (B, C, H, W); noise predictor eps_theta;
#         alphas (T,): cumulative products of (1 - beta_t)  [§2.1]
# Output: scalar loss

def p2_loss(x0, eps_theta, alphas, k, gamma):
    # Draw one diffusion step per image, uniformly over the T steps
    # [code: guided_diffusion/resample.py#L61-L67]
    t = uniform_int(T, size=B)
    # Corrupt x0 in closed form to get the noisy image x_t  [eq. 2]
    eps = normal(shape=x0.shape)
    a = reshape(alphas[t], (B, 1, 1, 1))     # broadcast over C, H, W
    x_t = sqrt(a) * x0 + sqrt(1 - a) * eps
    # Predict the noise from the noisy image and its step  [§2.2]
    eps_hat = eps_theta(x_t, t)
    # Signal-to-noise ratio of x_t  [eq. 4]
    snr = a / (1 - a)
    # P2 factor; it multiplies the simple loss, which already carries
    # lambda_t  [eq. 8] [eq. 7]
    weight = 1 / (k + snr) ** gamma
    # Squared error averaged over pixels, one value per image  [eq. 6]
    mse = mean((eps - eps_hat) ** 2, over=(C, H, W))
    # Average the weighted loss over the batch
    return mean(weight * mse, over=B)

Key equations:

\text{SNR}(t) = \frac{\alpha_t}{1 - \alpha_t} eq. 4

L_{simple} = \sum_t \mathbb{E}_{x_0,\epsilon}\left[\left\lVert \epsilon - \epsilon_\theta(x_t, t)\right\rVert^2\right] eq. 6

The paper's training objective is a hybrid, L_hybrid = L_simple + c L_vlb with c = 1e-3, which adds Nichol et al.'s learned variance sigma_t §2.2. The authors keep this hybrid and change only L_simple §2.2 §3.4.

With k = 1, the identity 1 / (1 + SNR(t)) = 1 - alpha_t (stated in §3.4) turns the P2 factor into (1 - alpha_t)^gamma.

Several design choices are set by the authors:

• k = 1 "for easy deployment" §3.4.
• gamma is 0.5 or 1. Values over 2 cause noise artifacts because almost no weight remains for the clean-up stage §3.4.
• T = 1000 for all experiments §3.4.
• The base network is ADM's U-Net, in a lighter version, with a learned variance §3.4 app. C.

Hyperparameters

Values are as stated in the paper unless marked "code" or "not stated".

Results

Author-reported. Setup: 256×256 images, the baseline and P2 weighting trained on the same data with the same lighter ADM-style U-Net §4.1 §3.4. FID and KID are computed against the whole training set §2.3. The paper compares against the baseline, and against the published numbers of other models tab. 2.

Table 1 (baseline vs P2, FID-50k and KID-50k, KID ×10³). tab. 1

Table 2 (comparison to prior literature, FID). tab. 2 The paper's own caption says FFHQ results are reproduced from other papers, Oxford Flowers from MSG-GAN, and CelebA-HQ from LSGM. Diffusion rows are P2 models or the baseline, as labelled.

The paper says "except for GANs, we achieve superior results" tab. 2. On FFHQ it says P2 is superior to most models except StyleGAN2 §4.2.

Table 3 (model configurations, FID-10k and KID-10k, 250 steps). tab. 3 The table does not state the KID scale.

Table 4 (sampling-schedule sweep, FID-10k and KID-10k). tab. 4 The schedule column lists steps assigned to one third of the process. The dataset is not stated in the table.

Other results, described in words:

• P2 is better than the baseline on every dataset and sampling step count tested, and the gap is largest on MetFaces (1k images) §4.1 tab. 1.
• The baseline shows colour-shift artifacts early in training (FFHQ) and even at convergence (MetFaces); P2 avoids them §4.1 fig. 6.
• Across sampling step counts (FFHQ), P2 with half the steps is better than the baseline §4.3 fig. 7.
• Sweeping sampling schedules improves the baseline slightly, but does not reach P2 §4.3 tab. 4.
• The cosine schedule is much worse than linear, even with P2 §3.3 fig. 4.
• Per-noise-level analysis holds on CelebA-HQ, LSUN-Church and CUB at 256² and 64² app. D fig. B.

Paper vs code

Code: https://github.com/jychoi118/P2-weighting at 3da0947ac350072e457c211401218175bc94e137. The repository is linked from the paper's footnote §3.4.

• Where the weight goes. Paper: "replacing sum_t lambda_t L_t with sum_t lambda'_t L_t" §3.2. Code: the P2 factor multiplies only the MSE term. The learned-variance vb term is added without the factor gaussian_diffusion.py gaussian_diffusion.py. The paper does not say which term of the hybrid objective receives the weight; the code gives the answer.
• Scale of the vb term. Paper: c = 1e-3 §2.2, units not stated. Code: vb is a per-step bits-per-dimension KL gaussian_diffusion.py; with --rescale_learned_sigmas, it is multiplied by num_timesteps / 1000 gaussian_diffusion.py, which is 1 at T=1000. The code's comment says this rescaling is for "equivalence with initial implementation". The code adds one sampled step's vb per example, so its expectation over uniform t is L_vlb / T. That equals c L_vlb with c = 1e-3 only if L_vlb is in bits per dimension.
• Baseline weight is implicit. The code uses the simple MSE and multiplies only by 1 / (k + SNR)^gamma gaussian_diffusion.py. This matches the paper's lambda'_t because L_simple is already lambda_t-weighted eq. 7. The paper's lambda_t expression is approximate app. A.2.
• Code defaults differ from the paper's setup. The argparse defaults are num_res_blocks=2, num_head_channels=-1, attention_resolutions="16,8", dropout=0.0, p2_gamma=0 (baseline), lr=1e-4, batch_size=1 script_util.py image_train.py. The README command for the paper's model sets one residual block, 64 head channels, attention at 16 only, lr 2e-5, batch size 8 and gamma 1 README.md. The paper's FFHQ and CelebA-HQ runs use gamma 0.5; the README shows that only as a flag README.md.
• Attention resolution. The default is "16,8" (both resolutions) script_util.py, the README uses 16 only README.md, and the paper's table says "16, bottle" tab. A.
• Model size. The paper says 94M §C; the README says 93M README.md.
• Training augmentation. The paper describes resize and center crop §4.1. The code's load_data defaults to random_flip=True, which image_train.py does not change image_datasets.py. The paper does not mention horizontal flips.
• Timestep sampling. The paper does not state how t is drawn. The code uses a uniform sampler resample.py with default schedule_sampler="uniform" image_train.py.
• Learning-rate schedule and weight decay. The paper does not state either. The code keeps the learning rate constant unless lr_anneal_steps is set (default 0) train_util.py, and weight decay defaults to 0.0 image_train.py.
• Evaluation code. The paper says it uses "a correctly implemented library" and cites CleanFID §2.3. The repository's evaluations/evaluator.py is the ADM TensorFlow Inception FID code evaluator.py. It contains no KID function, and the paper does not say which implementation produced the reported numbers.
• Sampler. The default in the sample script is ancestral sampling; DDIM is selected by a flag image_sample.py. The README's 250-step command does not use DDIM, and the ddim25 option does README.md.
• Licence file. The repository's LICENSE is MIT, with the copyright line from guided-diffusion (OpenAI, 2021) LICENSE. The README says the repository is "heavily based on" guided-diffusion README.md.
• Pre-trained models. The README links checkpoints on external storage (OneDrive and Google Drive). They were not downloaded or checked here.

No discrepancies were found in the P2 formula itself: the code's 1 / (k + SNR)^gamma matches eq. 8 with lambda_t absorbed into L_simple.

Pitfalls

• Which term gets the weight. The paper's formula is written for the sum of lambda_t L_t. In the hybrid objective actually trained (L_simple + c L_vlb), the code weights only the simple term. Re-implementing the paper's lambda'_t on the hybrid loss would weight the learned-variance term too, which the code does not do gaussian_diffusion.py.
• Do not multiply by lambda_t again. The simple loss already contains the baseline weighting eq. 7. Multiplying by lambda_t a second time changes the objective.
• Near-zero weight at low noise. With k = 1, the P2 factor is (1 - alpha_t)^gamma. For the linear schedule at T=1000, the first step has beta_1 = 1e-4 gaussian_diffusion.py, so its weight is about 1e-4 for gamma = 1 and about 1e-2 for gamma = 0.5. The paper intends this suppression §3.2 fig. 3, but it means the loss at the lowest noise levels is nearly ignored. Use gamma above 2 with caution: the paper reports noise artifacts §3.4.
• k changes the shape, not only the scale. Paper: k prevents exploding weights for tiny SNR and "determines sharpness" §3.2. The paper gives no ablation on k; the only value used is 1 §3.4. Do not assume other k values are tuned.
• The derivation of lambda_t is approximate. eq. 7 is an approximation that holds as T goes to infinity app. A.2. The schedule-based expression should not be used as an exact weight for T=1000 without checking.
• Schedule and weight are coupled. The paper says weighting and noise schedule are "correlated but not equivalent", and that the cosine schedule is weaker overall §3.3 fig. 4.
• Learned variance and the vb scaling. The paper gives c = 1e-3 but no units for L_vlb. The code's choice (bits per dimension, scaled by T/1000) is the only evidence; see "Paper vs code".
• FID-10k and FID-50k are different numbers. Ablations (Table 3, Table 4) use FID-10k; headline results use FID-50k §2.3. Do not compare the two across tables.
• KID scale. Table 1 is in ×10³ units tab. 1. Table 3 gives no scale.
• Table A column mapping. The "Images (M)" row is given as "18, 4.4" and "4.8, 4.4, 1.6". The mapping to datasets is inferred from column order, not labelled.
• Attention wording. The appendix says "a self-attention at a single resolution (16×16)" §C while the table says "16, bottle" tab. A. Two readings are possible. The code default uses both.
• Dataset list. §4.1 lists four datasets for the baseline comparison, but Table 1 also includes Flowers tab. 1.
• Comparisons with literature are not like-for-like. The paper compares P2 at 250 or 500 steps and reports other models' numbers reproduced from other papers tab. 2. The "state-of-the-art" claim is the authors' own §4.2.
• Ablation numbers are plotted. Figure 4 (training curves) and Figure 7 (step counts) are graphs with no values in the text. Use the tables for numbers; treat the figure values as approximate.
• Costs. The weight is a closed-form function of each step's SNR, so it adds no trainable parameters §3.2 gaussian_diffusion.py. The paper does not report the training cost of P2 vs the baseline. Sampling is the expensive part: the paper says a high-resolution image takes "more than 10 min" with a modern GPU at 1000 steps §4.3, and the limitations section says at least 25 forward passes are needed with DDIM app. B.1.
• Source not used. The source tree contains tables/rebuttal-fid.tex, which the compiled paper does not \input. It is not part of the paper and is not used here.

Lineage

• Builds on: DDPM, arXiv:2006.11239 (the baseline weighting lambda_t, §2.2); Improved DDPM, arXiv:2102.09672 (hybrid objective with learned variance and c = 1e-3, §2.2); ADM, arXiv:2105.05233 (the network and codebase, §C README.md); Variational Diffusion Models, arXiv:2107.00630 (SNR view of noise schedules, §2.1).
• Sampling: DDIM, arXiv:2010.02502 §4.3.
• Related: score-based models and the SDE view, arXiv:2011.13456, which the paper says can be written with different noise schedules §5.1.
• Compared against: the GAN, VAE and flow models in tab. 2, including StyleGAN2, LSGM and VQGAN.
• Superseded by: not known from content already in this wiki.