Diffusion详解及PyTorch代码
首先附上几个大佬的讲解
lilianweng-diffusion-models
这篇博客借鉴了上述博客和视频，同时加上个人的理解整合了一下，整个推导过程非常详细，希望能使每个人都看懂
结合之前讲过的VAE和GAN模型，Diffusion Model和他们的区别就是latent code和原图是同尺寸大小的。如下图所示，Diffusion Model分为前向过程和反向过程，前向过程将输入图片 x 0 x_{0} x0变为纯高斯噪声 x T x_{T} xT（就是一个不断加噪的过程），反向过程就是将噪声 x T x_{T} xT还原为图片 x 0 x_{0} x0的过程（就是一个不断去噪的过程）
知道Diffusion Model在做什么之后，我们分别对Diffusion的每个过程做详细的分析
Diffusion的前向过程
给定真实图片 x 0 ∼ q ( x ) x_{0} \sim q(x) x0∼q(x)，前向过程中diffusion model对其添加了 T T T次高斯噪声，分别得到图 x 1 , x 2 , x 3 , . . . , x T x_{1},x_{2},x_{3},...,x_{T} x1,x2,x3,...,xT，这个过程用下面式子表示
q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}}x_{t-1}, \beta_{t}I) q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)
整个前向加噪过程是马尔科夫过程，即 t t t时刻的状态只与 t − 1 t-1 t−1时刻有关，并且在不断加噪的过程中， x t x_{t} xt不断接近纯噪声， T → ∞ T\rightarrow \infty T→∞， x t x_{t} xt变为完全的高斯噪声，在论文中 β t \beta_{t} βt是从0.0001到0.02线性插值的， T = 1000 T=1000 T=1000，也就是说 β t \beta_{t} βt是不断增加的， 1 − β t 1-\beta_{t} 1−βt是不断减小的。因此我们再看上述分布 N ( x t ; 1 − β t x t − 1 , β t I ) \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}}x_{t-1}, \beta_{t}I) N(xt;1−βtxt−1,βtI)，随着 t t t增加， x t x_{t} xt的均值是 x t − 1 x_{t-1} xt−1的 1 − β t < 1 \sqrt{1-\beta_{t}} <1 1−βt<1倍，因此最终 x t x_{t} xt的均值不断变小，趋近于 0 0 0
为了推导方便，根据原论文我们令 α t = 1 − β t \alpha_{t} = 1-\beta_{t} αt=1−βt，令 α ‾ t = ∏ i = 1 T α i \overline{\alpha}_{t} = \prod_{i=1}^{T}\alpha_{i} αt=∏i=1Tαi，并用重参数化的方法来表示前向过程每一步的数据分布（不知道重参数化方法的直接先看最后），这里我们由 q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}}x_{t-1}, \beta_{t}I) q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)得
x t = 1 − β t x t − 1 + β t z 1 , w h e r e z 1 , z 2 , . . . , ∼ N ( 0 , I ) = α t x t − 1 + 1 − α t z 1 = α t ( α t − 1 x t − 2 + 1 − α t − 1 z 2 ) + 1 − α t z 1 = α t α t − 1 x t − 2 + α t 1 − α t − 1 z 2 + 1 − α t z 1 = α t α t − 1 x t − 2 + 1 − α t α t − 1 z ‾ 2 , z ‾ 2 ∼ N ( 0 , I ) = α ‾ t x 0 + 1 − α ‾ t z ‾ t \begin{aligned} x_{t} &= \sqrt{1-\beta_{t}}x_{t-1} + \sqrt{\beta_{t}}z_{1}, ~~~~ where~z_{1},z_{2},...,\sim \mathcal{N}(0, I) \\ &= \sqrt{\alpha_{t}}x_{t-1} + \sqrt{1-\alpha_{t}}z_{1} \\ &= \sqrt{\alpha_{t}}(\sqrt{\alpha_{t-1}}x_{t-2} + \sqrt{1-\alpha_{t-1}}z_{2}) + \sqrt{1-\alpha_{t}}z_{1} \\ &= \sqrt{\alpha_{t}\alpha_{t-1}}x_{t-2} + \sqrt{\alpha_{t}}\sqrt{1-\alpha_{t-1}}z_{2} + \sqrt{1-\alpha_{t}}z_{1} \\ &= \sqrt{\alpha_{t}\alpha_{t-1}}x_{t-2} + \sqrt{1-\alpha_{t}\alpha_{t-1}}\overline{z}_{2}, ~~~~ \overline{z}_{2}\sim \mathcal{N}(0, I) \\ &= \sqrt{\overline{\alpha}_{t}}x_{0} + \sqrt{1-\overline{\alpha}_{t}}\overline{z}_{t} \end{aligned} xt=1−βtxt−1+βtz1, where z1,z2,...,∼N(0,I)=αtxt−1+1−αtz1=αt(αt−1xt−2+1−αt−1z2)+1−αtz1=αtαt−1xt−2+αt1−αt−1z2+1−αtz1=αtαt−1xt−2+1−αtαt−1z2, z2∼N(0,I)=αtx0+1−αtzt
其中 α t 1 − α t − 1 z 2 + 1 − α t z 1 \sqrt{\alpha_{t}}\sqrt{1-\alpha_{t-1}}z_{2} + \sqrt{1-\alpha_{t}}z_{1} αt1−αt−1z2+1−αtz1到 1 − α t α t − 1 z ‾ 2 \sqrt{1-\alpha_{t}\alpha_{t-1}}\overline{z}_{2} 1−αtαt−1z2用到了高斯分布的独立可加性，即 N ( 0 , σ 1 2 I ) + N ( 0 , σ 2 2 I ) ∼ N ( 0 , ( σ 1 2 + σ 2 2 ) I ) \mathcal{N}(0, \sigma^{2}_{1}I) + \mathcal{N}(0, \sigma^{2}_{2}I) \sim \mathcal{N}(0, (\sigma^{2}_{1}+\sigma^{2}_{2})I) N(0,σ12I)+N(0,σ22I)∼N(0,(σ12+σ22)I)
α t 1 − α t − 1 z 2 ∼ N ( 0 , α t ( 1 − α t − 1 ) I ) 1 − α t z 1 ∼ N ( 0 , ( 1 − α t − 1 ) I ) α t ( 1 − α t − 1 ) z 2 + 1 − α t z 1 ∼ N ( 0 , [ α t ( 1 − α t − 1 ) + ( 1 − α t − 1 ) ] I ) ∼ N ( 0 , ( 1 − α t α t − 1 ) I ) → 1 − α t α t − 1 z ‾ 2 \begin{aligned} & \sqrt{\alpha_{t}}\sqrt{1-\alpha_{t-1}}z_{2} \sim \mathcal{N}(0, \alpha_{t}(1-\alpha_{t-1})I) \\ & \sqrt{1-\alpha_{t}}z_{1} \sim \mathcal{N}(0, (1-\alpha_{t-1})I) \\ & \sqrt{\alpha_{t}(1-\alpha_{t-1})}z_{2} + \sqrt{1-\alpha_{t}}z_{1} \sim \mathcal{N}(0, [\alpha_{t}(1-\alpha_{t-1}) + (1-\alpha_{t-1})]I) \\ & \sim \mathcal{N}(0, (1-\alpha_{t}\alpha_{t-1})I) \rightarrow \sqrt{1-\alpha_{t}\alpha_{t-1}}\overline{z}_{2} \end{aligned} αt1−αt−1z2∼N(0,αt(1−αt−1)I)1−αtz1∼N(0,(1−αt−1)I)αt(1−αt−1)z2+1−αtz1∼N(0,[αt(1−αt−1)+(1−αt−1)]I)∼N(0,(1−αtαt−1)I)→1−αtαt−1z2
有一个疑问，为什么 x t x_{t} xt的分布是 q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}}x_{t-1}, \beta_{t}I) q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)呢？因为这个公式是作者直接给出的，并没有一个推导，公式表明在加噪的过程中均值要乘上 1 − β t \sqrt{1-\beta_{t}} 1−βt，如果要保证均值最后为0的话，只需要每次乘的值小于1就可以了（虽然方差可能并不是 I I I），通过上述推导我们可以发现 x t x_{t} xt的最终等于 α ‾ t x 0 + 1 − α ‾ t z ‾ t \sqrt{\overline{\alpha}_{t}}x_{0} + \sqrt{1-\overline{\alpha}_{t}}\overline{z}_{t} αtx0+1−αtzt，即 T → ∞ , x t ∼ N ( 0 , I ) T \rightarrow \infty, x_{t} \sim \mathcal{N}(0, I) T→∞,xt∼N(0,I)，也就是说 N ( x t ; 1 − β t x t − 1 , β t I ) \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}}x_{t-1}, \beta_{t}I) N(xt;1−βtxt−1,βtI)这个分布能够保证 x t x_{t} xt最终收敛为标准高斯分布，但是具体这个式子怎么得到的，我依然不了解
Diffusion的反向过程
在前向加噪过程中，我们的加噪表达式为 q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t}; \sqrt{1-\beta_{t}}x_{t-1}, \beta_{t}I) q(xt∣xt−1)=N(xt;1−βtxt−1,βtI)，反向过程就是将上述过程进行逆转，得到 q ( x t − 1 ∣ x t ) q(x_{t-1}|x_{t}) q(xt−1∣xt)的分布，通过不断的去噪从 x t ∼ N ( 0 , I ) x_{t} \sim \mathcal{N}(0, I) xt∼N(0,I)中还原出原图 x 0 x_{0} x0，文中中证明了如果 q ( x t ∣ x t − 1 ) q(x_{t}|x_{t-1}) q(xt∣xt−1)满足高斯分布并且 β t \beta_{t} βt足够小， q ( x t − 1 ∣ x t ) q(x_{t-1}|x_{t}) q(xt−1∣xt)仍然是一个高斯分布。但是我们无法简单推断 q ( x t − 1 ∣ x t ) q(x_{t-1}|x_{t}) q(xt−1∣xt)，因此我们使用深度学习模型（参数为 θ \theta θ，结构一般为U-net+attention结构）来预测他的分布
p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_{\theta}({x_{t-1}|x_{t}}) = \mathcal{N}(x_{t-1}; \mu_{\theta}(x_{t}, t), \Sigma_{\theta}(x_{t}, t)) pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))
虽然我们无法得到反向过程的分布 q ( x t − 1 ∣ x t ) q(x_{t-1}|x_{t}) q(xt−1∣xt)，但是如果知道 x 0 x_{0} x0，是可以通过贝叶斯公式得到 q ( x t − 1 ∣ x t , x 0 ) q(x_{t-1}|x_{t}, x_{0}) q(xt−1∣xt,x0)为
q ( x t − 1 ∣ x t , x 0 ) = ( x t − 1 ; μ ~ ( x t , x 0 ) , β ~ t I ) q(x_{t-1}|x_{t}, x_{0}) = \mathcal(x_{t-1}; \tilde{\mu}(x_{t}, x_{0}), \tilde{\beta}_{t}I) q(xt−1∣xt,x0)=(xt−1;μ~(xt,x0),β~tI)
推导过程如下，首先利用贝叶斯公式将反向过程均变为前向过程 x t − 1 → x t x_{t-1} \rightarrow x_{t} xt−1→xt， x 0 → x t − 1 x_{0}\rightarrow x_{t-1} x0→xt−1以及 x 0 → x t x_{0}\rightarrow x_{t} x0→xt
q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) q ( x t ∣ x 0 ) q(x_{t-1}|x_{t}, x_{0}) = q(x_{t}|x_{t-1}, x_{0}) \frac{q(x_{t-1}|x_{0})}{q(x_{t}|x_{0})} q(xt−1∣xt,x0)=q(xt∣xt−1,x0)q(xt∣x0)q(xt−1∣x0)
根据高斯分布的概率密度函数的指数部分 N ∼ ( μ , σ 2 ) ∝ exp ( − ( x − μ ) 2 2 σ 2 ) \mathcal{N}\sim (\mu, \sigma^{2}) \propto \exp(-\frac{(x-\mu)^2}{2\sigma^{2}}) N∼(μ,σ2)∝exp(−2σ2(x−μ)2)以及前向推导公式 x t = α t x t − 1 + 1 − α t z 1 = α ‾ t x 0 + 1 − α ‾ t z ‾ t x_{t} = \sqrt{\alpha_{t}}x_{t-1} + \sqrt{1-\alpha_{t}}z_{1} = \sqrt{\overline{\alpha}_{t}}x_{0} + \sqrt{1-\overline{\alpha}_{t}}\overline{z}_{t} xt=αtxt−1+1−αtz1=αtx0+1−αtzt得
q ( x t − 1 ∣ x t , x 0 ) = q ( x t ∣ x t − 1 , x 0 ) q ( x t − 1 ∣ x 0 ) 1 q ( x t ∣ x 0 ) = [ α t x t − 1 + 1 − α t z 1 ] × [ α ‾ t − 1 x 0 + 1 − α ‾ t − 1 z ‾ t − 1 ] × [ 1 α ‾ t x 0 + 1 − α ‾ t z ‾ t ] ∝ exp ( − 1 2 ( ( x t − α t x t − 1 ) 2 β t + ( x t − 1 − α ‾ t − 1 x 0 ) 2 1 − α ‾ t − 1 − ( x t − α ‾ t x 0 ) 2 1 − α ‾ t ) ) = exp ( − 1 2 ( ( α t β t + 1 1 − α ‾ t − 1 ) x t − 1 2 ⏟ − ( 2 α t β t x t + 2 α ‾ t − 1 1 − α ‾ t − 1 x 0 ) x t − 1 ⏟ + C ( x t , x 0 ⏟ ) ) ) \begin{aligned} q(x_{t-1}|x_{t}, x_{0}) &= q(x_{t}|x_{t-1}, x_{0}) q(x_{t-1}|x_{0}) \frac{1}{q(x_{t}|x_{0})} \\ &= [\sqrt{\alpha_{t}}x_{t-1} + \sqrt{1-\alpha_{t}}z_{1}] \times [\sqrt{\overline{\alpha}_{t-1}}x_{0} + \sqrt{1-\overline{\alpha}_{t-1}}\overline{z}_{t-1}] \times [\frac{1}{\sqrt{\overline{\alpha}_{t}}x_{0} + \sqrt{1-\overline{\alpha}_{t}}\overline{z}_{t}}] \\ &\propto \exp(-\frac{1}{2}(\frac{(x_{t} - \sqrt{\alpha_{t}}x_{t-1})^2}{\beta_{t}} + \frac{(x_{t-1}-\sqrt{\overline{\alpha}_{t-1}}x_{0})^{2}}{1-\overline{\alpha}_{t-1}} - \frac{(x_{t}-\sqrt{\overline{\alpha}_{t}}x_{0})^{2}}{1-\overline{\alpha}_{t}})) \\ &= \exp(-\frac{1}{2}(( \underbrace{\frac{\alpha_{t}}{\beta_{t}}+\frac{1}{1-\overline{\alpha}_{t-1}})x^{2}_{t-1}} - \underbrace{(\frac{2\sqrt{\alpha_{t}}}{\beta_{t}}x_{t} + \frac{2\sqrt{\overline{\alpha}_{t-1}}}{1-\overline{\alpha}_{t-1}}x_{0})x_{t-1}} + \underbrace{C(x_{t}, x_{0}}))) \end{aligned} q(xt−1∣xt,x0)=q(xt∣xt−1,x0)q(xt−1∣x0)q(xt∣x0)1=[αtxt−1+1−αtz1]×[αt−1x0+1−αt−1zt−1]×[αtx0+1−αtzt1]∝exp(−21(βt(xt−αtxt−1)2+1−αt−1(xt−1−αt−1x0)2−1−αt(xt−αtx0)2))=exp(−21(( βtαt+1−αt−11)xt−12− (βt2αtxt+1−αt−12αt−1x0)xt−1+ C(xt,x0)))
根据 exp ( − ( x − μ ) 2 2 σ 2 ) = exp ( − 1 2 ( 1 σ 2 x 2 − 2 μ σ 2 x + μ 2 σ 2 ) ) \exp(-\frac{(x-\mu)^{2}}{2\sigma^{2}}) = \exp(-\frac{1}{2}(\frac{1}{\sigma^{2}}x^{2}-\frac{2\mu}{\sigma^{2}}x + \frac{\mu^{2}}{\sigma^{2}})) exp(−2σ2(x−μ)2)=exp(−21(σ21x2−σ22μx+σ2μ2))，上式的最后括号部分，我们分别进行化简能够得到 q ( x t − 1 ∣ x t , x 0 ) = ( x t − 1 ; μ ~ ( x t , x 0 ) , β ~ t I ) q(x_{t-1}|x_{t}, x_{0}) = \mathcal(x_{t-1}; \tilde{\mu}(x_{t}, x_{0}), \tilde{\beta}_{t}I) q(xt−1∣xt,x0)=(xt−1;μ~(xt,x0),β~tI)的均值和方差，如下
{ 1 σ 2 = 1 β ~ t = ( α t β t + 1 1 − α ‾ t − 1 ) 2 μ σ 2 = 2 μ ~ t ( x t , x 0 ) β ~ t = ( 2 α t β t x t + 2 α ‾ t − 1 1 − α ‾ t − 1 x 0 ) \left\{ \begin{array}{ll} \frac{1}{\sigma^{2}} = \frac{1}{\tilde{\beta}_{t}} = (\frac{\alpha_{t}}{\beta_{t}} + \frac{1}{1-\overline{\alpha}_{t-1}}) \\ ~~ \\ \frac{2\mu}{\sigma^{2}} = \frac{2\tilde{\mu}_{t}(x_{t}, x_{0})}{\tilde{\beta}_{t}} = (\frac{2\sqrt{\alpha_{t}}}{\beta_{t}}x_{t} + \frac{2\sqrt{\overline{\alpha}_{t-1}}}{1-\overline{\alpha}_{t-1}}x_{0}) \end{array} \right. ⎩ ⎨ ⎧σ21=β~t1=(βtαt+1−αt−11) σ22μ=β~t2μ~t(xt,x0)=(βt2αtxt+1−αt−12αt−1x0)
化简得
{ β ~ t = 1 − α ‾ t − 1 1 − α ‾ t ⋅ β t μ ~ t ( x t , x 0 ) = α t ( 1 − α ‾ t − 1 ) 1 − α ‾ t x t + α ‾ t − 1 β t 1 − α ‾ t x 0 \left\{ \begin{array}{ll} \tilde{\beta}_{t} = \frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_{t}} \cdot \beta_{t} \\ ~~ \\ \tilde{\mu}_{t}(x_{t}, x_{0}) = \frac{\sqrt{\alpha_{t}}(1-\overline{\alpha}_{t-1})}{1-\overline{\alpha}_{t}}x_{t} + \frac{\sqrt{\overline{\alpha}_{t-1}}\beta_{t}}{1-\overline{\alpha}_{t}}x_{0} \end{array} \right. ⎩ ⎨ ⎧β~t=1−αt1−αt−1⋅βt μ~t(xt,x0)=1−αtαt(1−αt−1)xt+1−αtαt−1βtx0
在前向过程中我们推导出 x t = α ‾ t x 0 + 1 − α ‾ t z ‾ t x_{t} = \sqrt{\overline{\alpha}_{t}}x_{0} + \sqrt{1-\overline{\alpha}_{t}}\overline{z}_{t} xt=αtx0+1−αtzt，由此得到 x 0 = 1 α ‾ t ( x t − 1 − α ‾ t z t ‾ ) x_{0} = \frac{1}{\sqrt{\overline{\alpha}_{t}}}(x_{t}-\sqrt{1-\overline{\alpha}_{t}}\overline{z_{t}}) x0=αt1(xt−1−αtzt)并替换上面均值得到
μ ~ t = 1 α t ( x t − β t 1 − α ‾ t z ‾ t ) \tilde{\mu}_{t} = \frac{1}{\sqrt{\alpha_{t}}}(x_{t} - \frac{\beta_{t}}{\sqrt{1-\overline{\alpha}_{t}}}\overline{z}_{t}) μ~t=αt1(xt−1−αtβtzt)
其中高斯分布 z ‾ t \overline{z}_{t} zt是深度学习模型所预测的噪声，可以看做 z θ ( x t , t ) z_{\theta}(x_{t}, t) zθ(xt,t)，由此得到均值为
μ θ ( x t , t ) = 1 α t ( x t − β t 1 − α ‾ t z θ ( x t , t ) ) \mu_{\theta}(x_{t}, t) = \frac{1}{\sqrt{\alpha_{t}}}(x_{t} - \frac{\beta_{t}}{\sqrt{1-\overline{\alpha}_{t}}}z_{\theta}(x_{t}, t)) μθ(xt,t)=αt1(xt−1−αtβtzθ(xt,t))
这样我们证明最初已知 x 0 x_{0} x0后的反向表达式了
补充
1. 重参数化技巧
重参数化技巧在VAE中被应用过，此技巧主要用来使采样可以进行反向传播，假设我们随机采样时从任意一个高斯分布 N ( μ , σ 2 ) \mathcal{N}(\mu, \sigma^{2}) N(μ,σ2)中采样，然后预测结果，最终结果是无法反向传播的（不可导），通常做法是使用标准高斯分布 N ( 0 , I ) \mathcal{N}(0, I) N(0,I)作为引导
具体做法是首先从标准高斯分布中采样一个变量，然后根据高斯分布的均值 μ \mu μ和方差 σ 2 \sigma^{2} σ2来对采样变量进行线性变换，如下
z = μ + σ ⊙ ϵ , ϵ ∼ N ( 0 , I ) z = \mu + \sigma \odot \epsilon, \epsilon \sim \mathcal{N}(0, I) z=μ+σ⊙ϵ,ϵ∼N(0,I)
重参数化之后得到的变量 z z z具有随机性的，满足均值为 μ \mu μ，方差为 σ 2 \sigma^{2} σ2的高斯分布，这样采样过程就可导了，随机性加到了 ϵ ∼ N ( 0 , I ) \epsilon \sim \mathcal{N}(0, I) ϵ∼N(0,I)上，而不是 N ( μ , σ 2 ) \mathcal{N}(\mu, \sigma^{2}) N(μ,σ2)
再通俗解释一下，就是如果我们直接对原来高斯分布采样的话，采样之后的所有计算是可以反向传播的，但是是传不到<在采样这个步骤之前的过程>的，因为数据具有随机性了，反向传播不能传播随机性的梯度，当我们将随机性转移到 ϵ ∼ N ( 0 , I ) \epsilon \sim \mathcal{N}(0, I) ϵ∼N(0,I)上时，因为标准高斯分布之前就没有数据了，所以不用继续传播了，传播的是对其采样之后的 μ \mu μ和 σ \sigma σ