大型扩散模型需要的不只是 EMA

arXiv:2610.22945v1 公告类型：新 摘要：尽管学习率衰减和 Schedule-Free 优化在大型语言模型训练中表现出色，大规模扩散模型通常仍采用恒定学习率（LR）和指数移动平均（EMA）进行训练。我们研究了 EMA 是否特别适合扩散模型训练，还是更强的替代方案尚未得到充分探索。我们的理论表明，对恒定步长轨迹进行平均存在固有的优化局限，而在标准非凸假设下，LR 衰减和 Schedule-Free 优化则能够提供趋于零的平稳性保证。实证结果显示，LR 衰减与 EMA 具有互补优势：在大多数测试的起始点上，将衰减与 EMA 结合都能进一步改善标准基线，但收益仍对衰减开始的时机较为敏感。Schedule-Free 则无需这一时机要求，并且在训练的各个阶段始终优于采用恒定 LR 的 EMA，同时实现更低的固定评估目标值和更高的生成质量。我们的分析还表明，这三种策略都能大幅降低参数变异性，但更强的稳定化并不总是更好：恒定 LR 下的 EMA 保持了相对更高的参数变异性，而激进的衰减与 EMA 结合虽然能够实现最低的检查点变异性，却也削弱了持续优化能力。Schedule-Free 在大幅降低参数变异性的同时保留了更多后期优化进展，在参数稳定性与持续优化之间实现了更好的平衡。最后，Schedule-Free 的优势不仅适用于预训练，也延伸到了 ReFL 和生成器侧 DMD，展现出其在不同扩散优化场景中的可迁移性。总体而言，我们的结果表明，大型扩散模型需要的不仅是参数平均：有效的优化需要在降低方差与持续优化之间取得平衡。