Hoe Schaal Je Je EMA Op

Samenvatting

Het behouden van trainingsdynamiek over verschillende batchgroottes is een belangrijk hulpmiddel voor praktisch machine learning, omdat het de afweging tussen batchgrootte en werkelijke rekentijd mogelijk maakt. Deze afweging wordt doorgaans mogelijk gemaakt door een schaalregel; bijvoorbeeld, bij stochastische gradiëntdaling zou de leerrate lineair geschaald moeten worden met de batchgrootte. Een ander belangrijk hulpmiddel voor praktisch machine learning is het Exponentieel Voortschrijdend Gemiddelde (EMA) van het model, wat een modelkopie is die geen gradiëntinformatie ontvangt, maar in plaats daarvan zijn doelmodel volgt met enige momentum. Dit model EMA kan de robuustheid en generalisatie-eigenschappen van supervised learning verbeteren, pseudo-labeling stabiliseren, en een leersignaal bieden voor Self-Supervised Learning (SSL). Eerdere werken hebben het model EMA apart behandeld van optimalisatie, wat leidde tot verschillende trainingsdynamieken over batchgroottes en lagere modelprestaties. In dit werk bieden we een schaalregel voor optimalisatie in aanwezigheid van model EMA's en demonstreren we de geldigheid ervan over een reeks architecturen, optimalisatoren en datamodaliteiten. We tonen ook de geldigheid van de regel aan waar het model EMA bijdraagt aan de optimalisatie van het doelmodel, waardoor we EMA-gebaseerde pseudo-labeling en SSL-methoden kunnen trainen bij zowel kleine als grote batchgroottes. Voor SSL maken we het mogelijk om BYOL te trainen tot een batchgrootte van 24.576 zonder prestatieverlies, wat optimaal een 6-voudige reductie in werkelijke rekentijd oplevert.

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.

Hoe Schaal Je Je EMA Op

How to Scale Your EMA

Samenvatting

Support