STORK: Faster Diffusion And Flow Matching Sampling By Resolving Both Stiffness And Structure-Dependence

📄 arXiv: 2505.24210v2 📥 PDF

作者: Zheng Tan, Weizhen Wang, Andrea L. Bertozzi, Ernest K. Ryu

分类: cs.CV, math.NA

发布日期: 2025-05-30 (更新: 2025-10-01)

🔗 代码/项目: GITHUB

💡 一句话要点

STORK：通过解决刚性和结构依赖性加速扩散模型和Flow Matching模型的采样

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 扩散模型 Flow Matching 快速采样 数值积分 ODE求解

📋 核心要点

  1. 扩散模型和Flow Matching模型采样过程计算量大，快速采样方法是研究热点。
  2. STORK方法旨在解决ODE的刚性和对半线性结构的依赖性，从而加速采样过程。
  3. 实验证明，STORK方法能够持续提高图像和视频生成的扩散和Flow Matching采样质量。

📝 摘要（中文）

扩散模型（DMs）和Flow Matching模型在图像和视频生成方面表现出了卓越的性能。然而，这些模型在采样过程中需要大量的函数评估次数（NFEs），导致推理成本高昂。因此，开发需要较少NFEs且能保持质量的快速采样方法一直是研究的热点。然而，现有的无训练采样方法无法同时解决两个关键挑战：ODE的刚性（即速度场的非线性）以及对DM ODE的半线性结构的依赖性（这限制了它们直接应用于Flow Matching模型）。在这项工作中，我们引入了稳定化泰勒正交Runge-Kutta（STORK）方法，解决了这两个设计问题。我们证明了STORK能够持续提高图像和视频生成的扩散和Flow Matching采样的质量。

🔬 方法详解

问题定义：扩散模型和Flow Matching模型在图像和视频生成领域取得了显著成果，但采样过程需要大量的函数评估（NFEs），导致计算成本高昂。现有的快速采样方法通常难以同时处理ODE的刚性（速度场的非线性）以及模型对半线性结构的依赖性，限制了其在不同模型上的通用性。

核心思路：STORK方法的核心在于设计一种新的数值积分方案，该方案能够有效地处理ODE的刚性，同时避免对特定模型结构的过度依赖。通过稳定化的泰勒展开和正交Runge-Kutta方法，STORK能够在保证采样质量的前提下，显著减少所需的函数评估次数。

技术框架：STORK方法基于常微分方程（ODE）求解器，用于对扩散模型或Flow Matching模型定义的轨迹进行积分。其主要流程包括：1) 使用稳定化的泰勒展开来近似速度场；2) 利用正交Runge-Kutta方法进行数值积分；3) 在每一步迭代中，根据稳定性和精度要求调整步长。整个框架旨在以更少的步骤逼近目标分布。

关键创新：STORK方法的关键创新在于其稳定化的泰勒展开和正交Runge-Kutta方法的结合。稳定化的泰勒展开能够有效地抑制数值不稳定性，而正交Runge-Kutta方法则能够提高积分精度。与现有方法相比，STORK方法对模型结构的依赖性更小，因此可以更广泛地应用于不同的扩散模型和Flow Matching模型。

关键设计：STORK方法的关键设计包括：1) 稳定化项的选择，用于抑制泰勒展开的高阶项带来的数值不稳定性；2) 正交Runge-Kutta方法的阶数和系数的选择，以平衡精度和计算成本；3) 步长自适应策略，根据局部误差估计动态调整步长，以保证采样质量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，STORK方法在图像和视频生成任务中，能够以更少的函数评估次数（NFEs）达到与现有方法相当甚至更高的生成质量。具体而言，在相同的生成质量下，STORK方法可以显著减少所需的NFEs，从而降低计算成本，提高采样效率。代码已开源。

🎯 应用场景

STORK方法可广泛应用于图像生成、视频生成、音频生成等领域，尤其适用于对计算资源有限或对实时性要求较高的场景。通过减少采样所需的函数评估次数，STORK方法可以显著降低推理成本，提高生成效率，从而加速相关应用的部署和普及。未来，该方法有望进一步扩展到其他生成模型和逆问题求解领域。

📄 摘要（原文）

Diffusion models (DMs) and flow-matching models have demonstrated remarkable performance in image and video generation. However, such models require a significant number of function evaluations (NFEs) during sampling, leading to costly inference. Consequently, quality-preserving fast sampling methods that require fewer NFEs have been an active area of research. However, prior training-free sampling methods fail to simultaneously address two key challenges: the stiffness of the ODE (i.e., the non-straightness of the velocity field) and dependence on the semi-linear structure of the DM ODE (which limits their direct applicability to flow-matching models). In this work, we introduce the Stabilized Taylor Orthogonal Runge--Kutta (STORK) method, addressing both design concerns. We demonstrate that STORK consistently improves the quality of diffusion and flow-matching sampling for image and video generation. Code is available at https://github.com/ZT220501/STORK.