文章

扩散模型

扩散模型

DDPM

加噪声的角度

DDPM的数学框架在《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》就已经完成了,但是《DDPM》在高分辨率图像上调试了出来。

DDPM的前向过程:

\[x_t=\alpha_tx_{t-1}+\beta_t\epsilon_t, \epsilon_t\sim\mathcal{N}(0,I) \tag3\]

这里$\alpha_t,\beta_t>0$且$\alpha_t^2+\beta_t^2=1$。于是有:

\[x_t=(\alpha_t\cdots\alpha_1)x_0+\sqrt{1-(\alpha_t\cdots\alpha_1)^2}\bar{\epsilon}_t,\bar\epsilon_t\sim\mathcal{N}(0,I)\tag4\]

记作

\[x_t=\bar{\alpha}_tx_0+\bar\beta_t\bar\epsilon_t\tag6\]

如果期待训练一个模型$\mu(x_t):x_t\mapsto x_{t-1}$,可以有

\[\mu(x_t)=\dfrac{1}{\alpha_t}(x_t-\beta_t\epsilon_\theta(x_t,t))\tag8\]

损失函数

\[||x_{t-1}-\mu(x_t)||^2=\dfrac{\beta_t^2}{\alpha_t^2}||\epsilon_t-\epsilon_\theta(x_t,t)\tag9||\]

里面那个东西可以写成

\[||\epsilon_t-\epsilon_\theta(\bar\alpha_tx_0+\alpha_t\bar\beta_{t-1}\bar\epsilon_{t-1}+\beta_t\epsilon_t,t)||\tag{11}\]

这里不用$\bar\epsilon_t$是因为它和前面的$\epsilon_t$不独立,只能采样$\bar\epsilon_{t-1}$,但是其实我们记$\bar\beta_t\epsilon=\alpha_t\bar\beta_{t-1}\bar\epsilon_{t-1}+\beta_t\epsilon_t$,可以知道它和$\beta_t\bar\epsilon_{t-1}-\alpha_t\bar\beta_{t-1}\epsilon_t=\bar\beta_t\omega$是独立的,于是我们求(11)式的期望不如直接求$\epsilon$和$\omega$这两个独立标准正态随机变量的期望,即代入(11):

\[\mathbb{E}_{\omega,\epsilon}[||\dfrac{\beta_t\epsilon-\alpha_t\bar\beta_{t-1}\omega}{\bar\beta_t}-\epsilon_\theta(\bar\alpha_tx_0+\bar\beta_t\epsilon,t)||^2]\tag13\]

含$\omega$的项的期望其实是常数,可以直接算出来,所以最后只用对$\epsilon$求期望就行了,于是最终的损失函数是:

\[||\epsilon-\dfrac{\bar\beta_t}{\beta_t}\epsilon_\theta(\bar\alpha_tx_0+\bar\beta_t\epsilon,t)||^2\tag{15}\]

VAE的角度

原论文其实是用VAE的角度去理解,把编码$x\rightarrow z$和生成$z\rightarrow x$变成需要$T$步才能完成,每步的编码分布$p(x_tx_{t-1})$和$q(x_{t-1}x_t)$只负责一个微小变化。相应地:
\[\begin{array}{} p(x_0,x_1,\cdots,x_T)=p(x_T|x_{T-1})\cdots p(x_1|x_0)p(x_0)\\q(x_0,x_1,\cdots,x_T)=q(x_0|x_1)\cdots p(x_{T-1}|T)q(x_T) \end{array}\tag3\]

VAE负责最小化KL散度,即:

\[KL(p||q)=\int p()\log\dfrac{p()}{q()}\]
接下来定下来$p()$和$q()$,当然是正态分布:$p(x_tx_{t-1})=\mathcal{N}(x_t;\alpha_tx_t,\beta_tI)$,$q(x_{t-1}x_t)=\mathcal{N}(x_{t-1};\mu(x_t),\sigma_t^2I)$,于是忽视掉可以直接计算积分的$p$的项和$q(x_T)$项,有
\[-\sum\int p()\log q(x_{t-1}|x_t)\]

对$t$来说,$x_{t-1}$是$x_t$的分布,所以采样的时候大于$t$的积分也能直接积成$1$,现在每一项可以写成:

\[-\int p(x_t|x_{t-1})\cdots p(x_1|x_0)p(x_0)\log q(x_{t-1}|x_t)\]
然后算$q(x_{t-1}x_t)$的时候也可以用不着$x_1,\cdot x_{t-2}$,只要直接有$p(x_{t-1}x_0)=\mathcal{N}(x_{t-1};\bar\alpha_{t-1}x_0,\bar\beta_{t-1}^2I)$,最后每一项就变成了:
\[-\int p(x_t|x_{t-1})p(x_{t-1}|x_0)p(x_0)\log q(x_{t-1}|x_t)\]

这个写开就和前文是一样的了。

贝叶斯公式的角度

最后一种理解方式是贝叶斯公式。

我们期望求$p(x_{t-1}x_t)$,但求不了,于是先假设知道$x_0$:
\[p(x_{t-1}|x_t,x_0)=\dfrac{p(x_t|x_{t-1})p(x_{t-1}|x_0)}{p(x_t|x_0)}\]

这个其实是可以计算的:

\[p(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\dfrac{\alpha_t\bar\beta_{t-1}^2}{\bar\beta_t^2}x_t+\dfrac{\alpha_{t-1}\bar\beta_t^2}{\bar\beta_t^2}x_0,\dfrac{\bar\beta_{t-1}^2\beta_t^2}{\bar\beta_t^2}I)\]
但是$x_0$是模型$\bar\mu(x_t)$的预测目标,我们期待$p(x_{t-1}x_t)\approx p(x_{t-1}x_t,x_0=\bar\mu(x_t))$。

我们假设$x_t=\bar\alpha_tx_0+\bar\beta_t\epsilon$,就有

\[\bar\mu(x_t)=\dfrac{1}{\bar\alpha_t}(x_t-\bar\beta_t\bar\epsilon_\theta(x_t,t))\]
就得到损失函数$ x_0-\bar\mu(x_t) ^2$了,等价于:
\[||\epsilon-\bar\epsilon_\theta(\bar\alpha_tx_0+\bar\beta_t\epsilon,t)||^2\]

简直一步到位。

DDIM

DDIM:《Denoising Diffusion Implicit Models》

回顾贝叶斯的推导过程,我们有$p(x_tx_{t-1})$,推导$p(x_tx_0)$,推导$p(x_{t-1}x_t,x_0)$,用神经网络近似$p(x_{t-1}x_t)$。
然而损失函数只依赖$p(x_tx_0)$,采样过程只依赖$p(x_{t-1}x_t)$,于是我们干脆想把$p(x_tx_{t-1})$去掉好了。
那么,期待$p(x_{t-1}x_t,x_0)$是一个正态分布,即
\[p(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1},\kappa_tx_t+\lambda_tx_0,\sigma_t^2I)\]

作为分布,它得满足边际分布条件。

\[\int p(x_{t-1}|x_t,x_0)p(x_t|x_0)dx_t=p(x_{t-1}|x_0)\]

假设

\[x_t=\bar{\alpha}_tx_0+\bar\beta_t\bar\epsilon_t\]

我们可以得到:

略,总之就是三个未知数两个方程,$\sigma_t$现在可以是个可变参数了。

取$\sigma_t=…$时,就变成了DDPM,取$\sigma=0$,这时就变成确定性变换了。

另外,DDPM的训练结果实质上包含了任意子序列的参数的训练结果。因为训练过程实际上就是由各个$\bar\alpha_t$确定的。因而我们可以在训练时使用一个$T$步的模型,但在生成过程中采用它的子序列。这样可以加速生成。

SGM through SDEs

把时间变成连续值,有

\[dx=f_t(x)dt+g_tdw\tag1\]

大致是说

\[x_{t+\Delta t}-x_t=f_t(x_t)\Delta t+g_t\sqrt{\Delta t}\epsilon\tag4\]

《Reverse-Time Diffusion Equation Models》给出了反向的SDE:

\[dx=[f_t-g_t^2\nabla_x\log p_t(x)]dt+g_tdw\]
取一些简单的$f$比如线性,$p(x_tx_0)$可以求出解析解。这时
\[p(x_t)=\int p(x_t|x_0)p(x_0)=\mathbb{E}_{x_0}[p(x_t|x_0)]\]

不过还是用神经网络$s_\theta(x,t)$来估计$\nabla_x\log p_t(x)$

\[\mathbb{E}_{x_0,x_t\sim p(x_t|x_0)p(x_0)}[||s_\theta(x_t,t)-\nabla_{x_t}\log p(x_t|x_0)||^2]\]

对(1)用狄拉克函数,展完泰勒求期望可以导出福克-普朗克方程:

\[\dfrac{\partial}{\partial t}p_t(x)=-\nabla_x\cdot[f_t(x)p_t(x)]+\dfrac12g_t^2\nabla_x\cdot\nabla_xp_t(x)\]

FP方程描述了粒子速度的概率密度函数在阻力和随机力的影响下的时间演化。

对于任意满足$\sigma_t^2\le g_t^2$的函数,有:

\[\dfrac{\partial}{\partial t}p_t(x)=-\nabla_x\cdot[(f_t(x)-\dfrac12(g_t^2-\sigma_t^2)\nabla_x\log p_t(x))p_t(x)]+\dfrac12\sigma_t^2\nabla_x\cdot\nabla_xp_t(x)\]

我们发现这个方程倒推回去是这个:

\[dx=(f_t(x)-\dfrac12(g_t^2-\sigma_t^2)\nabla_x\log p_t(x))dt+\sigma_tdw\tag{12}\]

也就是说(1)和(12)有着相同的$p_t(x)$:存在不同方差的前向过程,它们产生的边际分布是一样的。

顺便有反向过程:

\[dx=(f_t(x)-\dfrac12(g_t^2+\sigma_t^2)\nabla_x\log p_t(x))dt+\sigma_tdw\]

如果此时令$\sigma_t=0$,就有:

\[dx=(f_t(x)-\dfrac12(g_t^2-\sigma_t^2)\nabla_x\log p_t(x))dt\tag{14}\]

这个被称为概率流ODE。而且它的正向过程和反向过程是一样的,这是一个确定性变换,而且是可逆的,所以就可以把流模型的相关结果拿走,总之就是好处很多。

顺带一提,(14)中$f_t$取线性函数,可以证明能得到DDIM的连续版本的ODE。

Flow Matching

流模型:想找到一个从数据分布到简单分布的可逆变换然后再求逆变换。

事实上微分方程是描述这一过程的很好工具,于是有了神经网络+微分方程。

\[\dfrac{dx_t}{dt}=f_t(x_t)\]

将$x_0$视作数据样本,$x_T$视作简单分布的样本,$p_t(x_t)$是一簇随着参数$t$连续变化的分布的概率密度函数。

利用雅可比行列式和泰勒近似,可以得到,满足如下连续性方程的$f_t(x_t)$都可以通过求解方程来实现数据分布和简单分布之间的变换。

\[\dfrac{\partial}{\partial t}p_t(x_t)=-\nabla_{x_t}\cdot(f_t(x_t)p_t(x_t))\]

而且由于这是一个$d$维变量,却只有一个方程,所以这是一个不定方程,原则上可以指定任意的$p_t(x_t)$(而不仅限定$t=0,T$两个边界)来求解$f_t(x_t)$。

记$u(t,x_t)=(p_t(x_t),f_t(x_t)p_t(x_t))$于是原先方程可以改写成:

\[\nabla_{(t,x_t)}\cdot u(t,x_t)=0\]

约束条件

\[u_1(0,x_0)=p_0(x_0)\] \[\int u_1(t,x_t)dx_t=1\]

给定目标概率密度路径$p_t$和对应的向量场$u(t,x_t)$,流匹配目标定义为:

\[\mathcal{L}_{FM}(\theta)=\mathbb{E}_{t,p_t(x)}||v(t,x_t)-u(t,x_t)||^2\tag{FM}\]

这里$\theta$表示CNF向量场$v_t$的可学习参数,$t$是均匀分布,$x\sim p_t(x)$。简而言之,FM损失用神经网络$v_t$回归向量场$u_t$,在达到零损失时,学习的CNF模型将生成$p_t$。

$x_1$是某个遵循未知数据分布$q(x_1)$的随机变量,我们只能访问满足其分布的数据样本而不能访问密度函数本身。流匹配是一个简单而有吸引力的目标,但从本质上讲,它很难在实践中使用,因为我们不知道什么是适当的$p_t$和$u_t$。有很多概率路径可以满足$p_1(x)\approx q(x)$,更重要的是,我们一般无法访问生成$p_t$的形式$u_t$。

Conditional Flow Matching

构造目标概率路径的一种简单方法是通过更简单的概率路径的混合:给定特定的数据样本$x_1$,用$p_t(xx_1)$表示条件概率路径,使$p_0(xx_1)=p(x)$。设计$p_1(xx_1)$是$x=x_1$附近的分布,例如$\mathcal{N}(xx_1,\sigma^2I)$,$\sigma$是个小正数。对全体$q(x_1)$积分有边际概率路径:
\[p_t(x)=\int p_t(x|x_1)q(x_1)dx_1\] \[p_1(x)=\int p_1(x|x_1)q(x_1)dx_1\approx q(x)\] \[u_t(x)=\int u_t(x|x_1)\dfrac{p_t(x|x_1)q(x_1)}{p_t(x)}dx_1\]

考虑条件流匹配目标:

\[\mathcal{L}_{CFM}=\mathbb{E}_{t,q(x_1),p_t(x|x_1)}||v_t(x)-u_t(x|x_1)||^2\tag{CFM}\]

定理:设$p_t(x)>0$,则

\[\nabla_\theta\mathcal{L}_{FM}(\theta)=\nabla_\theta\mathcal{L}_{CFM}(\theta)\]

(把平方项展开作差,结果发现和$\theta$无关)

直线生成

我们希望,随机选定$x_0\sim p_0(x_0)$,$x_T\sim p_T(x_T)$,假设它们按照轨迹

\[x_t=\varphi(x_0,x_T)\]

进行变换,这个变换可以自行设计,写出微分方程:

\[\dfrac{dx_t}{dt}=\dfrac{\partial \varphi(x_0,x_T)}{\partial t}\]

但是计算的时候不知道$x_0$,希望学习$s_\theta(x_t,t)$来逼近右端。

\[\mathbb{E}_{x_0\sim p_0,x_T\sim p_T(x_T)}\left[||s_\theta(x_t,t)-\dfrac{\partial \varphi(x_0,x_T)}{\partial t}||^2\right]\]

推理过程就可以使用

\[\dfrac{dx_t}{dt}=s_\theta(x_t,t)\]

以直线为例,假设

\[x_t=\varphi_t(x_0,x_1)=(x_1-x_0)t+x_0\]

那么

\[\dfrac{\partial \varphi(x_0,x_T)}{\partial t}=x_1-x_0\]

于是训练目标

\[\mathbb{E}_{x_0\sim p_0,x_T\sim p_T(x_T)}\left[||s_\theta((x_1-x_0)t+x_0,t)-(x_1-x_0)||^2\right]\]

\[\mathbb{E}_{x_0,x_t\sim p_0(x_0)p_t(x_t|x_0)}\left[||s_\theta(x_t,t)-\dfrac{x_t-x_0}{t}||^2\right]\]

最优运输

最优运输问题:

\[\inf_{\pi\in\Pi(\mu,\nu)}\int c(x,y)d\pi(x,y)\]
$c(x,y)=\dfrac12 x-y ^2$为转移损失。

由于

\[\dfrac12||x-y||^2=\inf_{\substack{f\in C^1([0,1],\mathbb{R}^n)\\f(0)=x,f(1)=y}}\int_0^1\dfrac12||\dot{f}||^2dt\]

解是$f(t)=(1-t)x+ty$,是连接$x,y$的直线。因为任何一个连接$x,y$的路径都一定给出更高值,这个问题变成了$C^1([0,1],\mathbb{R}^n)$上:

\[\inf_{P\in\mathbb{D}^1(\mu,\nu)}\mathbb{E}_P(\int_0^1\dfrac12||\dot x||^2dt)\]

其中$\mathbb{D}$是概率测度空间。

考虑流体力学的图景(?)

\[\dfrac12||x-y||^2=\inf_{v\in\mathcal{V}_y}\int_0^1\dfrac12||v(x^v(t),t)||^2dt\] \[\dot x^v(t)=v(x^v(t),t), x(0)=x\]

这是一个最优控制问题,要求满足$x^v(1)=y$

所以最优运输问题等价于如下的随机控制问题

\[\inf_{v\in\mathcal{V}}\int_0^1\dfrac12||v(x^v(t),t)||^2dt\] \[\dot x^v(t)=v(x^v(t),t), x(0)\sim\mu, x(1)\sim\nu\]

设$\rho_0$,$\rho_1$是概率密度,$\mathbb{D}(\rho_0,\rho_1)$是规定边缘密度的分布集,给定$P\in\mathbb{D}$,如下问题:

\[\min H(Q,P),Q\in\mathbb{D}(\rho_0,\rho_1)\]

其中$H$是相对熵(KL散度):

\[H(Q,P)=\mathbb{E}_Q(\ln\dfrac{dQ}{dP})\]

这种问题在两个边际间寻找最大可能的演化。如果解存在,就称为相对于$P$的薛定谔桥。

薛定谔桥:考虑两个分布$p_\text{data}$和$p_\text{prior}$,和参考随机过程$\pi_\text{ref}$,希望找到一个随机过程$\pi^*$使得:

\[\pi^*=\argmin\{D_{KL}(\pi||\pi_\text{ref}):\pi_0=p_\text{data},\pi_N=p_\text{prior}\}\]

可以用IPF (Iterative Proportional Fitting)来解SB问题:

\[\pi^{2n+1}=\argmin\{D_{KL}(\pi||\pi^{2n}):\pi_N=p_\text{prior}\}\\ \pi^{2n+2}=\argmin\{D_{KL}(\pi||\pi^{2n+1}):\pi_0=p_\text{data}\}\]
本文由作者按照 CC BY 4.0 进行授权