Cómo Escalar tu EMA

Resumen

Preservar la dinámica de entrenamiento a través de diferentes tamaños de lote es una herramienta importante para el aprendizaje automático práctico, ya que permite equilibrar el tamaño del lote y el tiempo de ejecución. Este equilibrio suele habilitarse mediante una regla de escalado; por ejemplo, en el descenso de gradiente estocástico, se debe escalar la tasa de aprendizaje linealmente con el tamaño del lote. Otra herramienta clave para el aprendizaje automático práctico es el Promedio Móvil Exponencial (EMA, por sus siglas en inglés) del modelo, que es una copia del modelo que no recibe información de gradientes, sino que sigue a su modelo objetivo con cierto impulso. Este EMA del modelo puede mejorar la robustez y las propiedades de generalización del aprendizaje supervisado, estabilizar el etiquetado pseudo-supervisado y proporcionar una señal de aprendizaje para el Aprendizaje Autosupervisado (SSL). Trabajos anteriores han tratado el EMA del modelo de manera separada a la optimización, lo que ha llevado a dinámicas de entrenamiento diferentes según el tamaño del lote y a un menor rendimiento del modelo. En este trabajo, proporcionamos una regla de escalado para la optimización en presencia de EMAs del modelo y demostramos su validez en una variedad de arquitecturas, optimizadores y modalidades de datos. También mostramos la validez de la regla cuando el EMA del modelo contribuye a la optimización del modelo objetivo, permitiéndonos entrenar métodos de etiquetado pseudo-supervisado y SSL basados en EMA con tamaños de lote pequeños y grandes. Para SSL, habilitamos el entrenamiento de BYOL hasta un tamaño de lote de 24,576 sin sacrificar el rendimiento, logrando óptimamente una reducción de 6 veces en el tiempo de ejecución.

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.

Cómo Escalar tu EMA

How to Scale Your EMA

Resumen

Support