[论文解读] Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging

Jean Kaddour|arXiv (Cornell University)|Sep 29, 2022

Topic Modeling被引用 4

一句话总结

本文提出最新权重平均（LAWA），一种简单、与模型和优化器无关的方法，通过在每个训练周期结束时收集的最新k个模型检查点的权重进行平均，从而加速训练收敛。LAWA在ImageNet上最多节省68 GPU小时，在WikiText-103上最多节省30 GPU小时，使模型能提前数十个周期达到目标损失和准确率水平，且仅需在检查点收集和平均操作上对训练循环进行微小修改。

ABSTRACT

Training vision or language models on large datasets can take days, if not weeks. We show that averaging the weights of the k latest checkpoints, each collected at the end of an epoch, can speed up the training progression in terms of loss and accuracy by dozens of epochs, corresponding to time savings up to ~68 and ~30 GPU hours when training a ResNet50 on ImageNet and RoBERTa-Base model on WikiText-103, respectively. We also provide the code and model checkpoint trajectory to reproduce the results and facilitate research on reusing historical weights for faster convergence.

研究动机与目标

• 为解决大规模视觉和语言模型训练的高计算成本和长时间训练问题，特别是为那些无法使用高端硬件的研究人员提供支持。
• 在不修改模型架构、优化器或超参数的前提下，提升训练收敛速度。
• 探究在训练中段（损失景观仍在演化时）对近期模型权重进行平均，是否能加速收敛。
• 提供一种实用、即插即用的方法，以加速深度学习研究中的迭代实验。

提出的方法

• LAWA在每个训练周期结束时收集模型权重，并在队列中维护一个包含k个最近检查点的滑动窗口。
• 在每个周期，LAWA集成模型的参数通过计算k个最新检查点的均匀平均值进行更新：θ^LAWA_E = (1/k) * Σ_{i=E-k+1}^{E} θ_i。
• 通过使用循环缓冲高效实现平均操作，以降低时间复杂度，且仅在需要时评估平均模型，从而最小化内存和计算开销。
• 通过在训练数据上执行完整推理过程，重新计算平均模型的批量归一化统计量，以确保批量统计的一致性。
• 该方法与模型架构、优化器（如SGD或Adam）和任务无关，因此可广泛应用于视觉和NLP基准测试。
• 超参数k（最新检查点数量）通过经验调优，k=6在ImageNet和WikiText-103实验中均表现出色。

实验结果

研究问题

• RQ1 在训练中段对k个最新检查点进行平均，是否能显著加速损失和准确率的收敛？
• RQ2 与使用SGD或Adam的标准训练相比，LAWA在大规模视觉和语言模型上的达到准确率所需时间是否更短？
• RQ3 LAWA的性能对k的选择是否敏感？其性能是否在特定k值下达到最优，且在不同训练阶段表现一致？
• RQ4 LAWA是否无需架构或优化器修改，即可有效应用于计算机视觉（ResNet50在ImageNet上）和自然语言处理（RoBERTa-Base在WikiText-103上）？
• RQ5 LAWA是否提升了最终模型的泛化性能，还是仅作为收敛加速技术？

主要发现

• 与标准SGD训练相比，LAWA将ResNet50/ImageNet达到指定验证损失的时间最多减少了68 GPU小时。
• 在RoBERTa-Base/WikiText-103上，LAWA使模型比Adam更早约45个周期达到最终验证损失，节省约30 GPU小时。
• LAWA平均模型的最终泛化性能始终优于最终的Adam优化模型，证实其具有更好的泛化能力。
• 当k适中（如k=6）时，LAWA效果最佳，当k超过16时性能下降，表明过多历史检查点的平均可能产生负面影响。
• 该方法在不同优化器和任务下均表现稳健，在视觉和语言建模基准测试中均观察到一致的收敛加速。
• 实证结果表明，早期训练阶段进行平均会降低性能，证实LAWA在损失景观仍在演化的训练中段最为有效。