Proxies a pequeña escala para las inestabilidades en el entrenamiento de Transformers a gran escala

Resumen

Los equipos que han entrenado modelos grandes basados en Transformadores han reportado inestabilidades durante el entrenamiento a gran escala que no aparecían al entrenar con los mismos hiperparámetros a escalas más pequeñas. Aunque las causas de estas inestabilidades son de interés científico, los recursos necesarios para reproducirlas han dificultado su investigación. En este trabajo, buscamos formas de reproducir y estudiar la estabilidad e inestabilidad del entrenamiento a escalas más reducidas. Primero, nos enfocamos en dos fuentes de inestabilidad descritas en trabajos previos: el crecimiento de los logits en las capas de atención (Dehghani et al., 2023) y la divergencia de los logits de salida respecto a las probabilidades logarítmicas (Chowdhery et al., 2022). Al medir la relación entre la tasa de aprendizaje y la pérdida a diferentes escalas, demostramos que estas inestabilidades también aparecen en modelos pequeños cuando se entrena con tasas de aprendizaje altas, y que las mitigaciones empleadas previamente a gran escala son igualmente efectivas en este régimen. Esto nos lleva a investigar hasta qué punto otras intervenciones conocidas en el optimizador y el modelo influyen en la sensibilidad de la pérdida final a cambios en la tasa de aprendizaje. Para ello, estudiamos métodos como el calentamiento inicial (warm-up), la decaimiento de pesos (weight decay) y el muParam (Yang et al., 2022), y combinamos técnicas para entrenar modelos pequeños que logran pérdidas similares a lo largo de órdenes de magnitud de variación en la tasa de aprendizaje. Finalmente, para concluir nuestra exploración, estudiamos dos casos en los que las inestabilidades pueden predecirse antes de que surjan al examinar el comportamiento de escalado de las normas de activación y gradiente del modelo.

English

Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training stability and instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the muParam (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.

Proxies a pequeña escala para las inestabilidades en el entrenamiento de Transformers a gran escala

Small-scale proxies for large-scale Transformer training instabilities

Resumen

Support