EMAをスケーリングする方法

要旨

バッチサイズ間でのトレーニングダイナミクスを維持することは、実用的な機械学習において重要なツールです。なぜなら、バッチサイズと実時間のトレードオフを可能にするからです。このトレードオフは通常、スケーリングルールによって実現されます。例えば、確率的勾配降下法（SGD）では、学習率をバッチサイズに比例して線形にスケーリングすべきです。実用的な機械学習におけるもう一つの重要なツールは、モデルの指数移動平均（EMA）です。これは勾配情報を受け取らないモデルのコピーであり、代わりに一定のモーメンタムでターゲットモデルに追従します。このモデルEMAは、教師あり学習のロバスト性と汎化性能を向上させ、擬似ラベリングを安定化し、自己教師あり学習（SSL）の学習信号を提供することができます。これまでの研究では、モデルEMAを最適化とは別に扱ってきたため、バッチサイズ間で異なるトレーニングダイナミクスが生じ、モデルの性能が低下していました。本研究では、モデルEMAが存在する場合の最適化のためのスケーリングルールを提供し、その有効性をさまざまなアーキテクチャ、オプティマイザ、データモダリティにわたって実証します。また、モデルEMAがターゲットモデルの最適化に寄与する場合にもこのルールが有効であることを示し、小規模および大規模なバッチサイズでEMAベースの擬似ラベリングおよびSSL手法をトレーニングできるようにします。SSLに関しては、BYOLのトレーニングをバッチサイズ24,576まで性能を犠牲にすることなく可能にし、最適な場合では実時間を6倍短縮します。

English

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 6times wall-clock time reduction.

EMAをスケーリングする方法

How to Scale Your EMA

要旨

Support