[論文レビュー] 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

ひとこと要約

本論文では、各訓練エポックの終了時に収集されたk個の最新モデルチェックポイントの重みを平均化することで、訓練収束を加速する、モデルおよび最適化手法に依存しない単純な手法であるLatest Weight Averaging (LAWA)を提案する。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.

研究の動機と目的

• 大規模なビジョンおよび言語モデルの高い計算コストと長時間の訓練を低減すること、特に高価なハードウェアにアクセスできない研究者を対象とする。
• モデルアーキテクチャ、最適化手法、ハイパーパrameterの変更なしに、訓練収束速度を向上させること。
• 損失の地形がまだ変化している中間段階の訓練中に、最近のモデル重みを平均化することで収束を加速できるかどうかを検証すること。
• 深層学習研究における迅速な反復的実験を可能にする、実用的で即挿し可能な手法を提供すること。

提案手法

• LAWAは各エポックの終了時にモデル重みを収集し、キューを用いてk個の最新チェックポイントのスライディングウインドウを維持する。
• 各エポックにおいて、LAWAアンサンブルのモデルパラメータは、k個の最新チェックポイントの均等平均によって更新される：θ^LAWA_E = (1/k) * Σ_{i=E-k+1}^{E} θ_i。
• 平均化処理は円形バッファを用いて効率的に実装され、時間計算量が削減され、平均化モデルは必要に応じてのみ評価されるため、メモリおよび計算オーバーヘッドが最小限に抑えられる。
• 平均化モデルのバッチ正規化統計量は、訓練データを完全に前方伝搬するインフェンスパスにより再計算され、一貫性のあるバッチ統計量が保証される。
• この手法はモデルアーキテクチャ、最適化手法（例：SGD や Adam）およびタスクに対してアーキテクチャに依存せず、ビジョンおよびNLPのベンチマークに広く適用可能である。
• ハイパーパrameterであるk（最新チェックポイント数）は経験的に調整され、ImageNetおよびWikiText-103の両方の実験でk=6が優れた性能を示した。

実験結果

リサーチクエスチョン

• RQ1 訓練の中間段階において、k個の最新チェックポイントを平均化することで、損失および精度の観点から収束が顕著に加速するか？
• RQ2 LAWAは、SGD や Adam を用いた標準的な訓練に比べ、大規模なビジョンおよび言語モデルの「精度到達までの時間」を短縮できるか？
• RQ3 LAWAの性能はkの選択にどれほど敏感であり、異なる訓練段階において特定のk値で最適性能を示すか？
• RQ4 LAWAは、アーキテクチャや最適化手法の変更なしに、コンピュータビジョン（ResNet50/ImageNet）およびNLP（RoBERTa-Base/WikiText-103）の両方に対して効果的に適用可能か？
• RQ5 LAWAは最終的なモデルの汎化性能を向上させるのか、それとも単に収束加速のための手法にとどまるのか？

主な発見

• LAWAは、ResNet50/ImageNetにおいて、標準的なSGD訓練と比較して、特定の検証損失に到達するまでの時間を最大68 GPU時間短縮した。
• RoBERTa-Base/WikiText-103では、LAWAがAdamの最終検証損失に約45エポックも早く到達でき、約30 GPU時間の節約が達成された。
• LAWAで平均化されたモデルの最終的な汎化性能は、最終的なAdam最適化モデルよりも一貫して優れており、汎化性能の向上が確認された。
• LAWAはkが中程度の値（例：k=6）のときに最も効果的であり、kが16を超えると性能が低下した。これは、過去のチェックポイントを多すぎると平均化が有害である可能性を示唆している。
• この手法は、異なる最適化手法やタスクに対してロバストであり、ビジョンおよび言語モデリングベンチマークの両方で一貫した収束スピードアップが観察された。
• 経験的結果から、初期訓練段階での平均化は性能を悪化させることが確認され、LAWAが損失地形がまだ変化している中間段階で最も効果的であることが裏付けられた。