[论文解读] Efficient Diffusion Training via Min-SNR Weighting Strategy

Tiankai Hang, Shuyang Gu+6 more|arXiv (Cornell University)|Mar 16, 2023

Domain Adaptation and Few-Shot Learning被引用 5

一句话总结

本文提出 Min-SNR-γ，一种简单而有效的损失加权策略，通过基于截断信噪比（SNR）分配时间步特定权重，将去噪扩散训练视为多任务学习问题，显著加速收敛，并在 256×256 的 ImageNet 上实现 2.06 的新 SOTA FID 得分，且模型参数更小。

ABSTRACT

Denoising diffusion models have been a mainstream approach for image generation, however, training these models often suffers from slow convergence. In this paper, we discovered that the slow convergence is partly due to conflicting optimization directions between timesteps. To address this issue, we treat the diffusion training as a multi-task learning problem, and introduce a simple yet effective approach referred to as Min-SNR-$γ$. This method adapts loss weights of timesteps based on clamped signal-to-noise ratios, which effectively balances the conflicts among timesteps. Our results demonstrate a significant improvement in converging speed, 3.4$ imes$ faster than previous weighting strategies. It is also more effective, achieving a new record FID score of 2.06 on the ImageNet $256 imes256$ benchmark using smaller architectures than that employed in previous state-of-the-art. The code is available at https://github.com/TiankaiHang/Min-SNR-Diffusion-Training.

研究动机与目标

• 为解决去噪扩散模型训练中因时间步间优化方向冲突而导致的收敛缓慢问题。
• 将扩散训练重新构架为多任务学习问题，其中每个时间步被视为一个独立任务。
• 开发一种稳定、高效且可扩展的损失加权策略，避免自适应帕累托优化带来的不稳定性与计算开销。
• 在不修改网络架构的前提下提升训练效率与生成质量，实现在 ImageNet 基准上的最先进性能。

提出的方法

• 在多任务学习框架中，将每个去噪时间步视为一个独立任务。
• 提出 Min-SNR-γ，一种预定义的全局损失加权策略，使用截断信噪比（SNR）作为加权因子。
• 对 SNR 应用截断函数，防止极端权重，从而稳定训练并减少梯度冲突。
• 使用超参数 γ 控制截断强度，使方法在不同训练设置下均具备鲁棒性。
• 通过使用固定、分段的损失权重调度，避免运行时自适应加权，提升效率与稳定性。
• 使用标准架构（UNet、ViT）配合新损失加权策略进行训练，无需修改网络结构。

实验结果

研究问题

• RQ1 尽管时间步间共享模型权重，为何扩散训练仍收敛缓慢？
• RQ2 能否通过一种合理的损失加权策略缓解时间步间的优化方向冲突？
• RQ3 与自适应多任务学习方法（如帕累托优化）相比，全局预定义损失加权策略在扩散训练中的表现如何？
• RQ4 基于 SNR 的简单加权方案能否在收敛速度与 FID 得分上超越现有基线？
• RQ5 Min-SNR-γ 在不同架构及生成目标（如 ε 与 x₀ 预测）之间是否具备泛化能力？

主要发现

• Min-SNR-γ 的收敛速度比以往的加权策略快 3.4 倍，显著加速训练过程。
• 该方法在 ImageNet 256×256 基准上实现了 2.06 的新 SOTA FID 得分，优于以往的 SOTA 方法。
• 即使使用更小的模型（如 395M 参数），该方法在 ImageNet 256×256 上仍取得 2.81 的 FID 得分，优于更大的 DiT-XL-2 模型。
• 使用 Min-SNR-5 训练的 ViT-XL 模型仅用 210 万次迭代即达到 FID 2.08，比 DiT 快 3.3 倍。
• 该方法在 UNet 和 ViT 主干网络上均表现优异，且对 ε 与 x₀ 预测目标均有效。
• 消融实验表明，该方法对超参数 γ 具备鲁棒性，在不同设置下性能下降极小。