使用开箱即用的开源方案加速扩散模型
某机构研究院 构建了一个名为 某机构 FastGen 的全新开源库,统一了最先进的视频扩散蒸馏技术。
2026年1月27日
作者:Weili Nie,Julius Berner,Arash Vahdat
近期大规模扩散模型的进展在从图像合成到音频生成、3D 资产创建、分子设计等多个领域彻底改变了生成式人工智能。这些模型在各类条件生成任务中展示了前所未有的高质量、多样化输出能力。
尽管取得了这些成功,采样效率低下仍是一个根本性瓶颈。标准扩散模型需要几十到几百次迭代去噪步骤,导致高推理延迟和巨大计算成本。这限制了它们在交互式应用、边缘设备和大规模生产系统中的实际部署。
视频生成面临着一个尤为关键的挑战。诸如某机构 Cosmos 等开源模型,以及商业化的文本到视频系统,已展现出惊人的文本到视频能力。然而,由于时间维度的存在,视频扩散模型的计算需求要高几个数量级。生成单个视频可能需要数分钟到数小时,这使得实时视频生成、交互式编辑以及用于智能体训练的世界建模变得非常困难。
在不牺牲质量和多样性的前提下加速扩散采样,已成为一个关键的开放挑战,而视频生成则是其中最 demanding 且最具影响力的应用之一。
本文介绍某机构 FastGen,这是一个开源库,统一了最先进的扩散蒸馏技术,用于将多步扩散模型加速为单步或少步生成器。文章回顾了基于轨迹和基于分布的蒸馏方法,展示了可复现的基准测试结果,在保持质量的同时实现了 10 到 100 倍的采样加速,并展示了 FastGen 对多达 140 亿参数的大型视频模型的可扩展性。文章还重点介绍了其在交互式世界建模中的应用,其中因果蒸馏能够实现实时视频生成。
主要的加速方法有哪些?
越来越多的研究探索了扩散蒸馏,旨在将长去噪轨迹压缩为少量推理步骤。现有方法大致分为两类:
- 基于轨迹的蒸馏——包括渐进蒸馏和一致性模型(如某机构的 iCT 和 sCM,以及某机构和某机构的 MeanFlow)——直接回归教师模型的去噪轨迹。
- 基于分布的蒸馏——如某机构的 LADD,以及某机构和某机构的 DMD——使用对抗性或变分目标来对齐学生模型和教师模型的分布。
这些方法已成功地将图像领域的扩散采样减少到一或两步。然而,每类方法都存在显著的权衡。基于轨迹的方法通常存在训练不稳定、收敛缓慢和扩展性挑战,而基于分布的方法则倾向于内存密集型、对初始化敏感且容易出现模式坍塌。此外,这些方法单独使用,都无法持续地为复杂数据(如真实世界视频)实现高保真度的单步生成。
这促使需要一个统一且可扩展的框架,以集成、比较和发展扩散蒸馏方法,实现稳定训练、高质量生成,并能扩展到大型模型和复杂数据。
FastGen 提供的功能
FastGen 是一个全新的、开源的、多功能的库,在通用的即插即用接口下汇集了最先进的扩散蒸馏方法。
统一且灵活的接口
FastGen 为跨不同任务的扩散模型加速提供了统一抽象。用户提供他们的扩散模型(以及可选的训练数据),并选择合适的蒸馏方法。然后 FastGen 处理训练和推理流程,以最小的工程开销将原始模型转换为单步或少步生成器。
图 1. FastGen 蒸馏流程
可复现的基准测试与公平比较
FastGen 在标准图像生成基准上复现了所有支持的蒸馏方法。历史上,扩散蒸馏方法是在不同的代码库中使用不同的训练方案提出和评估的,这使得公平比较变得困难。通过统一实现和超参数选择,FastGen 实现了透明的基准测试,并作为少步扩散社区的通用评估平台。
下表展示了在 CIFAR-10 和 ImageNet-64 基准上各类蒸馏方法性能的全面比较,证明了 FastGen 的可复现性。该表显示了 FastGen 统一实现实现的单步图像生成质量,以及各原始论文中报告的结果(括号内)。每种方法按其蒸馏方法分类:沿扩散轨迹优化的基于轨迹的方法(ECT、TCM、sCT、sCD、MeanFlow)和直接匹配生成分布的基于分布的方法(LADD、DMD2、f-distill)。
表 1. 单步图像生成质量(Fréchet 起始距离 FID 衡量)
| 加速方法 | CIFAR-10 | ImageNet-64 |
|---|---|---|
| 基于轨迹的蒸馏 | ||
| ECT | FastGen: 2.92 (论文: 3.60) | FastGen: 4.05 (论文: 4.05) |
| TCM | 2.70 (2.46) | 2.23 (2.20) |
| sCT | 3.23 (2.85) | – |
| sCD | 3.23 (3.66) | – |
| MeanFlow | 2.82 (2.92) | – |
| 基于分布的蒸馏 | ||
| LADD | – | – |
| DMD2 | 1.99 (2.13)* | 1.12 (1.28) |
| f-distill | 1.85 (1.92)* | 1.11 (1.16) |
超越视觉任务
虽然在本文中我们以视觉任务演示 FastGen,但该库具有足够的通用性,可以加速不同领域的任何扩散模型。一个特别值得关注的领域是人工智能 for 科学应用,其中样本质量通常与样本多样性同等重要。
通过将蒸馏方法与网络定义解耦,FastGen 使得添加新模型变得简单且即插即用。例如,我们已成功使用 ECT 在某个 PhysicsNeMo 平台中对某机构的天气降尺度模型 Corrector Diffusion (CorrDiff) 进行了蒸馏,实现了单步公里尺度的大气降尺度。
如下图 2 所示,蒸馏模型匹配了 CorrDiff 的预测(在技能和离散度方面),同时推理速度提升了 23 倍。
图 2. 台风 “灿都” 期间 10 米处东风观测(左上),蒸馏单步模型的四个预测(右上),以及 18 步 CorrDiff 模型(右下)
可扩展且高效的基础设施
FastGen 还提供了一个高度优化的训练基础设施,用于将扩散蒸馏扩展到大型模型。支持的技术包括:
- 全分片数据并行 v2 (FSDP2)
- 自动混合精度 (AMP)
- 上下文并行 (CP)
- 灵活注意力
- 高效 KV 缓存管理
- 自适应有限差分 JVP 估计
借助这些优化,FastGen 可以高效地蒸馏大规模模型。例如,我们使用 DMD2 成功地将一个 140 亿参数的 Wan2.1 文本到视频模型蒸馏成一个少步生成器,在 64 个某机构 H100 GPU 上耗时 16 小时收敛。
图 3 显示了使用改进的 DMD2 方法蒸馏 Wan2.1-T2V-14B 时,50 步教师模型与两步蒸馏学生模型的视觉比较。尽管学生模型在采样速度上比教师模型快 50 倍,但其生成质量与教师模型非常接近。
图 3. 使用改进的 DMD2 蒸馏 Wan2.1-T2V-14B:50 步教师模型(CFG=6, NFE=100)(左)与两步蒸馏学生模型(NFE=2)(右)的视觉比较。NFE 表示生成过程中的函数评估次数。
面向交互式世界建模的 FastGen
交互式世界模型旨在模拟环境动态,并实时连贯地响应用户动作或智能体干预。它们需要:
- 高采样效率
- 长时程时间一致性
- 动作条件的可控性
视频扩散模型因其捕捉丰富视觉动态的能力而为世界建模提供了坚实基础,但其多步采样过程和被动公式阻碍了实时交互。
为了解决这个问题,最近的工作探索了因果蒸馏,它将双向视频扩散模型转换为少步、块状自回归模型。这种自回归结构实现了实时交互,并已成为交互式世界模型的一个有前景的基础。
FastGen 为多种因果蒸馏方法实现了训练和推理方案,包括 CausVid 和 Self-Forcing,其中默认方案主要是基于分布的。
基于轨迹的蒸馏尚未在因果蒸馏中得到广泛应用,原因在于双向教师模型与自回归学生模型之间的性能下降和轨迹错位。FastGen 通过以下两种方式解决了这些挑战:
- 热启动因果蒸馏:在应用基于分布的目标之前,可以使用基于轨迹的方法来初始化学生模型。
- 通过扩散强制的因果有监督微调 (SFT):FastGen 提供了一个因果有监督微调方案,首先训练一个多步块状自回归模型,然后将其作为新的教师模型进行基于轨迹的蒸馏。
这些组件实现了混合蒸馏流程,结合了基于轨迹方法的稳定性和基于分布目标的灵活性。
在应用方面,FastGen 支持广泛的开源视频扩散模型,包括 Wan2.1、Wan2.2 和某机构 Cosmos-Predict2.5,并为多种视频合成场景提供端到端加速:
- 文本到视频 (T2V)
- 图像到视频 (I2V)
- 视频到视频 (V2V)
用户可以灵活定制因果蒸馏流程,例如,扩展到从 20 亿到 140 亿参数的模型,为 I2V 添加首帧条件,或为 V2V 任务融入结构先验(如深度引导的驱动视频)。
因此,FastGen 为推进交互式世界模型提供了必要的基础设施——实现了快速、可控且时间一致的生成,这对于将扩散模型从被动合成器转变为实时交互系统至关重要。
开始使用
FastGen 旨在不仅仅是一个蒸馏技术的集合——它是一个用于加速扩散模型的统一研究和工程平台。通过在一个可扩展且可复现的框架下整合基于轨迹和基于分布的方法,FastGen 降低了尝试少步扩散模型的门槛,并实现了跨方法的公平基准测试。
立即尝试 FastGen——插入您自己的扩散模型,选择一种蒸馏方法,见证一个多步生成器转变为单步执行者。无论您的目标是加速视觉合成、推动科学发现,还是赋能交互式世界模型,FastGen 都能提供所需的灵活性和可复现性,助您以创纪录的速度将想法转化为实现。FINISHED