Hotdry.

Article

扩散模型积分表示学习:Flow Map 原理与实践指南

深入解析扩散模型在连续时间域的积分表示方法,探讨 Flow Map 如何实现无需离散采样的生成与推理,涵盖一致性规则、训练策略与工程实践。

2026-05-06ai-systems

扩散模型作为当前最具影响力的生成式人工智能技术之一,其核心通过迭代去噪过程将噪声样本逐步转化为目标数据样本。然而,这种逐 step 的采样方式带来了显著的计算开销。近年来,研究者们开始探索一种更为直接的路径:学习扩散模型的积分表示。这种被称为 Flow Map 的方法,能够在单次前向传播中完成从噪声到数据的映射,从而大幅提升生成效率。本文将系统性地介绍 Flow Map 的理论基础、一致性训练策略以及实际应用。

从扩散模型到 Flow Map 的演进

理解 Flow Map 的前提在于重新审视扩散模型的本质。在标准扩散模型的采样过程中,模型在每个时间步预测当前输入对应的干净数据期望值,这一预测实际上指向了数据流形上的切线方向。通过沿着这一方向执行有限步长的更新,采样算法逐渐将噪声分布引向目标数据分布,在输入空间中勾勒出一条连接两者的路径。

关键洞察在于:给定路径上的任意一点,扩散模型仅能提供局部信息 —— 即该点处的切线方向。要获得目标数据,需要反复执行这种局部更新,整个过程犹如在黑暗中依赖即时反馈进行导航,被称为 “死 reckoning”(航位推算)。这种方法的固有局限在于:每个采样步骤都是一次近似,且误差会随着步骤数量累积。

Flow Map 则采取了一种更为全局的视角:与其逐步积分切线方向,不如直接学习路径上的任意两点之间的映射关系。形式化地,对于时间步骤 s(源)和 t(目标),Flow Map 定义为 F(x_s, s, t) = x_t,即从路径上任意源点直接跳转到任意目标点。这种能力使其能够绕过迭代式采样的固有局限,实现更高效的生成过程。

三种一致性规则:理论框架

Flow Map 训练的核心在于确保一致性,即 Flow Map 对路径的描述必须满足某些数学性质。Boffi、Albergo 和 Vanden-Eijnden 等研究者提出了三种相互等价的一致性规则,它们从不同角度约束了 Flow Map 的行为。

组合性(Compositionality)是最直观的一致性规则:分两步从 s 跳转到 u 等价于直接从 s 跳转到 u。数学表达为 F(F(x_s, s, t), t, u) = F(x_s, s, u)。这一性质允许我们将小时间间隔的预测 “拼接” 成大时间间隔的预测,是许多蒸馏方法的核心依据。

拉格朗日一致性(Lagrangian Consistency)从另一个角度刻画 Flow Map:固定源点 x_s,观察目标时间 t 连续变化时 Flow Map 输出如何演变。输出对时间 t 的导数恰好等于该点处的速度向量,即 d/dt F(x_s, s, t) = v(F(x_s, s, t), t)。这建立了 Flow Map 与标准扩散模型之间的桥梁。

欧拉一致性(Eulerian Consistency)则固定目标时间 t,考察源点变化时的影响。当源时间 s 发生无穷小变化时,输出应保持不变,因为目标位置与源点的具体取值无关。这一约束涉及 Jacobian-Vector Product(JVP)的计算,为训练提供了第三种等价的损失函数形式。

这三种规则在数学上彼此等价,任何一种都足以完整刻画 Flow Map 的行为。选择哪种规则主要取决于实现便利性和计算效率的权衡。

训练策略:从蒸馏到自训练

基于上述一致性规则,研究者们发展出多种训练 Flow Map 的策略。根据是否依赖预训练扩散模型,这些方法大致可分为两类:蒸馏训练从零训练

蒸馏训练(Distillation)首先训练一个标准的扩散模型作为 “教师”,然后通过一致性损失函数将教师模型的知识迁移到 Flow Map 学生模型。具体而言,以教师模型在目标时间步的预测作为监督信号,学生模型学习在源时间步直接预测该目标位置的 Flow Map。这种方法充分利用了预训练扩散模型已经学到的丰富知识,训练过程相对稳定。

代表性方法包括 Lagrangian Map Distillation(LMD)和 Eulerian Map Distillation(EMD)。LMD 通过拉格朗日一致性损失训练 Flow Map,计算 ∂F/∂t 与教师速度向量之间的均方误差;EMD 则采用欧拉一致性视角,通过 ∂F/∂s + ∇F · v = 0 构建损失。这两种方法各有权衡:LMD 需要计算高阶导数,而 EMD 需要计算空间 Jacobian。

自蒸馏(Self-Distillation)进一步放松了对教师模型的依赖,允许学生模型同时扮演教师角色,通过 “自己教自己” 的方式逐步提升能力。这种方法的核心是确保 Flow Map 在任意时间点的预测都保持一致:先执行一个小跳获得中间点,再用 Flow Map 从该中间点预测目标,两者应给出相同结果。代表性工作包括 Consistency ModelsMeanFlow,后者通过巧妙的参数化设计,无需显式教师即可实现从零训练。

Marginal-from-Conditional 学习是 MeanFlow 提出的关键技巧。其核心洞察在于:标准扩散训练的 MSE 损失本质上是在学习条件期望 E[x_0 | x_t],但通过巧妙的 stop-gradient 操作,可以将其转化为对边缘速度的学习。这一技巧使 Flow Map 能够像标准扩散模型一样从数据中直接学习,而无需依赖任何预训练教师。

工程实践:参数配置与优化技巧

将 Flow Map 应用于实际生产环境时,需要关注若干工程细节。首先是时间步采样策略:由于 Flow Map 需要处理源时间步和目标时间步两个维度,采样策略的设计直接影响训练效率和最终性能。实践中常采用课程学习(Curriculum Learning)策略,初期聚焦短时间间隔以建立基础能力,后期逐步扩展到长时间间隔以学习长程依赖。

Stop-gradient 操作在自蒸馏训练中至关重要。通过阻止梯度流经目标计算图中的某些部分,可以避免高阶导数计算带来的内存开销和数值不稳定问题。典型的做法是将教师预测或中间计算结果包裹在 stop-gradient 中,使梯度仅流经学生模型的可学习参数。

损失函数加权同样需要细致调优。由于不同时间间隔的预测难度差异显著,简单均匀加权可能导致某些区间被过度关注而其他区间训练不足。实践中常采用时间相关的加权方案,例如与信噪比(SNR)相关联的权重调整。

对于大规模模型如 SDXLFLUX 的蒸馏,研究者们还引入了教师引导(Teacher Guidance)技术:不仅使用教师模型的预测作为监督信号,还将其作为条件信息输入学生模型。这种条件化策略使学生能够更好地捕捉教师的输出分布特性。

应用场景与延伸方向

Flow Map 的核心价值在于将扩散模型的迭代式采样转化为单次前向传播,这一特性在以下场景中尤为突出:

实时生成是最直接的应用领域。通过将完整扩散模型蒸馏为 Flow Map,可以在消费级硬件上实现数十毫秒级别的图像生成。LCM-LoRALatent Consistency Models 已经展示了 1-4 步生成的能力,FID 指标接近原始模型的数百步采样结果。

可控生成与编辑同样受益于 Flow Map 的全局视角。由于 Flow Map 建立了噪声空间与数据空间之间的确定性映射,可以直接对噪声向量进行操作以实现属性编辑、风格迁移等任务,而无需逐 step 执行复杂的引导策略。

离散数据建模是近期活跃的研究方向。传统的离散扩散模型在蒸馏为少步生成器时会面临严重的独立假设问题。连续扩散方法配合 Flow Map 训练提供了一种优雅的解决方案,Flow Map Language ModelsDiscrete Flow Maps 等工作已经在语言建模任务上展现出潜力。

非欧几里得空间的推广也值得关注。流形上的扩散模型和 Flow Map 拓展已在蛋白质设计等科学应用中得到验证,Riemannian MeanFlow 证明了该方法在黎曼流形上的可行性。

结语:超越离散采样

Flow Map 代表了扩散模型蒸馏领域的范式转变:从试图改进离散采样算法,转向学习数据流形上的全局积分映射。尽管当前方法仍需依赖预训练扩散模型进行知识迁移,但随着 MeanFlow 等从零训练方法的发展,完全端到端的 Flow Map 训练正在变得可行。

从更高视角审视,Flow Map 与一致性模型、Rectified Flow 等技术共同构成了 “少步生成” 研究的核心线索。这些方法虽然具体实现各异,但都指向同一个目标:在保持生成质量的前提下,大幅降低扩散模型的计算成本。随着硬件算力的持续提升和训练技术的不断成熟,我们有理由相信,Flow Map 将在下一代生成式人工智能系统中扮演关键角色。

参考资料:本文核心内容基于 Sander Dieleman 的博客文章 "Learning the integral of a diffusion model" 以及 Boffi 等人发表的 Flow Map Matching 相关论文。

ai-systems