Des substituts à petite échelle pour les instabilités de l'entraînement des Transformers à grande échelle

Résumé

Les équipes ayant entraîné des modèles de grande taille basés sur des Transformers ont rapporté des instabilités lors de l'entraînement à grande échelle, qui n'apparaissaient pas lors de l'entraînement avec les mêmes hyperparamètres à des échelles plus réduites. Bien que les causes de ces instabilités soient d'un intérêt scientifique, les ressources nécessaires pour les reproduire ont rendu leur investigation difficile. Dans ce travail, nous cherchons des moyens de reproduire et d'étudier la stabilité et l'instabilité de l'entraînement à des échelles plus petites. Tout d'abord, nous nous concentrons sur deux sources d'instabilité décrites dans des travaux antérieurs : la croissance des logits dans les couches d'attention (Dehghani et al., 2023) et la divergence des logits de sortie par rapport aux log-probabilités (Chowdhery et al., 2022). En mesurant la relation entre le taux d'apprentissage et la perte à différentes échelles, nous montrons que ces instabilités apparaissent également dans les petits modèles lors de l'entraînement à des taux d'apprentissage élevés, et que les atténuations précédemment employées à grande échelle sont tout aussi efficaces dans ce régime. Cela nous incite à étudier dans quelle mesure d'autres interventions connues sur l'optimiseur et le modèle influencent la sensibilité de la perte finale aux variations du taux d'apprentissage. À cette fin, nous étudions des méthodes telles que le warm-up, la décroissance de poids et le muParam (Yang et al., 2022), et combinons des techniques pour entraîner des petits modèles qui atteignent des pertes similaires sur plusieurs ordres de grandeur de variation du taux d'apprentissage. Enfin, pour conclure notre exploration, nous étudions deux cas où les instabilités peuvent être prédites avant leur apparition en examinant le comportement de mise à l'échelle des normes des activations et des gradients du modèle.

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.

Des substituts à petite échelle pour les instabilités de l'entraînement des Transformers à grande échelle

Small-scale proxies for large-scale Transformer training instabilities

Résumé

Support