Score Matching and Diffusion Models - II
CSE 849: Deep Learning
Vishnu Boddeti
Recap: Core Idea
- Interpolating between two distributions:
- The data distribution is denoted pdata∈P(Rd).
- The easy-to-sample distribution is denoted pref∈P(Rd).
- pref is usually the standard multivariate Gaussian.
- Going from the data to the easy-to-sample distribution: noising process.
- Going from the easy-to-sample to the data distribution: generative process.
- How to invert the forward noising process?
Recap: Forward and Reverse Diffusion SDE
- Forward diffusion: dxt=f(xt,t)dt+g(t)dwt
- Backward diffusion: yt=xT−t dyt=[−f(yt,T−t)+g(T−t)2∇xlogpT−t(yt)]dt+g(T−t)dwt
- Same diffusion coefficient.
- Opposite drift term with additional distribution correction: g(T−t)2∇xlogpT−t(yt) drives the diffusion in regions with high pT−t probability.
- x↦∇xlogpt(x) is called the (Stein) score of the distribution.
- Can we simulate this backward diffusion using Euler-Maruyama? Yn+1=Yn+h[−f(Yn,T−t)+g(T−t)2∇xlogpT−t(Yn)]+√hg(T−t)Zn
Learning the score function: Denoising score matching
- Goal: Estimate the score x↦∇xlogpt(x) using only available samples (x0,xt).
- For the models of interest, xt=αtx0+βtZt is a rescaled noisy version of x0.
- Explicit conditional distribution: pt|0(xt|x0)=N(αtx0,β2tId). pt(xt)=∫Rdp0,t(x0,xt)dx0=∫Rdpt|0(xt|x0)p0(x0)dx0 ∇xtpt(xt)=∫Rd∇xtpt|0(xt|x0)p0(x0)dx0∇xtlogpt(xt)=∇xtpt(xt)pt(xt)=∫Rd∇xtlogpt|0(xt|x0)p0(x0)pt(xt)dx0=∫Rd[∇xtlogpt|0(xt|x0)]pt|0(xt|x0)p0(x0)pt(xt)dx0=∫Rd[∇xtlogpt|0(xt|x0)]p0|t(x0|xt)dx0
- Conclusion: ∇xtlogpt(xt)=Ex0∼p0|t(x0|xt)[∇xtlogpt|0(xt|x0)]=E[∇xtlogpt|0(xt|x0)|xt]
Learning the score function: Denoising score matching continued...
∇xtlogpt(xt)=Ex0∼p0|t(x0|xt)[∇xtlogpt|0(xt|x0)]=E[∇xtlogpt|0(xt|x0)|xt]- ∇xtlogpt|0(xt|x0) is explicit (forward transition): For xt|x0∼N(αtx0,β2tId) ∇xtlogpt|0(xt|x0)=∇xt[12β2t‖xt−αtx0‖2+C]=1β2t(xt−αtx0)=1βt(Zt)
- But the distribution p0|t(x0|xt) is not explicit (backward conditional)! E[∇xtlogpt|0(xt|x0)|xt]=1β2t(xt−αtE[x0|xt])
- E[x0|xt] is the best estimate of the initial noise-free x0 given its noisy version xt.
Learning the score function: Denoising score matching continued...
∇xtlogpt(xt)=Ex0∼p0|t(x0|xt)[∇xtlogpt|0(xt|x0)]=E[∇xtlogpt|0(xt|x0)|xt]- We use the following properties of the conditional expectation.
- Y=E[X|F] if and only if Y=argminE‖X−Z‖,Z∈L2(F).
- Y∈σ(X) iff ∃f:Rd→Rd (measurable) with Y=f(X).
- Y=E[X|U] if Y=f(U) with f=argminE‖X−f(U)‖,f∈L2(U).
- Hence the function xt↦∇xtlogpt(xt) is the solution of: ∇xtlogpt=argminEp0,t‖f(xt)−∇xtlogpt|0(xt|x0)‖2,f∈L2(pt)
- We obtain a loss function to learn the function f using Monte Carlo approximation with samples (x0,xt) for the expectation.
Learning the score function: Denoising score matching continued...
∇xtlogpt(xt)=Ex0∼p0|t(x0|xt)[∇xtlogpt|0(xt|x0)]=E[∇xtlogpt|0(xt|x0)|xt]- f:Rd→Rd will be approximated with a neural network such as a (complex) U-Net (Ho et al., 2020).
- But we need to have an approximation of ∇xtlogpt for all time t (at least for the times tn in our Euler-Maruyama scheme).
- In practice we share the same network architecture for all time t: one learns a network sθ(x,t) such that sθ(x,t)≈∇xlogpt(x),x∈Rd,t∈[0,T].
- Final loss for denoising score matching: (Song et al., 2021b) θ∗=argminEt(λtE(x0,xt)‖sθ(xt,t)−∇xtlogpt|0(xt|x0)‖2) where t is chosen uniformly in [0,T] and t↦λt is a weighting term to balance the importance of each t.
Diffusion SDE: Reversed
dyt=[−f(yt,T−t)+g(T−t)2∇xlogpT−t(yt)]dt+g(T−t)dwt
Summary: Generative Modeling with SDE
Score Architecture
θ∗=argminEt(λtE(x0,xt)‖sθ(xt,t)−∇xtlogpt|0(xt|x0)‖2)- sθ:Rd×[0,T]→Rd is a (complex) U-net (Ronneberger et al., 2015), eg in (Ho et al., 2020) "All models have two convolutional residual blocks per resolution level and self-attention blocks at the 16x16 resolution between the convolutional blocks".
- Diffusion time t is specified by adding the Transformer sinusoidal position embedding into each residual block (Vaswani et al., 2017).
Exponential moving average
- Several choices for t↦λt (Kingma and Gao, 2023).
- Training using Adam algorithm (Kingma and Ba, 2015), but still unstable.
- To regularize: Exponential Moving Average (EMA) of weights. ˉθn+1=(1−m)ˉθn+mθm
- Typically m=10−4 (more than 104 iterations are averaged).
- The final averaged parameters ˉθK are used at sampling.
Exponential moving average
Sampling Strategy
- The score function of a distribution is generally used for Langevin sampling. Xn+1=Xn+γ∇xlogp(Xn)+√2γZn
- (Song et al., 2021b) propose to add one step of Langevin diffusion (same t=tn) after each step Euler-Maruyama step (tn to tn+1).
- This means that we jump from one trajectory to the other, but we correct some defaults from the Euler scheme.
- This is called a Predictor-Corrector sampler.
Sampling Strategy
Results
- (Song et al., 2021b) achieved SOTA in terms of FID for CIFAR-10 unconditional sampling.
- Very good results for 1024x1024 portrait images.
- See also "Diffusion Models Beat GANs on Image Synthesis"
Many approximations
- Many approximations in the full generative pipelines:
- The final distribution pT is not exactly a normal distribution.
- The learnt U-Net model sθ is far from being the exact score function:
- Sample-based, limitations from the architecture...
- Discrete sampling scheme (Euler-Maruyama, Predictor-Corrector,...).
- Score function may behave badly near t=0 (irregular density in case of manifold hypothesis).
- But we do have theoretical guarantees if all is well controlled.
Theorem (Convergence guarantees (De Bortoli, 2022)) Let p0 be the data distribution having a compact manifold support and let qT be the generator distribution from the reversed diffusion. Under suitable hypotheses, the 1-Wasserstein distance W1(p0,qT) can be explicitly bounded and tends to zero when all the parameters are refined (more Euler steps, better score learning, etc.).
Sampling via an ODE
- We derived the Fokker-Plank equation for qt=pT−t of reversed diffusion yt=xT−t. ∂tqt(x)=−∂tpT−t(x)=∇x(f(x,T−t)pT−t(x))−12g(T−t)2∇2xpT−t(x)=∇x(f(x,T−t)qt(x))−12g(T−t)2∇2xqt(x)=∇x(f(x,T−t)pT−t(x))+(−1+12)12g(T−t)2∇2xqt(x)=∇x([f(x,T−t)+g(T−t)2∇logqt(x)]qt(x))+12g(T−t)2∇2xqt(x)
- This is the Fokker-Planck equation associated with the diffusion SDE: dyt=[−f(yt,T−t)+g(T−t)2∇xlogpT−t(yt)]dt+g(T−t)dwt
Sampling via an ODE
- We derived the Fokker-Plank equation for qt=pT−t of reversed diffusion yt=xT−t. ∂tqt(x)=−∂tpT−t(x)=∇x(f(x,T−t)pT−t(x))−12g(T−t)2∇2xpT−t(x)=∇x(f(x,T−t)qt(x))−12g(T−t)2∇2xqt(x)=∇x(f(x,T−t)pT−t(x))+(−12+0)12g(T−t)2∇2xqt(x)=−∇x([−f(x,T−t)+12g(T−t)2∇logqt(x)]qt(x))
- This is the Fokker-Planck equation associated with the diffusion SDE: dyt=[−f(yt,T−t)+12g(T−t)2∇xlogpT−t(yt)]dt which is an Ordinary Differential Equation (ODE) (no stochastic term).
Reverse Diffusion via an ODE
Probability flow ODE: dyt=[−f(yt,T−t)+12g(T−t)2∇xlogpT−t(yt)]dt
- We get a deterministic mapping between initial noise and generated images.
- We do not simulate the (chaotic) path of the stochastic diffusion but we still have the same marginal distribution pt.
- We can use any ODE solver, with higher order than Euler scheme.
Reverse Diffusion via an ODE
Probability flow ODE: dyt=[−f(yt,T−t)+12g(T−t)2∇xlogpT−t(yt)]dt
- From (Karras et al., 2022) "Through extensive tests, we have found Heun's 2nd order method (a.k.a. improved Euler, trapezoidal rule) [...] to provide an excellent tradeoff between truncation error and NFE."
- Requires much less NFE than stochastic samplers (eg around 50 steps instead of 1000), see also Denoising Diffusion Implicit Models (DDIM) (Song et al., 2021a) for a deterministic approach.
Q & A
Score Matching and Diffusion Models - II CSE 849: Deep Learning Vishnu Boddeti