[論文レビュー] Stochastic Weight Averaging in Parallel: Large-Batch Training that Generalizes Well

Vipul Gupta, Santiago Akle Serrano+1 more|arXiv (Cornell University)|Jan 7, 2020

Advanced Neural Network Applications参考文献 26被引用数 15

ひとこと要約

本稿では、大規模ミニバッチによる高速収束の後に、複数の独立して訓練された小規模ミニバッチモデルの重みを平均化することで、深層ニューラルネットワークの大バッチ学習を高速化する手法であるStochastic Weight Averaging in Parallel（SWAP）を提案する。SWAPは、小バッチ学習と同等の一般化性能を達成するが、著しく短い時間で実行され、CIFAR10では最先端の手法と比較して最大68%の訓練時間を短縮する。

ABSTRACT

We propose Stochastic Weight Averaging in Parallel (SWAP), an algorithm to accelerate DNN training. Our algorithm uses large mini-batches to compute an approximate solution quickly and then refines it by averaging the weights of multiple models computed independently and in parallel. The resulting models generalize equally well as those trained with small mini-batches but are produced in a substantially shorter time. We demonstrate the reduction in training time and the good generalization performance of the resulting models on the computer vision datasets CIFAR10, CIFAR100, and ImageNet.

研究の動機と目的

• 大バッチディープラーニング学習における訓練速度と一般化性能のトレードオフを解消すること。
• 大バッチ初期化の後に小バッチモデルの並列訓練を活用することで、モデルの一般化性能を損なわずに高速化を実現すること。
• 大規模DNN学習に計算リソースを効率的に活用できる、実用的でチューニングが少ない手法を開発すること。
• 大バッチ学習に続く並列的小バッチの微調整によって、小バッチ学習と同等の一般化性能を持つモデルが得られることを示すこと。

提案手法

• まず、大規模ミニバッチを用いて、訓練損失が低い領域に素早く到達するようにモデルを訓練する。
• 次に、大バッチモデルから開始し、複数の独立したモデルを小規模ミニバッチで並列に訓練する。
• 最終的なモデルは、これらの独立して訓練された小バッチモデルの重みを平均することで得られる。
• 小バッチ段階では、平坦な極小値への収束を促進するために周期的学習率スケジュールを用いる。
• 大バッチから小バッチへの移行は、ハイパーパrameterによって制御され、実験ではグリッドサーチにより選定される。
• 本手法はさまざまな最適化手法と互換性があり、通信オーバーヘッドを低減するための技術（例：ローカルSGD）と組み合わせて利用可能である。

実験結果

リサーチクエスチョン

• RQ1 一般化性能を損なわずに、大バッチ学習を高速化できるか？
• RQ2 複数の独立して訓練された小バッチモデルの重みを平均化することで、標準的大バッチ学習に比べて一般化性能が向上するか？
• RQ3 大バッチ事前学習と並列的小バッチ微調整の組み合わせによって、小バッチ学習と同等の一般化性能を達成できるか？
• RQ4 提案手法の訓練時間は、標準的小バッチおよび大バッチ学習戦略と比較してどの程度か？
• RQ5 性能と効率の最適化を図るための、大バッチと小バッチ段階の間の臨界的遷移点は存在するか？

主な発見

• CIFAR100では、SWAPは241.54秒でテスト精度79.11%を達成し、小バッチSWAと同等の性能を示したが、3.5倍の高速化を実現した。
• CIFAR10では、SWAPはDAWNBench優勝エントリの68%の時間で訓練を完了させながら、同等の一般化性能を達成した。
• 大バッチ学習に続く小バッチSWAは、SWAPと同等のテスト精度（CIFAR100で78.12% vs. 78.18%）を達成するまでに3倍以上時間がかかった。
• 40エポックの小バッチエポックを用いたSWAPは、CIFAR100で241.54秒で79.11%のテスト精度を達成し、大バッチSWA（76.00%）および大バッチに続く小バッチSWA（78.12%）よりも速度と精度の両面で優れていた。
• SWAPの最終的な平均化モデルは、標準的なSGDモデルよりも損失盆地の中心に近いことが判明し、より高いロバスト性を示唆している。
• 可視化結果から、大バッチモデルは、微調整された小バッチモデルが収束する領域と重複する盆地に到達していることが示され、ほぼ凸な盆地の存在を支持する。