生成扩散模型漫谈(十七):构建ODE的一般步骤(下)
By 苏剑林 | 2023-02-23 | 76244位读者 |历史总是惊人地相似。当初笔者在写《生成扩散模型漫谈(十四):构建ODE的一般步骤(上)》(当时还没有“上”这个后缀)时,以为自己已经搞清楚了构建ODE式扩散的一般步骤,结果读者 @gaohuazuo 就给出了一个新的直观有效的方案,这直接导致了后续《生成扩散模型漫谈(十四):构建ODE的一般步骤(中)》(当时后缀是“下”)。而当笔者以为事情已经终结时,却发现ICLR2023的论文《Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow》又给出了一个构建ODE式扩散模型的新方案,其简洁、直观的程度简直前所未有,令人拍案叫绝。所以笔者只好默默将前一篇的后缀改为“中”,然后写了这个“下”篇来分享这一新的结果。
直观结果 #
我们知道,扩散模型是一个$\boldsymbol{x}_T\to \boldsymbol{x}_0$的演化过程,而ODE式扩散模型则指定演化过程按照如下ODE进行:
\begin{equation}\frac{d\boldsymbol{x}_t}{dt}=\boldsymbol{f}_t(\boldsymbol{x}_t)\label{eq:ode}\end{equation}
而所谓构建ODE式扩散模型,就是要设计一个函数$\boldsymbol{f}_t(\boldsymbol{x}_t)$,使其对应的演化轨迹构成给定分布$p_T(\boldsymbol{x}_T)$、$p_0(\boldsymbol{x}_0)$之间的一个变换。说白了,我们希望从$p_T(\boldsymbol{x}_T)$中随机采样一个$\boldsymbol{x}_T$,然后按照上述ODE向后演化得到的$\boldsymbol{x}_0$是$\sim p_0(\boldsymbol{x}_0)$的。
原论文的思路非常简单,随机选定$\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T)$,假设它们按照轨迹
\begin{equation}\boldsymbol{x}_t = \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)\label{eq:track}\end{equation}
进行变换。这个轨迹是一个已知的函数,是我们自行设计的部分,理论上只要满足
\begin{equation}\boldsymbol{x}_0 = \boldsymbol{\varphi}_0(\boldsymbol{x}_0, \boldsymbol{x}_T),\quad \boldsymbol{x}_T = \boldsymbol{\varphi}_T(\boldsymbol{x}_0, \boldsymbol{x}_T)\end{equation}
的连续函数都可以。接着我们就可以写出它满足的微分方程:
\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}\label{eq:fake-ode}\end{equation}
但这个微分方程是不实用的,因为我们想要的是给定$\boldsymbol{x}_T$来生成$\boldsymbol{x}_0$,但它右端却是$\boldsymbol{x}_0$的函数(如果已知$\boldsymbol{x}_0$就完事了),只有像式$\eqref{eq:ode}$那样右端只含有$\boldsymbol{x}_t$的ODE(单从因果关系来看,理论上也可以包含$\boldsymbol{x}_T$,但我们一般不考虑这种情况)才能进行实用的演化。那么,一个直观又“异想天开”的想法是:学一个函数$\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$尽量逼近上式右端!为此,我们优化如下目标:
\begin{equation}\mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T)}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}\right\Vert^2\right] \label{eq:objective}
\end{equation}
由于$\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$尽量逼近了$\frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}$,所以我们认为将方程$\eqref{eq:fake-ode}$的右端替换为$\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$也是成立的,这就得到实用的扩散ODE:
\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\label{eq:s-ode}\end{equation}
简单例子 #
作为简单的例子,我们设$T=1$,并设变化轨迹是直线
\begin{equation}\boldsymbol{x}_t = \boldsymbol{\varphi}_t(\boldsymbol{x}_0,\boldsymbol{x}_1) = (\boldsymbol{x}_1 - \boldsymbol{x}_0)t + \boldsymbol{x}_0\end{equation}
那么
\begin{equation}\frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t} = \boldsymbol{x}_1 - \boldsymbol{x}_0\end{equation}
所以训练目标$\eqref{eq:objective}$就是:
\begin{equation}\mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_T\sim p_T(\boldsymbol{x}_T)}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}\big((\boldsymbol{x}_1 - \boldsymbol{x}_0)t + \boldsymbol{x}_0, t\big) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right\Vert^2\right]\end{equation}
或者等价地写成
\begin{equation}\mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{x}_t\sim p_0(\boldsymbol{x}_0)p_t(\boldsymbol{x}_t|\boldsymbol{x}_0)}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \frac{\boldsymbol{x}_t - \boldsymbol{x}_0}{t}\right\Vert^2\right]\end{equation}
这就完事了!结果跟《生成扩散模型漫谈(十四):构建ODE的一般步骤(中)》的“直线轨迹”例子是完全一致的,也是原论文主要研究的模型,被称为“Rectified Flow”。
从这个直线例子的过程也可以看出,通过该思路来构建扩散ODE的步骤只有寥寥几行,相比之前的过程是大大简化了,简单到甚至让人有种“颠覆了对扩散模型的印象”的不可思议之感。
证明过程 #
然而,迄今为止前面“直观结果”一节的结论只能算是一个直观的猜测,因为我们还没有从理论上证明优化目标$\eqref{eq:objective}$所得到的方程$\eqref{eq:s-ode}$的确实现了分布$p_T(\boldsymbol{x}_T)$、$p_0(\boldsymbol{x}_0)$之间的变换。
为了证明这一结论,笔者一开始是想证明目标$\eqref{eq:objective}$的最优解满足连续性方程:
\begin{equation}\frac{\partial p_t(\boldsymbol{x}_t)}{\partial t} = -\nabla_{\boldsymbol{x}_t}\cdot\big(p_t(\boldsymbol{x}_t)\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)\big)\end{equation}
如果满足,那么根据连续性方程与ODE的对应关系(参考《生成扩散模型漫谈(十二):“硬刚”扩散ODE》、《测试函数法推导连续性方程和Fokker-Planck方程》),方程$\eqref{eq:s-ode}$确实是分布$p_T(\boldsymbol{x}_T)$、$p_0(\boldsymbol{x}_0)$之间的一个变换。
但仔细想一下,这个思路似乎有点迂回了,因为根据文章《测试函数法推导连续性方程和Fokker-Planck方程》,连续性方程本身就是由ODE通过
\begin{equation}\mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] = \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t + \boldsymbol{f}_t(\boldsymbol{x}_t)\Delta t)\right]\label{eq:base}\end{equation}
推出的,所以按理说$\eqref{eq:base}$更基本,我们只需要证明$\eqref{eq:objective}$的最优解满足它就行。也就是说,我们想要找到一个纯粹是$\boldsymbol{x}_t$的函数$\boldsymbol{f}_t(\boldsymbol{x}_t)$满足$\eqref{eq:base}$,然后发现它正好是$\eqref{eq:objective}$的最优解。
于是,我们写出(简单起见,$\boldsymbol{\varphi}_t(\boldsymbol{x}_0,\boldsymbol{x}_T)$简写为$\boldsymbol{\varphi}_t$)
\begin{equation}\begin{aligned}
\mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] =&\, \mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\phi(\boldsymbol{\varphi}_{t+\Delta t})\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\phi(\boldsymbol{\varphi}_t) + \Delta t\,\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{\varphi}_t}\phi(\boldsymbol{\varphi}_t)\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\
\end{aligned}\end{equation}
其中第一个等号是因为式$\eqref{eq:track}$,第二个等号是泰勒展开到一阶,第三个等号同样是式$\eqref{eq:track}$,第四个等号就是因为$\boldsymbol{x}_t$是$\boldsymbol{x}_0,\boldsymbol{x}_T$的确定性函数,所以关于$\boldsymbol{x}_0,\boldsymbol{x}_T$的期望就是关于$\boldsymbol{x}_t$的期望。
我们看到,$\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$是$\boldsymbol{x}_0,\boldsymbol{x}_T$的函数,接下来我们再做一个假设:式$\eqref{eq:track}$关于$\boldsymbol{x}_T$是可逆的。这个假设意味着我们可以从式$\eqref{eq:track}$中解出$\boldsymbol{x}_T=\boldsymbol{\psi}_t(\boldsymbol{x}_0,\boldsymbol{x}_t)$,这个结果可以代入$\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$,使它变为$\boldsymbol{x}_0,\boldsymbol{x}_t$的函数。所以我们有
\begin{equation}\begin{aligned}
\mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] =&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_T}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0, \boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}_t)\right] + \Delta t\,\mathbb{E}_{\boldsymbol{x}_t}\left[\underbrace{\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]}_{\boldsymbol{x}_t\text{的函数}}\cdot\nabla_{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right] \\
=&\, \mathbb{E}_{\boldsymbol{x}_t}\left[\phi\left(\boldsymbol{x}_t + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\right)\right]
\end{aligned}\end{equation}
其中第二个等号是因为$\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$已经改为$\boldsymbol{x}_0,\boldsymbol{x}_t$的函数,所以第二项期望的随机变量改为$\boldsymbol{x}_0,\boldsymbol{x}_t$;第三个等号则是相当于做了分解$p(\boldsymbol{x}_0,\boldsymbol{x}_t)=p(\boldsymbol{x}_0|\boldsymbol{x}_t)p(\boldsymbol{x}_t)$,此时$\boldsymbol{x}_0,\boldsymbol{x}_t$不是独立的,所以要注明$\boldsymbol{x}_0|\boldsymbol{x}_t$,即$\boldsymbol{x}_0$是依赖于$\boldsymbol{x}_t$的。注意$\frac{\partial \boldsymbol{\varphi}_t}{\partial t}$原本是$\boldsymbol{x}_0,\boldsymbol{x}_t$的函数,现在对$\boldsymbol{x}_0$求期望后,剩下的唯一自变量就是$\boldsymbol{x}_t$,后面我们会看到它就是我们要找的纯粹是$\boldsymbol{x}_t$的函数!第四个等号,就是利用泰勒展开公式将两项重新合并起来。
现在,我们得到了
\begin{equation}\mathbb{E}_{\boldsymbol{x}_{t+\Delta t}}\left[\phi(\boldsymbol{x}_{t+\Delta t})\right] = \mathbb{E}_{\boldsymbol{x}_t}\left[\phi\left(\boldsymbol{x}_t + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\right)\right]\end{equation}
对于任意测试函数$\phi$成立,所以这意味着
\begin{equation}\boldsymbol{x}_{t+\Delta t} = \boldsymbol{x}_t + \Delta t\,\mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\quad\Rightarrow\quad\frac{d\boldsymbol{x}_t}{dt} = \mathbb{E}_{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\label{eq:real-ode}\end{equation}
就是我们要寻找的ODE。根据
\begin{equation}\mathbb{E}_{\boldsymbol{x}}[\boldsymbol{x}] = \mathop{\text{argmin}}_{\boldsymbol{\mu}}\mathbb{E}_{\boldsymbol{x}}\left[\Vert \boldsymbol{x} - \boldsymbol{\mu}\Vert^2\right]\label{eq:mean-opt}\end{equation}
式$\eqref{eq:real-ode}$的右端正好是训练目标$\eqref{eq:objective}$的最优解,这就证明了优化训练目标$\eqref{eq:objective}$得出的方程$\eqref{eq:s-ode}$的确实现了分布$p_T(\boldsymbol{x}_T)$、$p_0(\boldsymbol{x}_0)$之间的变换。
读后感受 #
关于“直观结果”中的构建扩散ODE的思路,原论文的作者还写了篇知乎专栏文章《[ICLR2023] 扩散生成模型新方法:极度简化,一步生成》,大家也可以去读读。读者也是在这篇专栏中首次了解到该方法的,并深深为之震惊和叹服。
如果读者读过《生成扩散模型漫谈(十四):构建ODE的一般步骤(中)》,那么就会更加体会到该思路的简单直接,也更能理解笔者为何如此不吝赞美之词。不怕大家笑话,笔者在写“中篇”(当时的“下篇”)的时候,是考虑过式$\eqref{eq:track}$所描述的轨迹的,但是在当时的框架下,根本没法推演下去,最后以失败告终,当时完全想不到它能以一种如此简捷的方式进行下去。所以,写这个扩散ODE系列真的让人有种“人比人,气死人”的感觉,“中篇”、“下篇”就是自己智商被一次次“降维打击”的最好见证。
读者可能想问,还会不会有更简单的第四篇,让笔者再一次经历降维打击?可能有,但概率真的很小了,真的很难想象会有比这更简单的构建步骤了。“直观结果”一节看上去很长,但实际步骤就只有两步:1、随便选择一个渐变轨迹;2、用$\boldsymbol{x}_t$的函数去逼近渐变轨迹对$t$的导数。就这样的寥寥两步,还能怎么再简化呢?甚至说,“证明过程”一节的推导也是相当简单的了,虽然写得长,但本质就是求个导,然后变换一下求期望的分布,比前两篇的过程简单了可不止一丁半点。总而言之,亲自完成过ODE扩散的前两篇推导的读者就能深刻感觉到,这一篇的思路是真的简单,简单到让我们觉得已经无法再简单了。
此外,除了提供构建扩散ODE的简单思路外,原论文还讨论了Rectified Flow跟最优传输之间的联系,以及如何用这种联系来加速采样过程,等等。但这部分内容并不是本文主要关心的,所以等以后有机会我们再讨论它们。
文章小结 #
本文介绍了Rectified Flow一文中提出的构建ODE式扩散模型的一种极其简单直观的思路,并给出了自己的证明过程。
转载到请包括本文地址:https://kexue.fm/archives/9497
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Feb. 23, 2023). 《生成扩散模型漫谈(十七):构建ODE的一般步骤(下) 》[Blog post]. Retrieved from https://kexue.fm/archives/9497
@online{kexuefm-9497,
title={生成扩散模型漫谈(十七):构建ODE的一般步骤(下)},
author={苏剑林},
year={2023},
month={Feb},
url={\url{https://kexue.fm/archives/9497}},
}
February 24th, 2023
在知乎原论文作者下看到苏神的评论
没过几天没想到苏神就发文章了
虽然被原作者打击到了,但还是要学习和分享的,要不然不进则退啊~
March 7th, 2023
原论文的得定理3.3好像就是为了证明这个?原论文证明看起来更简约,但是看不懂,大佬可以讲讲吗?
原论文的证明是“简短”,不是“简单”。原论文的证明过程详细展开的话,就是用 https://kexue.fm/archives/9461 中的方法先证明连续性方程,然后把本文的证明进一步往连续性方程靠,更加迂回。理论上不大可能有比本文更简单的证明思路了。
March 9th, 2023
在系列(15)https://kexue.fm/archives/9379中看到轨迹簇的写法$\boldsymbol{\varphi}_t(\boldsymbol{x}_t|\boldsymbol{x}_0) = \boldsymbol{x}_T$时候就感觉很奇怪,这里的新的表示方式$\boldsymbol{x}_t = \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)$更加显式了。
而且其实本文中的$\frac{d\boldsymbol{x}_t}{dt} = \frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}$(通过带入系列(15)$(7)$式并对比系列(15)$(8)$式)应该和系列(15)中的$\boldsymbol{f}_t(\boldsymbol{x}_t|\boldsymbol{x}_0)$是一回事,所以两篇文章之中的训练目标应该也一样?(本文优化目标$(5)$似乎要加一个对时间t的采样?)
其实本文的推导也用到了“从$\boldsymbol{x}_t = \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)$解出$\boldsymbol{x}_T = \cdots$“这一步,所以单看这个设定来说,两者半斤八两,没有谁优谁劣了。
理论上来说,两篇文章的结果应该可以互推,但很遗憾,我尝试了,实在推不下去(捂脸)。
March 20th, 2023
如果这个模型能一步完成生成,那么有没有可能用VAE,换成图像大小级别的潜在空间,也能达到差不多的效果啊?
已经测试过了,虽然换成图像大小级别的潜在空间,VAE的效果有所提升,不过由于MSE损失的存在,生成效果还是很模糊。换用GAN后(同样是图像大小的潜在空间),效果就比较符合预期了,不过这样也只是实现了高斯分布到图像分布的单射,没什么实际作用,本来想换成o-gan的,不过不知道为啥,换成图像大小的潜在空间后,它就没法收敛了。
1、VAE“换成图像大小级别的潜在空间”,随机采样效果还可以吗?GAN居然能成?
2、O-GAN可能要调一下模型结构(判别器可能要换成U-Net)和相关系数的权重吧。
1. VAE的编码器和解码器其实用的就是U-Net架构,生成效果其实一般,不如NVAE,也许把编码器换成NVAE就行?。GAN生成器用的U-Net,判别器还是分类器架构,这样确实能成,生成效果也不错,没有一直训练下去,不知道它会不会模式崩溃。
2. O-GAN确实用的U-Net架构,这时相关系数我是把每个像素的相关系数都求了一遍再求和,判别器我也是对每个像素的特征向量求均值,不过这样设置似乎没有生效
可能维度大了不好控制吧,我也没有经验,我直觉也是无法成功的。
May 5th, 2023
宋彪博士新出的一致性模型感觉和这个一样走直线,但是却可以从头直接训练一步模型,不用进行蒸馏训练。它们之间有什么内在联系吗
reflow也可以直接训练,训练出来的模型也可以当作一步模型,只不过没有蒸馏后的效果好而已。
我反复粗看了几次relfow和songyang的模型(没细看倒是真的),都没发现这两个模型在数学上有什么区别。
reflow因为首次进行了成对采样,这样会产生两条直线的交叉点,所以至少需要1次蒸馏后才会得到不错的1步模型。我看了他们的代码,发现一致性模型通过用s(xt,t)==s_ema(xt+1, t+1)来进行约束,随着训练进行,t+1和t之间的间隔越来越小,相当于在训练中“包含了”蒸馏,所以可以达到训练一次就能获得一个还不错的一步模型的效果。把这个技巧放在reflow上面,它们还真的没什么差别了。
感谢指点,解惑了。
整理一下方法,方便大家进行纠错和验证,可以把训练目标换成s(x1,t1)=s_ema(x2, t2), t2 > t1, 即在直线运动的假设下,一条轨迹上不同时间点的速度是相同的,同时轨迹的起点终点相同,t=1时到达终点。那么有x1 + s(x1, t1) * (1 - t1) = x2 + s_ema(x2, t2) * (1-t2),最终的训练目标变为s(x1,t1) = ((x2 - x1 + s_ema(x2, 2)*(1 - t2) ) / (1 - t1)。
比起一致性模型,这个训练目标算是非常简洁了,我也试了一下在人脸上的效果,可以在5步的时候达到不错的效果,但是2步的时候效果较差,可能模型太小吧
有点疑惑,“把训练目标换成s(x1,t1)=s_ema(x2, t2)”是您提出的方案,还是一致性模型的方案?看你前一个回复,好像说的是一致性模型的方案,但这个回复你又说“比起一致性模型”,感觉又像是您自己提出的方案?
其实二者的出发点是一样的,只是一致性模型在分数模型视角下导向了不同的损失函数,另外可能为了稳定训练,它的代码里面引入了过多的超参数,既要设计dt的变化过程,又引入了与t相关的插值函数。相比较而言,这里的训练目标没有任何超参数的设置,同时能极快收敛,“简洁性”是对此而言的。
两者应该是一样的,因为我训练完后,用一致性模型的采样方式来生成图片,也就是在ode采样器的基础上,下一步的dt是上一步的dt的2倍
好的,迟点我也好好读读,感谢指点。
@o_glay|comment-21673
您好,我不太明白您说的“相当于在训练中包含了蒸馏”,您可以继续解释一下吗?
我认为consistency model(CM)一步就可以达到比较好的效果是因为reflow的配对是随机生成的,中间存在交叉、不合理的配对,所以reflow先需要优化配对,并在优化了的配对的基础上再reflow才能产生直的轨迹。
相比之下,CM仍然是提前设定好的Gaussian schedule,本身配对较为合理,所以无需多次训练。
我认为您说的“一致性模型走直线”也有一些偏差,CM的轨迹并不是直线的。
CM之所以可以一步生成是靠$x_0=f(x_t,t)$这个定义和对应这一目标的loss保证的,也就是说无论是蒸馏或不蒸馏的训练方式,都可以认为是将原本轨迹是$\frac{dx_t}{dt}=-ts_\theta(x_t,t)\triangle t$的常规扩散模型蒸馏成直接输出$x_0$的模型。具体来说假设常规模型的参数是$\phi$,它可以通过某种方式得到一个一致性模型$f_\theta(x,t_n;\phi)$,这种方式可以是从$x_t$一步步采样得到$x_0$,不管什么方式,总之得到$f_\theta(x,t_n;\phi)$的结果是困难的,所以想学习一个$f_\theta(x,t_n)$来逼近它,原文证明了他设定的$f(x_{t-1},t-1)==f_{ema}(x_t,t)$能达到上面的目标。
所以CM应该是不能直接优化s(xt,t)==s_ema(xt+1, t+1)的。
另外,我猜想s(x1,t1)==s_ema(x2, t2)可能使reflow一次训练(您的实验有效果),是因为这种目标函数规避了显式的配对$X_0$和$X_1$。
如有错误,还请苏神和这位大佬@o_glay指正,谢谢~
确实之前没细读CM论文,我也再次考虑了下我这里的方法为何生效,由于x1=(x-z)*t1 + z, x2 = (x-z) *t2 + z, 整个优化目标就成了 ((x - z) * dt + (1 - t2) * s_ema) / (1 - t1),(其中 dt = t2 - t1,t1和t2都是随机采样的)。它似乎在动态平衡 x - z 和 s_ema,从这里也可以看出,其实是包含了x和z的配对的。至于去拟合x-z和s_eam两个速度的加权和为何有效,我还没有想明白,它可能是将原模型需要两步完成的事用1步去完成。
谢谢回复!我也再想想
感谢两位,学到不少。
我再次做了实验,确实是很大程度上避免了显式配对X0和X1,如果按照我之前的方式生成x1=(x-z)*t1 + z, x2 = (x-z) *t2 + z, t2 > t1。一定程度上同样会造成x和z的配对,这也是为什么我此前的实验中2步生成的效果不好。而如果按照x1 = x + t1 * z,x2 = x + t2 * z,t2 < t1的方式生成,这样就基本上规避了交叉点的产生,生成时也就只需要1步就能得到好结果,我在实验中也证明了这点。如此看来,我的方法其实是假定了整个轨迹走的是条直线,大大简化了步骤,逻辑上应该也是可以用CM那套来解释的。
感谢解惑,这样就讲得通了。
可以看一下这篇arXiv新放上去的论文--Catch-Up Distillation: You Only Need to Train Once for Accelerating Sampling,应该是投今年的Neurips,粗看了一眼貌似就是一样的想法,实在是出的太快了
粗粗看了一下论文和代码,不说完全一样,也基本一致了,哈哈
佩服大佬直觉。
CM只是一种训练/蒸馏的方法,就和flow matching一样,通过选择不同的schedule,即alpha和sigma,可以学习不同类型的轨迹。
换句话说,o_glay所做的就是提出了一个自定义的轨迹,然后使用CM的方式训练。
是这个样子的,感谢
还有,大佬你把宋飏博士的名字打错了orz
May 26th, 2023
另开一楼斗胆为楼上@o_glay与@花一两位大佬的讨论做一个小结还有记录一下我的问题。
首先在 @o_glay|comment-21662 中,o_glay提出假设在CM中考虑直线的轨迹,可以将训练目标设定为各个点的速度(或者score)相等(此处相等包括方向和大小,类似于reflow,不同于CM中使用的VESDE轨迹,VESDE中速度的大小不被保证恒定),也就是$v(x_{t1},t_1)=v_{ema}(x_{t2}, t_2)$。进一步具体化o_glay提出的直线轨迹为$x_t+v(x_t,t)(1-t)=x_1$,和复杂化之后的新的训练目标$v(x_{t1},t_1) = ((x_{t2} - x_{t1} + v_{ema}(x_{t2}, t_2)*(1 - t_2) ) / (1 - t_1)$。
这里提出的轨迹其实就是reflow中的形式,也就是说,因为速度是一个常量,我们可以有$v(x_t,t)=x_1-x_0$. 然而在训练方式上,不同于reflow直接构造$(x_0,x_1)$ pair来训练,这里采用了CM形式的训练方式(因为理论上速度的大小方向恒定,这里还可以摆脱CM中$t_1, t_2$最终必须无限接近的约束)。
在@o_glay|comment-21681中,我想过为什么不直接使用$v(x_{t1},t_1)=v_{ema}(x_{t2}, t_2)$这个loss而是$v(x_{t1},t_1) = ((x_{t2} - x_{t1} + v_{ema}(x_{t2}, t_2)*(1 - t_2) ) / (1 - t_1)$这个loss。在CM的设定中,有一个重要的条件就是在$t=\sigma_{min}$的时候有$f_\theta(x_t,t)=x_0$。而在$v(x_{t1},t_1) = ((x_{t2} - x_{t1} + v_{ema}(x_{t2}, t_2)*(1 - t_2) ) / (1 - t_1)$这个loss中我们可以直接隐含的包含(例如当$t_2=1$) $v(x_{t1},t_1) = (x_{1} - x_{t1}) / (1 - t_1)$这样的类似(至少固定了一边的)边界条件。
或者我们可以从另外一个角度看新的loss为什么work。这里我感觉拟合$x-z$和$v_{ema}$两个速度的加权和这件事情,和在CM的eq(5)中引入的$c_{skip}$和$c_{out}$有异曲同工之妙。(当然了,这时候再看是否更加简洁就是另外一个问题了?)
在@o_glay|comment-21704中,o_glay大佬换了$x_t=x+tz$(也就是VESDE)形式的轨迹来获得训练的pair,但我不能够理解为什么这样就能够规避交叉点的产生?
同样还有@花一|comment-21676这里花一大佬提出的,“CM仍然是提前设定好的Gaussian schedule,本身配对较为合理”我也不是很理解,在实际操作中不都是随机的采样构造出 数据噪声对 吗?
以上内容如有错误,还请苏神和两位大佬指正,谢谢~
PS: 我在cifar10上试过$x_t=(1-t)x_0 + tz$这种形式的轨迹,当然我是和CM一样直接预测$x_0$而不是速度,不过我的实验效果一步生成好像也还可以(和原本CM中VESDE的轨迹相比)。
至于$x_t=x+tz$为什么较大程度上规避交叉产生,更多是实验上的结果吧,而且直觉上比起在$x$和$z$之间插值,在$x$基础上添加噪声应该更加不易产生交叉
我大概知道怎么解释了。
在构造训练数据对的时候,不管选择什么样的path(我是说VESDE,或者线性插值),只要是用diffusion的方式训练,也就是训练数据对是$(x_0,x_t)$(其中$x_0$是clean image),就意味着模型是在逐渐学习一整个轨迹(背后的probabilistic flow ODE),因此本身不会交叉。(参考很多理论文章里面画的ODE轨迹图)
例如在$x_0$和$z$之间插值也是一样的,训练数据对也是$(x_0,x_t)$或者$(x_t,x_{t+1})$,理论上不同$x$开始的轨迹都不会交叉。
而在用reflow的方式训练的时候,因为训练数据对是随机配对的$(x_0,z)$, 没有中间过程(learning 过程和ODE flow无关),所以交叉也会是不可避免的。
不知道这样解释行不行 0.0
奇怪了,我这里的实验也表明直接预测$x0$的效果要好于预测速度,不知道为啥
感觉预测的representation是什么很难说哪一种理论上更优越(我目前还没有看到相关的),可能别的超参数也会有影响吧
最近没时间看扩散,后面再向几位学习。
May 26th, 2023
谢谢总结,很到位,综合@花一 提供的论文中的描述来看,整个目标:$\frac{v\cdot dt + v_{est}\cdot (1 - t_2)}{1 - t_1}$基本上就是原始reflow的目标$v=x-z$和一个预估的目标$v_{est}$的加权和,至于$v_{est}$怎么来,可以通过$f_{ema}(x_2, t_2)$获得,也可以通过$k1 = f_{ema}(x_1, t_1), k2 =f_{ema}(k1*dt + x_1, t_2), v_{est}=0.5 * (k1 + k2)$获得(即二阶龙格-库塔算法,由此可以推得3阶,4阶...),阶数越高效果越好。不过@花一提及的论文比较简单粗暴,直接用$0.5 * v + 0.5 * v_{est}$作为目标,并没有动态改变加权系数。
纠正一下,之前没仔细看论文附带的代码,它的逻辑实际上是用n阶龙格-库塔算法来估计当前点$x_t$的速度$v_{est}$,然后对$v=x-z$和$v_{est}$进行加权,代码里面直接采用的$0.5 * v + 0.5 * v_{est}$做目标函数,和咱的目标函数$\frac{v\cdot dt + v_{ema}\cdot (1 - t_2)}{1 - t_1}$还是有着微妙的不同的,我把$v_{ema}$替换成用龙哥-库塔估计的$v_{est}$实际上不起作用。
June 27th, 2023
看了测试函数法(https://spaces.ac.cn/archives/9461),从(15)推出(16)好像推不出来,能详细写一下证明过程吗?
就是测试函数法的基本定义啊(那篇文章的式$(8)$)
那篇文章从(8)式推出的是f(x)ϕ(x)=g(x)ϕ(x),怎么推出自变量也相等xt+Δt=xt+ΔtEx0|xt[∂φt/∂t]搞不明白
抱歉,那篇文章多打了一个$\phi(\boldsymbol{x})$,应该是$f(\boldsymbol{x})=g(\boldsymbol{x})$(不过结果等价)。
我有点搞混了,原来你问的是这篇文章的公式$(15)$、$(16)$...
这个依据就是对于任意$\phi$,如果$\mathbb{E}_x[\phi(x)]=\mathbb{E}_y[\phi(y)]$,那么$x=y$是其中一个解,但不一定是唯一解,而本文只需要求出一个解而不是唯一解。
还是没反应过来,(15)式换个写法应该是Ex[ϕ(x)]=Ey[ϕ(y+Δ)],没法用x=y做为它的一个解吧?
如果$\mathbb{E}_x[\phi(x)]=\mathbb{E}_y[\phi(y+2)]$,那么$x=y+2$是不是一个解?或者反过来,$x=y+2$能不能推出$\mathbb{E}_x[\phi(x)]=\mathbb{E}_y[\phi(y+2)]$?
Ex[ϕ(x)]=Ey[ϕ(y+2)]可写成∫p(x)ϕ(x)dx=∫g(y)ϕ(y+2)dy, p(x),g(y)分别是x,y的概率密度。假如x=y+2成立,则∫p(x)ϕ(x)dx=∫p(y+2)ϕ(y+2)dy=∫g(y)ϕ(y+2)dy,因为ϕ任意,所以p(y+2)=g(y)即x(xt+Δt)的概率密度是y(对应xt)分密度函数的一个平移,xt不断演进而密度形状保持不变,这显然不对!
@elzyyzl|comment-22154
新开一楼吧。没看到什么显然不对的地方。
June 28th, 2023
[...]在笔者看来,《生成扩散模型漫谈(十七):构建ODE的一般步骤(下)》所介绍的ReFlow是理解扩散ODE的最直观方案,既然可以从扩散ODE的角度理解GAN,那么必定存在一个从ReFlow理解GAN的角度。经过一番尝试,笔者成功从ReFlow推出了类似WGAN-GP的结果。[...]
November 22nd, 2023
请问苏神,我的理解是这种线性ODE可以加速采样,那有没有可能有其他的效果,对下游任务有帮助的?就是说粒子轨迹没有重叠,有没有其他优良的性质?
ODE的初值决定了终值,所以ODE的轨迹必然不会有重叠的呀。除了加速之后,更优良的性质应该就是数值稳定性吧,不容易出现混沌现象。