前言
扩散模型再怎么牛,也只能在同一个数据集下生成类似的图像,不可能随心所欲地命令其生成想要的图像。但是实际上文生图早已实现,而且不是什么难事,这是怎么做到的呢?这便是条件扩散模型的功劳,是它实现了文本等语义信息的编码嵌入,控制扩散模型生成指定内容的图像,以条件扩散模型为基础,衍生出了一系列的控制模型如ControlNet等,能够更加精细地控制生成图像的内容,这才给扩散模型强大的生成能力插上了任意飞翔的翅膀。掌握在扩散模型内加入条件的设计思想,才是真正掌握了扩散模型的精髓。
条件扩散模型
扩散模型的雏形最早于2018年就已经被提出,集大成版本DDPM于2020年提出、实用版本DDIM于2020年底提出,但是真正火起来的时间在2021年前后,中间的时间差,都在等它!
术语解释
条件扩散模型在扩散模型的采样步骤上进行修改,采样是指,从已知的概率分布模型(或确定性模型中)抽取一个图像样本的过程。例如,对于DDPM而言,逆过程中每一步都是从一个已知的分布N \mathcal{N} N ;而对于DDIM而言,它的逆过程不是随机的,是从已知的上一步样本直接推断下一步样本。
扩散模型回顾
可参考DDPM的博客 和DDIM的博客 ,尤其注意采样部分的内容:
p θ ( x t ∣ x t + 1 ) = N [ x t ; μ ϕ ( x t + 1 ) , Σ ϕ ] \begin{equation}p_{\bm{\theta}}(\mathbf{x}_t|\mathbf{x}_{t+1})=\mathcal{N}[\mathbf{x}_t;\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}_{t+1}),\boldsymbol{\Sigma}_{\boldsymbol{\phi}}]\end{equation}
p θ ( x t ∣ x t + 1 ) = N [ x t ; μ ϕ ( x t + 1 ) , Σ ϕ ]
x t = α t α [ x t + 1 − 1 − α t + 1 ϵ θ ( x t + 1 , t ) ] + 1 − α t − σ t + 1 2 ϵ θ ( x t + 1 , t ) + σ t + 1 z \begin{equation}\mathbf{x}_t=\sqrt{\frac{\alpha_t}{\alpha}}\left[\mathbf{x}_{t+1}-\sqrt{1-\alpha_{t+1}}\bm{\epsilon_\theta}(\mathbf{x}_{t+1},t)\right]+\sqrt{1-\alpha_t-\sigma^2_{t+1}}\bm{\epsilon_\theta}(\mathbf{x}_{t+1},t)+\sigma_{t+1}\mathbf{z}\end{equation}
x t = α α t [ x t + 1 − 1 − α t + 1 ϵ θ ( x t + 1 , t ) ] + 1 − α t − σ t + 1 2 ϵ θ ( x t + 1 , t ) + σ t + 1 z
其中z ∼ N ( 0 , I ) \mathbf{z}\sim\mathcal{N}(\mathbf{0,I}) z ∼ N ( 0 , I ) 。
条件引导
在论文[1] 指出,当需要使用预训练的图像分类器 作为引导模型时,基于DDPM和DDIM的采样过程可以修改为以下两个非常简洁的条件引导公式:
DDPM模型的条件引导:p θ ( x t ∣ x t + 1 ) = N [ x t ; μ ϕ ( x t + 1 ) + s Σ ∇ x t + 1 log p ϕ ( y ∣ x t + 1 ) , Σ ϕ ] p_{\bm{\theta}}(\mathbf{x}_t|\mathbf{x}_{t+1})=\mathcal{N}[\mathbf{x}_t;\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}_{t+1})+s\Sigma\nabla_{x_{t+1}}\log p_\phi(\mathbf{y|x_{t+1}}),\boldsymbol{\Sigma}_{\boldsymbol{\phi}}] p θ ( x t ∣ x t + 1 ) = N [ x t ; μ ϕ ( x t + 1 ) + s Σ ∇ x t + 1 log p ϕ ( y∣ x t + 1 ) , Σ ϕ ] ;
DDIM的条件引导:ϵ ^ θ ( x t ) = ϵ θ ( x t + 1 ) − 1 − α ˉ t + 1 ∇ x t + 1 log p ϕ ( y ∣ x t + 1 ) \hat{\epsilon}_{\bm{\theta}}(\mathbf{x}_t)=\bm{\epsilon_\theta}(\mathbf{x}_{t+1})-\sqrt{1-\bar{\alpha}_{t+1}}\nabla_{\mathbf{x}_{t+1}}\log p_\phi(\mathbf{y|x}_{t+1}) ϵ ^ θ ( x t ) = ϵ θ ( x t + 1 ) − 1 − α ˉ t + 1 ∇ x t + 1 log p ϕ ( y∣x t + 1 ) 。
式中,s s s 是一个超参数项,用于控制条件项的影响程度。相比DDPM需要使用超参数控制条件引导,DDIM的条件引导显得更加自然。
两项都是针对逆过程而言的,其中第一项是DDPM在对逆过程后一个时刻的图像x t \mathbf{x}_t x t 生成的概率分布作修改,然后以条件引导的概率模型采样获取前一个时刻的图像x t − 1 \mathbf{x}_{t-1} x t − 1 ;第二项是DDIM在对逆过程后一个时刻的图像样本x t \mathbf{x}_t x t 向前一个时刻的图像样本x t − 1 \mathbf{x}_{t-1} x t − 1 递推过程的修改,加入了条件引导的梯度项。综上可知:
这个公式非常简洁优美,也非常方便使用,当然也必须要注意,这里的引导项是分类器,即直观的理解是让扩散模型生成的图像符合分类器的判断,因此生成的模型将逼近分类器 预训练的图像分布 。
数学原理
一定要使用分类器才能得到条件引导的结果吗?数字上讲其实其他一切形式的引导都可以,哪怕是一个图像的梯度也能用于引导扩散模型生成。在已知扩散模型的生成过程 p θ ( x t ∣ x t + 1 ) p_{\bm{\theta}}(\mathbf{x}_t|\mathbf{x}_{t+1}) p θ ( x t ∣ x t + 1 ) 、扩散过程 q ( x t + 1 ∣ x t ) q(\mathbf{x}_{t+1}|\mathbf{x}_t) q ( x t + 1 ∣ x t ) 和条件分布 q ϕ ( y ∣ x t ) q_{\bm{\phi}}(\mathbf{y|x}_t) q ϕ ( y∣x t ) 的前提下,可以定义条件扩散过程为q ^ ≈ q \hat{q}\approx q q ^ ≈ q ,在假设条件扩散过程与扩散过程的原始图像、单步扩散条件概率相等的情况下时,可以证明,条件扩散过程的扩散概率分布与扩散过程是完全相等的,即有:q ^ ( x t + 1 ∣ x t ) = q ( x t + 1 ∣ x t ) \hat{q}(\mathbf{x}_{t+1}|\mathbf{x}_t) = q(\mathbf{x}_{t+1}|\mathbf{x}_t) q ^ ( x t + 1 ∣ x t ) = q ( x t + 1 ∣ x t ) ,由此得到条件扩散模型的单步生成条件概率 :
q ^ ( x t ∣ x t + 1 , y ) = q ( x t ∣ x t + 1 ) q ^ ( y ∣ x t ) q ^ ( y ∣ x t + 1 ) \hat{q}(\mathbf{x}_t|\mathbf{x}_{t+1},y)=\frac{q(\mathbf{x}_t|\mathbf{x}_{t+1})\hat{q}(\mathbf{y|x}_t)}{\hat{q}(\mathbf{y|x}_{t+1})}
q ^ ( x t ∣ x t + 1 , y ) = q ^ ( y∣x t + 1 ) q ( x t ∣ x t + 1 ) q ^ ( y∣x t )
分别使用拟合 的生成过程p θ p_{\bm{\theta}} p θ 和分类器模型p ϕ p_{\bm{\phi}} p ϕ 替换上式的分子项,而分母项在单步条件生成过程中与x t \mathbf{x}_t x t 无关,不影响单步生成,则需要优化的生成过程将修改为:
p θ , ϕ ( x t ∣ x t + 1 , y ) = Z p θ ( x t ∣ x t + 1 ) p ϕ ( y ∣ x t ) p_{\bm{\theta,\phi}}(\mathbf{x}_t|\mathbf{x}_{t+1},\mathbf{y})=Zp_{\bm{\theta}}(\mathbf{x}_t|\mathbf{x}_{t+1})p_{\bm{\phi}}(\mathbf{y|x}_t)
p θ , ϕ ( x t ∣ x t + 1 , y ) = Z p θ ( x t ∣ x t + 1 ) p ϕ ( y∣x t )
其中Z Z Z 是常数,不影响最终的结果。第一项就是扩散模型的生成过程,而第二项是条件引导项,这里是分类器。
DDPM生成推导
DDPM的生成过程中,前一个时刻的样本从后一个时刻的样本分布中采样,单步生成过程服从高斯分布,可见于公式(1) 。
对于条件引导项,使用Taylor公式在x t = μ ( x t + 1 ) \mathbf{x}_t=\bm{\mu}(\mathbf{x}_{t+1}) x t = μ ( x t + 1 ) 处展开对数概率密度,取一阶近似,有:
log p ϕ ( y ∣ x t ) ≈ log p ϕ ( y ∣ x t ) x t = μ ( x t + 1 ) + ∇ x t T log p ϕ ( y ∣ x t ) x t = μ ( x t + 1 ) ⋅ [ x t − μ ( x t + 1 ) ] = [ x t − μ ( x t + 1 ) ] T ∇ x t log p ϕ ( y ∣ x t ) x t = μ ( x t + 1 ) + C 1 \begin{aligned}
\log p_{\bm{\phi}}(\mathbf{y|x}_t)&\approx\log p_{\bm{\phi}}(\mathbf{y|x}_{t})_{\mathbf{x}_t=\bm{\mu}(\mathbf{x}_{t+1})}+\nabla_{\mathbf{x}_{t}}^\mathrm{T}\log p_{\bm{\phi}}(\mathbf{y|x}_{t})_{\mathbf{x}_t=\bm{\mu}(\mathbf{x}_{t+1})}\cdot[\mathbf{x}_t-\bm{\mu}(\mathbf{x}_{t+1})]\\
&=[\mathbf{x}_t-\bm{\mu}(\mathbf{x}_{t+1})]^\mathrm{T}\nabla_{\mathbf{x}_{t}}\log p_{\bm{\phi}}(\mathbf{y|x}_{t})_{\mathbf{x}_t=\bm{\mu}(\mathbf{x}_{t+1})}+C_1
\end{aligned}
log p ϕ ( y∣x t ) ≈ log p ϕ ( y∣x t ) x t = μ ( x t + 1 ) + ∇ x t T log p ϕ ( y∣x t ) x t = μ ( x t + 1 ) ⋅ [ x t − μ ( x t + 1 )] = [ x t − μ ( x t + 1 ) ] T ∇ x t log p ϕ ( y∣x t ) x t = μ ( x t + 1 ) + C 1
如取g = ∇ x t T log p ϕ ( y ∣ x t ) x t = μ ( x t + 1 ) \mathbf{g}=\nabla_{\mathbf{x}_{t}}^\mathrm{T}\log p_{\bm{\phi}}(\mathbf{y|x}_{t})_{\mathbf{x}_t=\bm{\mu}(\mathbf{x}_{t+1})} g = ∇ x t T log p ϕ ( y∣x t ) x t = μ ( x t + 1 ) ,代入至待优化的生成过程,算出条件引导的生成过程分布为:
log [ p θ ( x t ∣ x t + 1 ) p ϕ ( y ∣ x t ) ] = − 1 2 ( x t − μ ( x t + 1 ) ) T Σ − 1 ( x t − μ ( x t + 1 ) ) + g T ( x t − μ ( x t + 1 ) ) + C 2 = − 1 2 ( x t − μ ( x t + 1 ) − Σ g ) T Σ − 1 ( x t − μ ( x t + 1 ) − Σ g ) + C 3 = log p ( ζ ) + C 4 , ζ ∼ N [ x t ; μ ( x t + 1 ) + Σ g , Σ ] \begin{aligned}
\log[p_{\bm{\theta}}(\mathbf{x}_t|\mathbf{x}_{t+1})p_{\bm{\phi}}(\mathbf{y|x}_t)]&=-\frac{1}{2}(x_t-\mu(x_{t+1}))^\mathrm{T}\Sigma^{-1}(x_t-\mu(x_{t+1}))+\mathbf{g}^\mathrm{T}(x_t-\mu(x_{t+1}))+C_2\\
&=-\frac{1}{2}(x_t-\mu(x_{t+1})-\Sigma\mathbf{g})^\mathrm{T}\Sigma^{-1}(x_t-\mu(x_{t+1})-\Sigma\mathbf{g})+C_3\\
&=\log p(\bm{\zeta})+C_4,\quad\bm{\zeta}\sim\mathcal{N}[\mathbf{x}_t;\mu(\mathbf{x}_{t+1})+\Sigma\mathbf{g},\Sigma]
\end{aligned}
log [ p θ ( x t ∣ x t + 1 ) p ϕ ( y∣x t )] = − 2 1 ( x t − μ ( x t + 1 ) ) T Σ − 1 ( x t − μ ( x t + 1 )) + g T ( x t − μ ( x t + 1 )) + C 2 = − 2 1 ( x t − μ ( x t + 1 ) − Σ g ) T Σ − 1 ( x t − μ ( x t + 1 ) − Σ g ) + C 3 = log p ( ζ ) + C 4 , ζ ∼ N [ x t ; μ ( x t + 1 ) + Σ g , Σ ]
由此原式得证。
DDIM生成推导
DDIM的生成过程与上述不同在于,其生成过程的前一个时刻样本直接由后一个时刻的样本取得,不需要经过概率分布采样,即单步生成条件概率为1,可见于公式(2) 。作者使用了一个小技巧,将基于概率模型的条件引导项转化为确定模型的分数函数 score funtion 梯度项[2] ,即:
∇ x t + 1 log p θ ( x t ) = − 1 1 − α ˉ t + 1 ϵ θ ( x t + 1 , t + 1 ) ∇ x t + 1 log p θ ( y , x t ) = − 1 1 − α ˉ t + 1 ϵ ^ θ ( x t + 1 , t + 1 ) \begin{aligned}
\nabla_{\mathbf{x}_{t+1}}\log p_{\bm{\theta}}(\mathbf{x}_t)&=-\frac{1}{\sqrt{1-\bar{\alpha}_{t+1}}}\bm{\epsilon_\theta}(\mathbf{x}_{t+1},t+1)\\
\nabla_{\mathbf{x}_{t+1}}\log p_{\bm{\theta}}(\mathbf{y,x}_t)&=-\frac{1}{\sqrt{1-\bar{\alpha}_{t+1}}}\bm{\hat{\epsilon}_\theta}(\mathbf{x}_{t+1},t+1)
\end{aligned}
∇ x t + 1 log p θ ( x t ) ∇ x t + 1 log p θ ( y , x t ) = − 1 − α ˉ t + 1 1 ϵ θ ( x t + 1 , t + 1 ) = − 1 − α ˉ t + 1 1 ϵ ^ θ ( x t + 1 , t + 1 )
其中第二行是条件引导的联合生成过程,将其使用条件概率公式展开,可得:
log [ p θ ( x t + 1 ) p ϕ ( y ∣ x t + 1 ) ] = ∇ x t + 1 log p θ ( x t + 1 ) + ∇ x t + 1 log p θ ( y ∣ x t + 1 ) − 1 1 − α ˉ t + 1 ϵ ^ θ ( x t + 1 , t + 1 ) = − 1 1 − α ˉ t + 1 ϵ θ ( x t + 1 , t + 1 ) + ∇ x t + 1 log p θ ( y ∣ x t + 1 ) ϵ ^ θ ( x t + 1 , t + 1 ) = ϵ θ ( x t + 1 , t + 1 ) − 1 − α ˉ t + 1 ∇ x t + 1 log p θ ( y ∣ x t + 1 ) \begin{aligned}
\log[p_{\bm{\theta}}(\mathbf{x}_{t+1})p_{\bm{\phi}}(\mathbf{y|x}_{t+1})]&=\nabla_{\mathbf{x}_{t+1}}\log p_{\bm{\theta}}(\mathbf{x}_{t+1})+\nabla_{\mathbf{x}_{t+1}}\log p_{\bm{\theta}}(\mathbf{y|x}_{t+1})\\
-\frac{1}{\sqrt{1-\bar{\alpha}_{t+1}}}\bm{\hat{\epsilon}_\theta}(\mathbf{x}_{t+1},t+1)&=-\frac{1}{\sqrt{1-\bar{\alpha}_{t+1}}}\bm{\epsilon_\theta}(\mathbf{x}_{t+1},t+1)+\nabla_{\mathbf{x}_{t+1}}\log p_{\bm{\theta}}(\mathbf{y|x}_{t+1})\\
\bm{\hat{\epsilon}_\theta}(\mathbf{x}_{t+1},t+1)&=\bm{\epsilon_\theta}(\mathbf{x}_{t+1},t+1)-\sqrt{1-\bar{\alpha}_{t+1}}\nabla_{\mathbf{x}_{t+1}}\log p_{\bm{\theta}}(\mathbf{y|x}_{t+1})
\end{aligned}
log [ p θ ( x t + 1 ) p ϕ ( y∣x t + 1 )] − 1 − α ˉ t + 1 1 ϵ ^ θ ( x t + 1 , t + 1 ) ϵ ^ θ ( x t + 1 , t + 1 ) = ∇ x t + 1 log p θ ( x t + 1 ) + ∇ x t + 1 log p θ ( y∣x t + 1 ) = − 1 − α ˉ t + 1 1 ϵ θ ( x t + 1 , t + 1 ) + ∇ x t + 1 log p θ ( y∣x t + 1 ) = ϵ θ ( x t + 1 , t + 1 ) − 1 − α ˉ t + 1 ∇ x t + 1 log p θ ( y∣x t + 1 )
直观上的理解是,由于DDIM实际上单步生成不需要采样,因此条件引导将直接作用于去噪网络ϵ θ \bm{\epsilon_\theta} ϵ θ ,因此条件引导将直接作用于去噪UNet网络,本方法更适用于单纯含有UNet网络的条件引导。
总结
关于DDPM和DDIM两种生成式模型的条件引导,可以使用下图来方便理解:
图中的综合点处按去噪公式综合,可见不同时间阶段的去噪图像通过某个迭代的高斯分布关联,而条件引导将直接作用于关联的高斯分布上。
而对于DDIM,其不同时间阶段的去噪图像直接关联,因此条件引导将直接作用于去噪UNet上,这里体现在与去噪结果的估计噪声ϵ t \bm{\epsilon}_t ϵ t 综合,如下图所示:
条件引导的扩散模型,是将语义与生成式模型的一次优美的尝试(虽然我不知道以前的GAN网络有无类似的),它的数学原理简洁而有力,因此其融合效果将比直接依靠花里胡哨的网络设计更加有效,这才是正儿八经的深度学习!
参考文献
[1] Song J. , Meng C. , Ermon S. .Denoising Diffusion Implicit Models[M/OL].arXiv,2022
[2] Yang S. , Jascha S. D. , Diederik P. K. , Abhishek K. , Stefano E. and Ben P. . Score-based generative modeling through stochastic differential equations[M/OL]. arXiv:2011.13456 , 2020.