扩散模型训练动态的分析与优化

扩散模型目前在数据驱动的图像生成领域占据主导地位，其在大规模数据集上的可扩展性表现尤为突出。本文中，我们识别并修正了主流ADM扩散模型架构在训练过程中存在的若干导致训练不均衡且低效的问题，而无需改变其高层结构。通过观察训练过程中网络激活值与权重的幅值变化失控及不平衡现象，我们重新设计了网络层，使其在期望意义上保持激活值、权重及更新量的幅值稳定。我们发现，系统性地应用这一设计原则可有效消除观测到的漂移与不平衡现象，在相同计算复杂度下显著提升模型性能。通过采用快速确定性采样方法，我们的改进将ImageNet-512图像生成任务的先前最优FID分数2.41提升至1.81。作为独立贡献，我们提出了一种在训练完成后（即事后）设定指数移动平均（EMA）参数的方法。该方法无需多次重复训练即可实现对EMA长度的精确调优，并揭示了EMA参数与网络架构、训练时长及引导机制之间出人意料的相互作用关系。