Wie Sie Ihr EMA skalieren können

Zusammenfassung

Die Beibehaltung der Trainingsdynamik über verschiedene Batch-Größen hinweg ist ein wichtiges Werkzeug für praktisches maschinelles Lernen, da es den Kompromiss zwischen Batch-Größe und Echtzeit ermöglicht. Dieser Kompromiss wird typischerweise durch eine Skalierungsregel ermöglicht, zum Beispiel sollte bei stochastischem Gradientenabstieg die Lernrate linear mit der Batch-Größe skaliert werden. Ein weiteres wichtiges Werkzeug für praktisches maschinelles Lernen ist das exponentielle gleitende Durchschnittsmodell (EMA), das eine Modellkopie darstellt, die keine Gradienteninformationen erhält, sondern stattdessen mit einem gewissen Momentum dem Zielmodell folgt. Dieses EMA-Modell kann die Robustheit und Generalisierungseigenschaften des überwachten Lernens verbessern, Pseudo-Labeling stabilisieren und ein Lernsignal für selbstüberwachtes Lernen (SSL) bereitstellen. Frühere Arbeiten haben das EMA-Modell getrennt von der Optimierung behandelt, was zu unterschiedlichen Trainingsdynamiken über verschiedene Batch-Größen hinweg und zu einer geringeren Modellleistung führte. In dieser Arbeit stellen wir eine Skalierungsregel für die Optimierung in Anwesenheit von EMA-Modellen bereit und demonstrieren deren Gültigkeit über eine Reihe von Architekturen, Optimierern und Datenmodalitäten hinweg. Wir zeigen auch die Gültigkeit der Regel, wenn das EMA-Modell zur Optimierung des Zielmodells beiträgt, was es uns ermöglicht, EMA-basierte Pseudo-Labeling- und SSL-Methoden sowohl bei kleinen als auch bei großen Batch-Größen zu trainieren. Für SSL ermöglichen wir das Training von BYOL bis zu einer Batch-Größe von 24.576 ohne Leistungseinbußen, was optimalerweise eine 6-fache Reduzierung der Echtzeit bedeutet.

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.

Wie Sie Ihr EMA skalieren können

How to Scale Your EMA

Zusammenfassung

Support