ニューラルネットワークの並列訓練の限界を測定(2/2)

１．ニューラルネットワークの並列訓練の限界を測定(2/2)まとめ

・有効な最大バッチサイズはある程度ワークロードのあらゆる側面に依存
・しかし大きいバッチサイズの恩恵を受ける事ができる条件はまだ不明
・最適化アルゴリズムのわずかな変更が訓練時間を劇的に短縮する可能性有

２．巨大なバッチサイズが有効になる条件

以下、ai.googleblog.comより「Measuring the Limits of Data Parallel Training for Neural Networks」の意訳です。元記事は2019年3月19日、 Chris ShallueさんとGeorge Dahlさんによる投稿です。これも地味ですがとても今後が楽しみな研究です。

バッチサイズとトレーニング時間の基本的な関係は普遍的であるように思われますが、異なるスケーリング段階間を遷移するポイントはニューラルネットワークアーキテクチャとデータセットによって劇的に異なることがわかりました。これは、単純なデータ並列処理が、いくつかのワークロードに対して、現在のハードウェア（例えば、Cloud TPU Pods）の限界を超える大きなスピードアップを提供できることを意味します。

現在の最大のハードウェアや、近い将来開発されるであろうハードウェアから恩恵を受けるためには、ワークロードによっては単純なデータ並列処理から移行する必要があるものもあります。たとえば、前述の図では、CIFAR-10のResNet-8は1,024を超えるバッチサイズの恩恵を受けることができませんが、ImageNetのResNet-50ではバッチサイズを少なくとも65,536まで増やすことで恩恵を受けることができます。

作業手順(ワークロード)の最適化
どのワークロードがデータ並列トレーニングから最も利益を得るかを予測できれば、利用可能なハードウェアを最大限に活用するようにワークロードを調整できます。

ただし、有効な最大バッチサイズは、少なくともある程度、ワークロードのあらゆる側面（ニューラルネットワークアーキテクチャ、データセット、およびオプティマイザ）に依存するため、これは簡単ではないことが多くの結果から示唆されています。

例えば、同じオプティマイザ(最適化手法)と同じデータセットを使用してトレーニングした場合を比較しても、ニューラルネットワークアーキテクチャによっては、他よりはるかに大きいバッチサイズの恩恵を受けることができる場合があります。

この効果は時々ネットワークの幅と深さに依存しますが、しかし、異なるタイプのネットワーク間で矛盾するケースも存在し、更に幾つかのネットワークは「幅」と「深さ」の明白な概念さえ持っていません。

また、一部のデータセットは他のものよりはるかに大きいバッチサイズの恩恵を受けることができますが、これらの違いは常にデータセットのサイズによって説明されるわけではありません。小さいデータセットの方が大きいデータセットよりも大きいバッチサイズの恩恵を受ける場合があります。

左：LM1Bデータセットではトランスフォーマは、LSTMよりもはるかに大きなバッチサイズの恩恵を受ける事ができます。右：Common Crawlデータセットは、LM1Bデータセットの1000倍以上のサイズですが、LM1Bデータセットより大きなバッチサイズの恩恵を受ける事ができません。

おそらく私たちの最も有望な発見は、確率的勾配降下法で運動量(momentum)を取り入れる事のような最適化アルゴリズムへのわずかな変更でさえ、バッチサイズの増加と訓練時間の短縮の関係を劇的に改善できることです。これにより、新しいデータオプティマイザを設計したり、今まで考慮していなかったオプティマイザのスケーリング特性をテストしたりして、データの並列処理を最大限に活用できるオプティマイザを見つけることができます。

今後の研究
バッチサイズを増やすことによって追加のデータ並列処理を利用することは、さまざまなワークロードにわたって貴重なスピードアップを実現するための簡単な方法ですが、試したすべてのワークロードについて、最先端のハードウェアの限界内でメリットは減少しました。

しかしながら、今回の研究結果は、いくつかの最適化アルゴリズムが多くのモデルとデータセットにまたがって完全なスケーリング段階に一貫して拡張することができるかもしれないことを示唆しています。

今後の作業では、他のオプティマイザを使って同じ測定を実行し、それらの既存のオプティマイザが様々なケースで完全なスケーリング段階を拡張できるかどうかを調べます。

謝辞
この研究の著者は、Chris Shallue、Jaehoon Lee、Joe Antognini、Jascha Sohl-Dickstein、Roy Frostig、およびGeorge Dahl（ChrisとJaehoonは同等の貢献）です。多くの研究者が私たちが基盤としてきたこの分野で研究をしてきたので、関連する研究の完全な詳細については私たちの論文を見てください。