Come Scalare la Tua EMA

Abstract

Preservare le dinamiche di addestramento attraverso diverse dimensioni del batch è uno strumento importante per il machine learning pratico, poiché consente di bilanciare la dimensione del batch e il tempo di esecuzione. Questo bilanciamento è tipicamente abilitato da una regola di scalatura; ad esempio, nella discesa del gradiente stocastica, si dovrebbe scalare il tasso di apprendimento linearmente con la dimensione del batch. Un altro strumento cruciale per il machine learning pratico è la Media Mobile Esponenziale (EMA) del modello, che è una copia del modello che non riceve informazioni sul gradiente, ma segue invece il modello target con un certo momento. Questa EMA del modello può migliorare la robustezza e le proprietà di generalizzazione dell'apprendimento supervisionato, stabilizzare il pseudo-labeling e fornire un segnale di apprendimento per l'Apprendimento Auto-Supervisionato (SSL). I lavori precedenti hanno trattato la EMA del modello separatamente dall'ottimizzazione, portando a dinamiche di addestramento diverse tra le dimensioni del batch e a prestazioni del modello inferiori. In questo lavoro, forniamo una regola di scalatura per l'ottimizzazione in presenza di EMA del modello e ne dimostriamo la validità attraverso una gamma di architetture, ottimizzatori e modalità di dati. Mostriamo inoltre la validità della regola nei casi in cui la EMA del modello contribuisce all'ottimizzazione del modello target, permettendoci di addestrare metodi di pseudo-labeling e SSL basati su EMA sia con batch piccoli che grandi. Per SSL, abilitiamo l'addestramento di BYOL fino a una dimensione del batch di 24.576 senza sacrificare le prestazioni, ottenendo una riduzione ottimale del tempo di esecuzione di 6 volte.

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.

Come Scalare la Tua EMA

How to Scale Your EMA

Abstract

Support