Training diffusion models

Training

• (Ho et al. (2020) ) Denoising loss: ${\mathcal{L}}_{simple}=||{ϵ}_{\theta }(x_t,t)-{ϵ}_t|{|}_2^2$
  • This objective can be seen as a reweighted form of ${\mathcal{L}}_{VLB}$ (without the terms affecting ${\Sigma }_{\theta }$ ). The authors found that optimizing this reweighted objective resulted in much better sample quality than optimizing ${\mathcal{L}}_{VLB}$ directly, and explain this by drawing a connection to generative score matching (Song & Ermon, 2019; 2020).

”Improved denoising diffusion probabilistic models” (TLDR: they learn ${\Sigma }_{\theta }(x_t,t))$)

Learning the variance

• We want to fit $p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$

• The variational lower bound loss is derived from VDM ${\mathcal{L}}_{VLB}(\theta )=-p(x_0|x_1)+{\sum }_t{D}_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))$

• To fit ${\mu }_{\theta }$, one can simply reparametrize ${\mu }_{\theta }$ as a noise prediction network ${ϵ}_{\theta }$, and use ${\mathcal{L}}_{simple}=||{ϵ}_{\theta }(x_t,t)-{ϵ}_t|{|}_2^2$

• However to train ${\Sigma }_{\theta }$, one must use the full ${\mathcal{L}}_{VLB}$

• They do $\mathcal{L}={\mathcal{L}}_{simple}+\lambda {\mathcal{L}}_{VLB}$

• ${\mu }_{\theta }$ is only trained using $L_{simple}$, thus they use a stop-grad when computing ${\mathcal{L}}_{VLB}$ (https://github.com/openai/improved-diffusion/blob/1bc7bbbdc414d83d4abf2ad8cc1446dc36c4e4d5/improved_diffusion/gaussian_diffusion.py#L679)

  • simply use mean_pred.detach() as the mean when computing the VLB.

• For sample quality, the first few steps of the diffusion don’t really matter i.e. very small details. HOWEVER, for maximizing log-likelihood, the first few steps of the diffusion process matter the most as they contribute the most to the variational lower bound (Fig.2 of the paper)

  • This is because the likelihood of a training sample $x_0$ with very little noise must still be well calibrated and high for this sample. This is not really taken into account when only doing noise matching loss.

• In “Improved denoising diffusion probabilistic models”, they characterized the variance as: ${\Sigma }_{\theta }(x_t,t)=exp(vlog{\beta }_t+(1-v)log{\tilde{\beta }}_t)$

• where ${\beta }_t$ is the variance schedule and ${\tilde{\beta }}_t$ is the variance of the posterior $q(x_{t-1}|x_t,x_0)$

• and $v$ is a vector containing one component per dimension

Better noise schedule

• cosine schedule
• they use a small offset $s$ to prevent ${\beta }_t$ to be too small near t=0
• They selected $s=0.008$ such that $\sqrt{{\beta }_0}$ was slightly smaller than the pixel bin size 1/127.5.
• Can be smaller for other modalities.

Reducing gradient noise

• ${\mathcal{L}}_{VLB}$ introduces a lot of gradient noise
• gradient noise = l2 norm of concatenated gradient
• Noting that different terms of Lvlb have greatly different magnitudes (Figure 2), we hypothesized that sampling t uniformly causes unnecessary noise in the Lvlb objective
• Simple importance sampling technique reduces this noise
  • ${\mathcal{L}}_{VLB}=\mathbb{E}[\frac{L_t}{p_t}]$ where $p_t\sim \sqrt{\mathbb{E}[L_t^2]}$ and $\sum p_t=1$
• We found that the importance sampling technique was not helpful when optimizing the less-noisy ${\mathcal{L}}_{hybrid}$ objective directly.

Wurtschen (learned noise gating trick)

• We have $x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t,\text{where}{ϵ}_t\sim \mathcal{N}(0,1)$
• To predict $\overline{ϵ}$, they do $\overline{ϵ}=\frac{x_t-A}{|1-B|+1e^{-5}}$
• with $A,B=f_{\theta }(x_t,t)$, $A$ and $B$ have the same dimension as the noise $ϵ$. The division is element-wise.
• It makes the training more stable. They hypothesize this occurs because the model parameters are initialized to predict 0 at the beginning, enlarging the difference to timesteps with a lot of noise. By reformulating to the $A\&B$ objective, the model initially returns the input, making the loss small for very noised inputs.
• noise prediction: https://github.com/dome272/Wuerstchen/blob/main/modules.py#L307
• diffusion implementation: https://github.com/pabloppp/pytorch-tools/blob/master/torchtools/utils/diffusion.py
• Additionally, they do p2 loss weighted noise matching
  • $\mathcal{L}=p_2(t)\cdot ||ϵ-\overline{ϵ}|{|}^2$ where $p_2(t)=\frac{1+{\bar{\alpha }}_t}{1-{\bar{\alpha }}_t}$
  • making higher noise levels contribute more to the loss

Analyzing and Improving the Training Dynamics of Diffusion Models

Noisy training signal

• The training dynamics of diffusion models remain challenging due to the highly stochastic loss function.
  • The final image quality is dictated by faint image details predicted throughout the sampling chain
  • small mistakes at intermediate steps can have snowball effects in subsequent iterations
• The network must accurately estimate the average clean image across a vast range of noise levels, Gaussian noise realizations, and conditioning inputs.
• Learning to do so is difficult given the chaotic training signal that is randomized over all of these aspects