How to Scale Your EMA

要約

バッチ サイズ全体でトレーニングのダイナミクスを維持することは、バッチ サイズと実時間の間のトレードオフを可能にするため、実際の機械学習にとって重要なツールです。
このトレードオフは通常、スケーリング ルールによって有効になります。たとえば、確率的勾配降下法では、バッチ サイズに応じて学習率を線形にスケーリングする必要があります。
もう 1 つの重要な機械学習ツールはモデル EMA です。これはターゲット モデルの機能コピーであり、そのパラメータは運動量ハイパーパラメータによってパラメータ化されたレートで指数移動平均 (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 machine learning tool is the model EMA, a functional copy of a target model, whose parameters move towards those of its target model according to an Exponential Moving Average (EMA) at a rate parameterized by a momentum hyperparameter. This model EMA can improve the robustness and generalization of supervised learning, stabilize pseudo-labeling, and provide a learning signal for Self-Supervised Learning (SSL). Prior works have not considered the optimization of the model EMA when performing scaling, 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 a model EMA and demonstrate the rule’s 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, a 6$\times$ wall-clock time reduction under idealized hardware settings.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google