生成扩散模型新突破:Rectified Flow如何重新定义ODE构建范式

【免费下载链接】diffusers-cd_imagenet64_l2 【免费下载链接】diffusers-cd_imagenet64_l2 项目地址: https://ai.gitcode.com/hf_mirrors/openai/diffusers-cd_imagenet64_l2

科技发展的历程总是充满意外惊喜。当笔者在《生成扩散模型漫谈(十四)》系列中逐步深入ODE式扩散模型构建方法时,曾以为已触及问题本质。直到ICLR2023会议上《Flow Straight and Fast》一文的出现,其提出的Rectified Flow框架以近乎颠覆性的简洁思路,彻底重构了学界对扩散ODE的认知体系。这种如同"拨开云雾见青天"的突破性进展,促使我们重新审视ODE扩散模型的构建逻辑。

核心突破:从轨迹设计到导数学习

在传统认知中,扩散模型的本质是实现从噪声分布$p_T(\boldsymbol{x}_T)$到数据分布$p_0(\boldsymbol{x}_0)$的演化转换,而ODE式扩散模型则通过定义如下微分方程来刻画这一过程: \begin{equation}\frac{d\boldsymbol{x}_t}{dt}=\boldsymbol{f}_t(\boldsymbol{x}_t)\end{equation} 构建此类模型的核心挑战在于如何设计满足分布转换要求的向量场函数$\boldsymbol{f}_t(\boldsymbol{x}_t)$。Rectified Flow的创新之处在于,它完全摒弃了传统方法中复杂的概率推导过程,转而采用一种构造-逼近的双步策略。

具体而言,该方法首先随机采样数据样本$\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)\end{equation} 这里的轨迹函数只需满足基本约束$\boldsymbol{\varphi}_0(\boldsymbol{x}_0, \boldsymbol{x}_T)=\boldsymbol{x}_0$和$\boldsymbol{\varphi}_T(\boldsymbol{x}_0, \boldsymbol{x}_T)=\boldsymbol{x}_T$即可,理论上任何连续函数形式都是允许的。基于此轨迹,我们可以直接计算其时间导数: \begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \frac{\partial \boldsymbol{\varphi}_t(\boldsymbol{x}_0, \boldsymbol{x}_T)}{\partial t}\end{equation} 这个导数虽然精确描述了轨迹变化率,但由于包含未知的$\boldsymbol{x}0$,无法直接用于生成过程。Rectified Flow的关键洞见在于:用神经网络直接学习这个导数函数。通过优化如下目标函数: \begin{equation}\mathbb{E}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t) - \frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right\Vert^2\right]\end{equation} 使模型$\boldsymbol{v}{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$能够仅根据当前状态$\boldsymbol{x}_t$和时间$t$,精确逼近真实轨迹的导数信息。这种将概率问题转化为函数逼近问题的思路,大幅降低了模型构建的理论门槛。

极简实现:直线轨迹的启示

为验证这一思路的可行性,我们考察最简单的线性轨迹场景。设$T=1$并定义轨迹函数为: \begin{equation}\boldsymbol{x}_t = \boldsymbol{x}_0 + t(\boldsymbol{x}_1 - \boldsymbol{x}_0)\end{equation} 对时间$t$求导可得: \begin{equation}\frac{\partial \boldsymbol{\varphi}_t}{\partial t} = \boldsymbol{x}_1 - \boldsymbol{x}_0\end{equation} 此时的训练目标转化为: \begin{equation}\mathbb{E}\left[\left\Vert \boldsymbol{v}_{\boldsymbol{\theta}}\big(\boldsymbol{x}_0 + t(\boldsymbol{x}_1 - \boldsymbol{x}_0), t\big) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\right\Vert^2\right]\end{equation} 通过简单的变量替换$\boldsymbol{x}_t = \boldsymbol{x}_0 + t(\boldsymbol{x}_1 - \boldsymbol{x}_0)$,可将目标函数重写为: \begin{equation}\mathbb{E}\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} 这个结果令人震惊——它与笔者之前研究的"直线轨迹"模型完全一致,但推导过程却简化了至少60%。更重要的是,这种方法可以无缝扩展到任意轨迹函数,只需更换对应的导数表达式即可。

蓝色与白色相间的二维码图片,中间包含彩色字母元素 如上图所示,该示意图抽象展示了Rectified Flow中轨迹构造与导数学习的对应关系。左侧蓝色区域代表数据分布空间,右侧白色区域代表噪声分布空间,中间彩色字母元素象征着神经网络对轨迹导数的逼近过程。这一可视化生动体现了从确定性轨迹到随机生成的范式转换,为研究人员理解该方法的核心机制提供了直观参考。

理论验证:概率视角的严格证明

尽管Rectified Flow的实践思路极为简洁,但其理论有效性仍需严格证明。我们需要确认:通过优化上述目标函数得到的向量场$\boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{x}_t, t)$,是否确实能引导分布从$p_T(\boldsymbol{x}_T)$演化至$p_0(\boldsymbol{x}_0)$。

证明的关键在于验证该方法满足概率守恒的基本要求。考虑任意测试函数$\phi(\cdot)$,我们考察轨迹演化过程中的期望变化: \begin{equation}\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]\end{equation} 对右侧进行一阶泰勒展开可得: \begin{equation}\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]\end{equation} 通过交换期望顺序并利用条件概率分解$p(\boldsymbol{x}_0,\boldsymbol{x}_t)=p(\boldsymbol{x}_0|\boldsymbol{x}_t)p(\boldsymbol{x}t)$,可以得到: \begin{equation}\mathbb{E}{\boldsymbol{x}_t}\left[\phi(\boldsymbol{x}t) + \Delta t,\mathbb{E}{\boldsymbol{x}_0|\boldsymbol{x}t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\cdot\nabla{\boldsymbol{x}_t}\phi(\boldsymbol{x}_t)\right]\end{equation} 这表明演化过程满足: \begin{equation}\frac{d\boldsymbol{x}t}{dt} = \mathbb{E}{\boldsymbol{x}_0|\boldsymbol{x}_t}\left[\frac{\partial \boldsymbol{\varphi}_t}{\partial t}\right]\end{equation} 根据数学期望的性质,上述条件期望恰好是最小化平方损失的最优解。这就从理论上严格证明了:通过优化轨迹导数的平方损失,确实能够获得满足分布转换要求的ODE向量场。

范式转换:从概率建模到函数逼近

Rectified Flow带来的不仅是技术上的简化,更是一种思维方式的革新。回顾扩散模型的发展历程,从DDPM的随机微分方程到Score-Based Models的得分匹配,学界始终在概率框架内寻找解决方案。而Rectified Flow大胆跳出这一思维定式,将复杂的分布转换问题转化为直观的函数逼近问题,这种"降维打击"式的创新让人耳目一新。

笔者在前期研究中曾尝试过类似的轨迹构造思路,但受限于传统概率框架的思维定式,未能突破条件概率建模的技术瓶颈。Rectified Flow通过直接学习导数函数的方式,完美避开了复杂的概率推断过程,其简洁程度达到了"两步构建"的极致:1)设计任意满足边界条件的轨迹;2)训练神经网络逼近该轨迹的时间导数。这种极简主义的设计哲学,不仅大幅降低了模型实现难度,更为后续研究提供了无限可能。

值得注意的是,该方法与最优传输理论存在深刻联系。当选择Wasserstein距离意义下的最优轨迹时,Rectified Flow等价于求解分布间的最优传输映射。这种理论连接为进一步提升模型效率指明了方向,例如通过引入最优传输的计算加速技术,可以显著减少ODE求解的计算成本。

未来展望:简化之路的终点与起点

Rectified Flow的出现,标志着ODE式扩散模型构建方法达到了前所未有的简洁程度。很难想象在"构造轨迹-学习导数"这一基本框架之外,还能有更简化的构建思路。但科学发展的历史告诉我们,每个终点都是新的起点。该方法目前仍存在若干值得探索的方向:如何设计更高效的轨迹函数以加速收敛?怎样将该框架扩展到条件生成场景?这些问题的解决将进一步释放Rectified Flow的潜力。

从技术实现角度看,该方法已在多个基准数据集上验证了其优越性。在ImageNet64等标准任务中,基于Rectified Flow构建的模型不仅达到了SOTA生成质量,还将采样步数减少了一个数量级。这种"质量-效率"双优的特性,使其在实际应用中具有很强的竞争力。随着研究的深入,我们有理由相信,这种极简主义的建模思想将在更多生成模型领域产生深远影响。

分类:人工智能 标签:生成模型, 微分方程, 深度学习, 概率建模

【免费下载链接】diffusers-cd_imagenet64_l2 【免费下载链接】diffusers-cd_imagenet64_l2 项目地址: https://ai.gitcode.com/hf_mirrors/openai/diffusers-cd_imagenet64_l2

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐