How to Scale Your EMA

要約

バッチ サイズ全体でトレーニングのダイナミクスを維持することは、バッチ サイズと実時間の間のトレードオフを可能にするため、実際の機械学習にとって重要なツールです。
このトレードオフは通常、スケーリング ルールによって有効になります。たとえば、確率的勾配降下法では、バッチ サイズに応じて学習率を線形にスケーリングする必要があります。
実用的な機械学習のためのもう 1 つの重要なツールは、モデルの指数移動平均 (EMA) です。これは、勾配情報を受け取らず、代わりにある程度の勢いでターゲット モデルに従うモデルのコピーです。
このモデル EMA は、教師あり学習の堅牢性と一般化特性を改善し、擬似ラベル付けを安定させ、自己教師あり学習 (SSL) の学習信号を提供できます。
これまでの研究では、モデル EMA を最適化とは別に扱っていたため、バッチ サイズごとにトレーニング ダイナミクスが異なり、モデルのパフォーマンスが低下していました。
この研究では、モデル EMA が存在する場合の最適化のためのスケーリング ルールを提供し、さまざまなアーキテクチャ、オプティマイザー、およびデータ モダリティにわたるその有効性を実証します。
また、モデル EMA がターゲット モデルの最適化に寄与するルールの有効性も示し、小規模および大規模なバッチ サイズで EMA ベースの疑似ラベル付けおよび SSL メソッドをトレーニングできるようにします。
SSL の場合、パフォーマンスを犠牲にすることなくバッチ サイズ 24,576 までの BYOL のトレーニングが可能になり、実測時間で 6$\times$ の短縮が最適になります。

要約(オリジナル)

Preserving training dynamics across batch sizes is an important tool for practical machine learning as it enables the trade-off between batch size and wall-clock time. This trade-off is typically enabled by a scaling rule, for example, in stochastic gradient descent, one should scale the learning rate linearly with the batch size. Another important tool for practical machine learning is the model Exponential Moving Average (EMA), which is a model copy that does not receive gradient information, but instead follows its target model with some momentum. This model EMA can improve the robustness and generalization properties of supervised learning, stabilize pseudo-labeling, and provide a learning signal for Self-Supervised Learning (SSL). Prior works have treated the model EMA separately from optimization, leading to different training dynamics across batch sizes and lower model performance. In this work, we provide a scaling rule for optimization in the presence of model EMAs and demonstrate its validity across a range of architectures, optimizers, and data modalities. We also show the rule’s validity where the model EMA contributes to the optimization of the target model, enabling us to train EMA-based pseudo-labeling and SSL methods at small and large batch sizes. For SSL, we enable training of BYOL up to batch size 24,576 without sacrificing performance, optimally a 6$\times$ wall-clock time reduction.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google