英伟达联合团队提出新型连续时间流图（flow map）模型 AYF：统一扩散与流模型的少步生成方案

由英伟达、多伦多大学及矢量研究所联合提出一种新型的连续时间流图（flow map）模型Align Your Flow（AYF），显著提升扩散模型和基于流的生成模型的采样效率。这些模型虽然在图像与文本生成任务中表现出色，但往往需要数百甚至上千次采样步骤，限制了其实际应用速度。

• 项目主页：https://research.nvidia.com/labs/toronto-ai/AlignYourFlow

AYF 通过将复杂模型蒸馏为高效的少步生成器，在仅需数步的情况下仍能保持高质量输出，从而大幅提升生成效率。

核心亮点概述 💡

• 高效少步生成：可在 4 步内完成高质量图像生成，显著快于传统方法。
• 统一理论框架：提出了新的训练目标，统一并拓展了现有的一致性模型和流匹配方法。
• 自引导与对抗微调：创新性地结合两种技术，增强生成质量的同时保持样本多样性。
• 多模态支持：不仅适用于图像生成，还可扩展至文本到图像等多模态任务。

主要功能一览

1. 高效少步生成

AYF 能够将原本需要大量采样步骤的扩散模型或流模型蒸馏为少步生成器，例如从 1000 步减少到 4 步，大幅缩短生成时间。

2. 高质量输出

即便在极少步骤下，AYF 的生成结果仍可媲美多步生成器，如在 ImageNet 上达到接近最佳 FID 分数。

3. 多模态能力

除了图像生成，AYF 还可用于文本到图像任务，根据文本描述快速生成高质量图像。

4. 自引导机制

AYF 利用低质量模型作为引导器，提升生成质量，避免传统分类器自由引导（CFG）带来的副作用。

技术特性解析

1. 连续时间流图（Flow Maps）

AYF 引入了连续时间流图的概念，能够直接连接任意两个噪声水平，适用于不同步数下的生成任务。

2. 新的训练目标

AYF 提出了两种新的训练目标：

• AYF-EMD（欧拉损失）：确保在固定时间 $ s $ 下，输出沿概率流 ODE 移动时保持一致。
• AYF-LMD（拉格朗日损失）：确保轨迹与 PF-ODE 对齐，提高模型稳定性。

3. 自引导技术

不同于传统 CFG，AYF 使用一个轻量级引导模型辅助训练，在不牺牲多样性的前提下提升生成质量。

4. 对抗性微调

通过引入对抗性微调策略，进一步提升生成效果，同时最小化对样本多样性的负面影响。

工作原理简析

AYF 的核心技术围绕“流图”这一神经网络结构展开：

1. 流图建模
   流图是一种映射函数，能够将任意噪声输入直接转换为目标数据分布中的点。AYF 训练该模型以学习从噪声到数据的路径。
2. 训练目标设计
   • AYF-EMD 和 AYF-LMD 目标分别从欧拉和拉格朗日视角出发，确保流图在不同时间步之间保持一致性。
   • 这些目标帮助模型更好地拟合真实的数据流动路径。
3. 自引导与对抗性微调
   • 在训练过程中使用低质量模型作为引导器，提升生成质量。
   • 后期加入对抗性微调，进一步优化生成效果，同时维持样本多样性。

实验验证成果

1. 图像生成性能

• 在 ImageNet 64x64 和 512x512 数据集上，AYF 均实现了当前最先进的少步生成性能。
• 例如，在 ImageNet 512x512 上，AYF 在 4 步生成时 FID 分数为 1.70，接近多步生成的最佳表现。

2. 文本到图像生成

• AYF 在文本条件合成任务中表现优异，优于所有非对抗性训练的少步采样器。
• 用户研究表明，AYF 生成的图像在质量与文本匹配度方面更具优势。

未来发展方向

AYF 的提出为高效生成模型开辟了新思路，未来可能的研究方向包括：

1. 跨领域扩展
   探索 AYF 在视频生成、分子结构预测等领域的应用潜力。
2. 性能优化
   通过更大容量的模型架构和更精细的训练策略，进一步提升生成质量。
3. 计算效率提升
   研究更高效的训练与推理方法，降低计算资源消耗，推动部署落地。