EMA를 확장하는 방법

초록

배치 크기 간 훈련 역학을 보존하는 것은 실용적인 머신러닝에서 중요한 도구입니다. 이는 배치 크기와 실제 소요 시간 간의 트레이드오프를 가능하게 하기 때문입니다. 이러한 트레이드오프는 일반적으로 스케일링 규칙을 통해 가능해지는데, 예를 들어 확률적 경사 하강법(SGD)에서는 학습률을 배치 크기에 선형적으로 스케일링해야 합니다. 실용적인 머신러닝을 위한 또 다른 중요한 도구는 모델 지수 이동 평균(EMA)입니다. 이는 그래디언트 정보를 받지 않지만 일정한 모멘텀으로 대상 모델을 따라가는 모델 복사본입니다. 이 모델 EMA는 지도 학습의 견고성과 일반화 성능을 향상시키고, 의사 레이블링을 안정화하며, 자기 지도 학습(SSL)을 위한 학습 신호를 제공할 수 있습니다. 기존 연구들은 모델 EMA를 최적화와 별도로 다루어 배치 크기 간 훈련 역학이 달라지고 모델 성능이 저하되는 문제가 있었습니다. 본 연구에서는 모델 EMA가 존재할 때의 최적화를 위한 스케일링 규칙을 제시하고, 다양한 아키텍처, 최적화 알고리즘, 데이터 모달리티에서 그 유효성을 입증합니다. 또한 모델 EMA가 대상 모델의 최적화에 기여하는 경우에도 이 규칙의 유효성을 보여, 작은 배치 크기와 큰 배치 크기 모두에서 EMA 기반 의사 레이블링 및 SSL 방법을 훈련할 수 있게 합니다. 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.

EMA를 확장하는 방법

How to Scale Your EMA

초록

Support