Как масштабировать экспоненциальное скользящее среднее (EMA)

How to Scale Your EMA

July 25, 2023

Авторы: Dan Busbridge, Jason Ramapuram, Pierre Ablin, Tatiana Likhomanenko, Eeshan Gunesh Dhekane, Xavier Suau, Russ Webb

cs.AI

Аннотация

Сохранение динамики обучения при различных размерах пакетов является важным инструментом в практическом машинном обучении, так как позволяет находить компромисс между размером пакета и временем выполнения. Этот компромисс обычно достигается с помощью правила масштабирования. Например, в стохастическом градиентном спуске следует масштабировать скорость обучения линейно с размером пакета. Ещё одним важным инструментом в практическом машинном обучении является экспоненциальное скользящее среднее модели (EMA) — копия модели, которая не получает информацию о градиентах, а вместо этого следует за целевой моделью с определённым импульсом. Такая модель EMA может улучшить устойчивость и обобщающие свойства обучения с учителем, стабилизировать псевдоразметку и предоставить обучающий сигнал для самообучения (SSL). В предыдущих работах модель EMA рассматривалась отдельно от оптимизации, что приводило к различной динамике обучения при разных размерах пакетов и снижению производительности модели. В данной работе мы предлагаем правило масштабирования для оптимизации в присутствии моделей EMA и демонстрируем его применимость для различных архитектур, оптимизаторов и типов данных. Мы также показываем, что это правило работает в случаях, когда модель EMA участвует в оптимизации целевой модели, что позволяет обучать методы псевдоразметки и SSL, основанные на EMA, как при малых, так и при больших размерах пакетов. Для 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.