通过自适应采样高梯度方差时间步,加速扩散模型训练并提升收敛性能。
Adaptive Non-uniform Timestep Sampling for Accelerating Diffusion Model Training
- 根据梯度更新对目标函数的影响,动态选择关键时间步进行采样。
- 训练速度提升显著,且收敛后性能优于传统均匀采样方法。
- 适用于多种数据集与模型架构,鲁棒性强,适合高效训练场景。
作为高度表达性的生成模型,扩散模型在图像生成、自然语言处理和组合优化等多个领域表现出色。然而,随着数据分布日益复杂,模型训练至收敛所需计算资源越来越多。传统扩散模型通常采用均匀时间步采样,但我们的研究发现,不同时间步的随机梯度方差差异显著,高方差时间步成为阻碍快速收敛的瓶颈。为此,我们提出一种非均匀时间步采样方法,优先采样对目标函数影响更大的关键时间步。该方法通过追踪每个时间步梯度更新对目标函数的影响,自适应选择最可能有效降低目标值的时间步。实验表明,该方法不仅显著加速训练过程,还能在收敛时取得更优性能。此外,该方法在多种数据集、调度策略及扩散模型架构上均表现稳健,优于以往缺乏此类鲁棒性的采样与加权启发式方法。
原文摘要 · Abstract (English)
As a highly expressive generative model, diffusion models have demonstrated exceptional success across various domains, including image generation, natural language processing, and combinatorial optimization. However, as data distributions grow more complex, training these models to convergence becomes increasingly computationally intensive. While diffusion models are typically trained using uniform timestep sampling, our research shows that the variance in stochastic gradients varies significantly across timesteps, with high-variance timesteps becoming bottlenecks that hinder faster convergence. To address this issue, we introduce a non-uniform timestep sampling method that prioritizes these more critical timesteps. Our method tracks the impact of gradient updates on the objective for each timestep, adaptively selecting those most likely to minimize the objective effectively. Experimental results demonstrate that this approach not only accelerates the training process, but also leads to improved performance at convergence. Furthermore, our method shows robust performance across various datasets, scheduling strategies, and diffusion architectures, outperforming previously proposed timestep sampling and weighting heuristics that lack this degree of robustness.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。