最优步长蒸馏（OSS）：通过优化采样步长来加速扩散模型的生成过程，同时保持生成质量

扩散模型是一种强大的生成模型，能够生成高质量的图像、视频等内容。然而，传统的扩散模型在采样过程中需要大量的步骤来逐步去除噪声并生成最终结果，这使得采样过程计算成本高昂。例如，在生成一张高质量的图像时，可能需要数百甚至上千个采样步骤，这在实际应用中会导致生成速度缓慢。

• GitHub：https://github.com/bebebe666/OptimalSteps

为了解决这一问题，中国科学院大学和腾讯混元研究院提出了一种新的方法——最优步长蒸馏（Optimal Stepsize Distillation, OSS）。该方法通过动态规划框架，从参考轨迹中提取最优的采样步长序列，从而在较少的采样步骤下生成与原始模型几乎相同的结果。

OSS通过优化采样步长来加速扩散模型的生成过程，同时保持生成质量。开发者已经提供了基于 DiT、FLUX、Open-Sora 和 Wan2.1 的算法的示例。

主要功能

1. 加速采样过程：通过优化采样步长，将原本需要数百步的采样过程加速到仅需数十步甚至更少，从而显著提高生成速度。
2. 保持生成质量：在减少采样步骤的同时，保持生成结果的质量，确保生成内容与原始模型的多步生成结果高度一致。
3. 通用性与鲁棒性：该方法适用于多种扩散模型架构、不同的噪声调度策略以及不同阶数的常微分方程（ODE）求解器，具有广泛的适用性。

主要特点

1. 理论最优性：通过动态规划算法，保证了在有限采样步骤下，生成结果与原始模型的连续轨迹尽可能接近，实现了全局误差最小化。
2. 架构无关性：该方法不依赖于具体的模型架构，无论是基于U-Net的模型还是基于Transformer的模型，都能有效应用。
3. 轻量级适配：在不同任务之间，仅需对采样步长进行轻量级校准，无需重新训练模型，即可实现5-10倍的速度提升，且性能损失极小。

工作原理

1. 知识蒸馏视角：将优化采样步长问题视为一种知识蒸馏问题，其中“学生”采样过程（有限步骤）试图逼近“教师”采样过程（大量步骤）的结果。
2. 动态规划算法：利用递归子结构特性，通过动态规划算法系统地推导出理论上的最优步长序列。具体来说，将问题分解为多个子任务，即用i个学生步骤逼近j个教师步骤的结果，并通过递归公式逐步求解。
3. 全局误差最小化：通过动态规划，确保在任何预定义的学生采样步骤下，生成结果与教师模型的输出尽可能接近，从而实现全局误差最小化。
4. 幅度校准：针对少步采样中可能出现的幅度偏差问题，提出了一种逐步仿射变换方法，通过调整学生输出的幅度，使其与教师模型的幅度特征对齐，从而提高生成结果的细节和真实性。

应用场景

1. 文本到图像生成：在文本到图像生成任务中，该方法可以将生成速度提升10倍，同时在GenEval基准测试中保持99.4%的性能，显著提高了生成效率。
2. 视频生成：在视频扩散模型中，优化后的采样步长能够实现10倍的速度提升，同时保持视频的视觉保真度，使得视频生成更加高效。
3. 多任务扩散模型：在多任务扩散模型中，最优步长调度可以被视为一种任务划分策略，更好地平衡不同噪声强度下的去噪任务，从而提高最终输出的质量。
4. 其他生成任务：该方法还可应用于其他需要高效采样的生成任务，如音频生成、3D模型生成等，具有广泛的应用前景。