历史总是惊人地相似。当初笔者在写《生成扩散模型漫谈（十四）：构建ODE的一般步骤（上）》（当时还没有“上”这个后缀）时，以为自己已经搞清楚了构建ODE式扩散的一般步骤，结果读者 @gaohuazuo 就给出了一个新的直观有效的方案，这直接导致了后续《生成扩散模型漫谈（十四）：构建ODE的一般步骤（中）》（当时后缀是“下”）。而当笔者以为事情已经终结时，却发现ICLR2023的论文《Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow》又给出了一个构建ODE式扩散模型的新方案，其简洁、直观的程度简直前所未有，令人拍案叫绝。所以笔者只好默默将前一篇的后缀改为“中”，然后写了这个“下”篇来分享这一新的结果。

直观结果 #

我们知道，扩散模型是一个xT→x0的演化过程，而ODE式扩散模型则指定演化过程按照如下ODE进行：
dxtdt=ft(xt)


而所谓构建ODE式扩散模型，就是要设计一个函数ft(xt)，使其对应的演化轨迹构成给定分布pT(xT)、p0(x0)之间的一个变换。说白了，我们希望从pT(xT)中随机采样一个xT，然后按照上述ODE向后演化得到的x0是∼p0(x0)的。

原论文的思路非常简单，随机选定x0∼p0(x0),xT∼pT(xT)，假设它们按照轨迹
xt=φt(x0,xT)


进行变换。这个轨迹是一个已知的函数，是我们自行设计的部分，理论上只要满足
x0=φ0(x0,xT),xT=φT(x0,xT)

的连续函数都可以。接着我们就可以写出它满足的微分方程：
dxtdt=∂φt(x0,xT)∂t

但这个微分方程是不实用的，因为我们想要的是给定xT来生成x0，但它右端却是x0的函数（如果已知x0就完事了），只有像式(1)那样右端只含有xt的ODE（单从因果关系来看，理论上也可以包含xT，但我们一般不考虑这种情况）才能进行实用的演化。那么，一个直观又“异想天开”的想法是：学一个函数vθ(xt,t)尽量逼近上式右端！为此，我们优化如下目标：
Ex0∼p0(x0),xT∼pT(xT)[‖vθ(xt,t)−∂φt(x0,xT)∂t‖2]

由于vθ(xt,t)尽量逼近了∂φt(x0,xT)∂t，所以我们认为将方程(4)的右端替换为vθ(xt,t)也是成立的，这就得到实用的扩散ODE：
dxtdt=vθ(xt,t)

简单例子 #

作为简单的例子，我们设T=1，并设变化轨迹是直线
xt=φt(x0,x1)=(x1−x0)t+x0


那么
∂φt(x0,xT)∂t=x1−x0

所以训练目标(5)就是：
Ex0∼p0(x0),xT∼pT(xT)[‖vθ((x1−x0)t+x0,t)−(x1−x0)‖2]

或者等价地写成
Ex0,xt∼p0(x0)pt(xt|x0)[‖vθ(xt,t)−xt−x0t‖2]

这就完事了！结果跟《生成扩散模型漫谈（十四）：构建ODE的一般步骤（中）》的“直线轨迹”例子是完全一致的，也是原论文主要研究的模型，被称为“Rectified Flow”。

从这个直线例子的过程也可以看出，通过该思路来构建扩散ODE的步骤只有寥寥几行，相比之前的过程是大大简化了，简单到甚至让人有种“颠覆了对扩散模型的印象”的不可思议之感。

证明过程 #

然而，迄今为止前面“直观结果”一节的结论只能算是一个直观的猜测，因为我们还没有从理论上证明优化目标(5)所得到的方程(6)的确实现了分布pT(xT)、p0(x0)之间的变换。

为了证明这一结论，笔者一开始是想证明目标(5)的最优解满足连续性方程：
∂pt(xt)∂t=−∇xt⋅(pt(xt)vθ(xt,t))


如果满足，那么根据连续性方程与ODE的对应关系（参考《生成扩散模型漫谈（十二）：“硬刚”扩散ODE》、《测试函数法推导连续性方程和Fokker-Planck方程》），方程(6)确实是分布pT(xT)、p0(x0)之间的一个变换。

但仔细想一下，这个思路似乎有点迂回了，因为根据文章《测试函数法推导连续性方程和Fokker-Planck方程》，连续性方程本身就是由ODE通过
Ext+Δt[ϕ(xt+Δt)]=Ext[ϕ(xt+ft(xt)Δt)]


推出的，所以按理说(12)更基本，我们只需要证明(5)的最优解满足它就行。也就是说，我们想要找到一个纯粹是xt的函数ft(xt)满足(12)，然后发现它正好是(5)的最优解。

于是，我们写出（简单起见，φt(x0,xT)简写为φt）
Ext+Δt[ϕ(xt+Δt)]=Ex0,xT[ϕ(φt+Δt)]=Ex0,xT[ϕ(φt)+Δt∂φt∂t⋅∇φtϕ(φt)]=Ex0,xT[ϕ(xt)]+ΔtEx0,xT[∂φt∂t⋅∇xtϕ(xt)]=Ext[ϕ(xt)]+ΔtEx0,xT[∂φt∂t⋅∇xtϕ(xt)]


其中第一个等号是因为式(2)，第二个等号是泰勒展开到一阶，第三个等号同样是式(2)，第四个等号就是因为xt是x0,xT的确定性函数，所以关于x0,xT的期望就是关于xt的期望。

我们看到，∂φt∂t是x0,xT的函数，接下来我们再做一个假设：式(2)关于xT是可逆的。这个假设意味着我们可以从式(2)中解出xT=ψt(x0,xt)，这个结果可以代入∂φt∂t，使它变为x0,xt的函数。所以我们有
Ext+Δt[ϕ(xt+Δt)]=Ext[ϕ(xt)]+ΔtEx0,xT[∂φt∂t⋅∇xtϕ(xt)]=Ext[ϕ(xt)]+ΔtEx0,xt[∂φt∂t⋅∇xtϕ(xt)]=Ext[ϕ(xt)]+ΔtExt[Ex0|xt[∂φt∂t]⏟xt的函数⋅∇xtϕ(xt)]=Ext[ϕ(xt+ΔtEx0|xt[∂φt∂t])]


其中第二个等号是因为∂φt∂t已经改为x0,xt的函数，所以第二项期望的随机变量改为x0,xt；第三个等号则是相当于做了分解p(x0,xt)=p(x0|xt)p(xt)，此时x0,xt不是独立的，所以要注明x0|xt，即x0是依赖于xt的。注意∂φt∂t原本是x0,xt的函数，现在对x0求期望后，剩下的唯一自变量就是xt，后面我们会看到它就是我们要找的纯粹是xt的函数！第四个等号，就是利用泰勒展开公式将两项重新合并起来。

现在，我们得到了
Ext+Δt[ϕ(xt+Δt)]=Ext[ϕ(xt+ΔtEx0|xt[∂φt∂t])]


对于任意测试函数ϕ成立，所以这意味着
xt+Δt=xt+ΔtEx0|xt[∂φt∂t]⇒dxtdt=Ex0|xt[∂φt∂t]

就是我们要寻找的ODE。根据
Ex[x]=argminμEx[‖x−μ‖2]

式(16)的右端正好是训练目标(5)的最优解，这就证明了优化训练目标(5)得出的方程(6)的确实现了分布pT(xT)、p0(x0)之间的变换。

读后感受 #

关于“直观结果”中的构建扩散ODE的思路，原论文的作者还写了篇知乎专栏文章《[ICLR2023] 扩散生成模型新方法：极度简化，一步生成》，大家也可以去读读。读者也是在这篇专栏中首次了解到该方法的，并深深为之震惊和叹服。

如果读者读过《生成扩散模型漫谈（十四）：构建ODE的一般步骤（中）》，那么就会更加体会到该思路的简单直接，也更能理解笔者为何如此不吝赞美之词。不怕大家笑话，笔者在写“中篇”（当时的“下篇”）的时候，是考虑过式(2)所描述的轨迹的，但是在当时的框架下，根本没法推演下去，最后以失败告终，当时完全想不到它能以一种如此简捷的方式进行下去。所以，写这个扩散ODE系列真的让人有种“人比人，气死人”的感觉，“中篇”、“下篇”就是自己智商被一次次“降维打击”的最好见证。

读者可能想问，还会不会有更简单的第四篇，让笔者再一次经历降维打击？可能有，但概率真的很小了，真的很难想象会有比这更简单的构建步骤了。“直观结果”一节看上去很长，但实际步骤就只有两步：1、随便选择一个渐变轨迹；2、用xt的函数去逼近渐变轨迹对t的导数。就这样的寥寥两步，还能怎么再简化呢？甚至说，“证明过程”一节的推导也是相当简单的了，虽然写得长，但本质就是求个导，然后变换一下求期望的分布，比前两篇的过程简单了可不止一丁半点。总而言之，亲自完成过ODE扩散的前两篇推导的读者就能深刻感觉到，这一篇的思路是真的简单，简单到让我们觉得已经无法再简单了。

此外，除了提供构建扩散ODE的简单思路外，原论文还讨论了Rectified Flow跟最优传输之间的联系，以及如何用这种联系来加速采样过程，等等。但这部分内容并不是本文主要关心的，所以等以后有机会我们再讨论它们。

文章小结 #

本文介绍了Rectified Flow一文中提出的构建ODE式扩散模型的一种极其简单直观的思路，并给出了自己的证明过程。

转载到请包括本文地址：https://spaces.ac.cn/archives/9497

更详细的转载事宜请参考：《科学空间FAQ》

如果您还有什么疑惑或建议，欢迎在下方评论区继续讨论。

如果您觉得本文还不错，欢迎分享/打赏本文。打赏并非要从中获得收益，而是希望知道科学空间获得了多少读者的真心关注。当然，如果你无视它，也不会影响你的阅读。再次表示欢迎和感谢！

如果您需要引用本文，请参考：

苏剑林. (Feb. 23, 2023). 《生成扩散模型漫谈（十七）：构建ODE的一般步骤（下） 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9497

@online{kexuefm-9497,
        title={生成扩散模型漫谈（十七）：构建ODE的一般步骤（下）},
        author={苏剑林},
        year={2023},
        month={Feb},
        url={\url{https://spaces.ac.cn/archives/9497}},
}