Le warmup est une technique d'entraînement utilisée en apprentissage automatique et en apprentissage profond qui augmente progressivement le taux d'apprentissage d'une petite valeur initiale à une valeur cible sur un nombre défini de premières étapes d'entraînement. Cette approche vise à stabiliser la phase précoce de l'entraînement des réseaux de neurones, en particulier des grands modèles tels que les grands modèles de langage et les transformers. En évitant un taux d'apprentissage élevé au début, le warmup prévient les gradients instables et aide le modèle à converger vers une meilleure solution finale.
La technique est souvent associée à un calendrier de décroissance ultérieur, comme l'annealing cosinusoïdal ou la décroissance linéaire, formant un calendrier de taux d'apprentissage complet. Le warmup est désormais un composant standard dans l'entraînement de nombreux modèles de pointe, y compris ceux développés par OpenAI, Anthropic et Google DeepMind. Il est généralement implémenté via un nombre spécifique de étapes de warmup ou un calendrier de warmup qui définit comment le taux augmente.
Motivation et Problème
Pendant la phase initiale de l'entraînement, les poids du modèle sont initialisés aléatoirement, souvent avec des méthodes comme l'initialisation Xavier ou He. À ce stade, les gradients peuvent être importants et bruités, surtout pour les réseaux profonds. Si un taux d'apprentissage élevé est appliqué immédiatement, l'optimiseur peut prendre des pas importants et erratiques, provoquant une divergence de la perte ou un repli du modèle dans un minimum local médiocre. Ce problème est exacerbé dans les modèles avec de nombreuses couches, comme les réseaux résiduels ou les transformers, où l'amplitude des gradients peut varier considérablement entre les couches.
Le warmup répond à cela en permettant au modèle de se déplacer d'abord dans une région du paysage de perte où les gradients sont plus stables et informatifs. Le petit taux d'apprentissage initial agit comme une forme de régularisation, empêchant des mises à jour drastiques qui pourraient détruire la configuration initiale des poids. Cela est particulièrement important pour les modèles qui utilisent des techniques comme la normalisation de couche ou la normalisation par lots, où les statistiques des activations ont besoin de temps pour se stabiliser.
Contexte Historique
Le concept de warmup n'est pas nouveau, mais il a gagné en importance avec l'essor de l'apprentissage profond. Les travaux précoces des années 2010 sur l'entraînement des réseaux de neurones profonds, comme ceux de l'université de Toronto et du laboratoire d'IA de Stanford, ont noté les avantages d'augmenter progressivement le taux d'apprentissage. Cependant, il est devenu une pratique largement adoptée après l'introduction de l'architecture Transformer dans l'article de 2017 « Attention Is All You Need » par des chercheurs de Google, dont Jakob Uszkoreit et Lukasz Kaiser. Cet article recommandait explicitement une phase de warmup de 4000 étapes pour entraîner les modèles transformer, ce qui est devenu une valeur par défaut courante.
Depuis lors, le warmup a été intégré dans de nombreux cadres d'entraînement et est un paramètre par défaut dans des bibliothèques comme les Transformers de Hugging Face et les planificateurs de taux d'apprentissage de PyTorch. Son adoption a grandi parallèlement à la mise à l'échelle des modèles, car les modèles plus grands sont plus sensibles aux choix d'hyperparamètres.
Comment Fonctionne le Warmup
Le warmup est généralement défini par un nombre d'étapes de warmup, noté \(W\). Pendant ces étapes, le taux d'apprentissage augmente linéairement ou quadratiquement d'une petite valeur (souvent proche de zéro) à un taux d'apprentissage cible \(\eta_{max}\). Un calendrier linéaire courant est :
\[ \eta_t = \eta_{max} \cdot \frac{t}{W} \]
pour \(t = 1, 2, \ldots, W\). Après la période de warmup, le taux d'apprentissage suit un calendrier de décroissance, comme la décroissance cosinusoïdale, la décroissance exponentielle ou une réduction par paliers. Certains calendriers utilisent un warmup non linéaire, comme une augmentation en racine carrée ou cubique, ce qui peut être bénéfique pour certaines architectures.
Le choix de \(W\) dépend de la taille du modèle et du jeu de données. Pour les petits modèles, le warmup peut durer seulement quelques centaines d'étapes, tandis que pour les grands modèles de langage, il peut s'étendre sur des milliers, voire des dizaines de milliers d'étapes. Par exemple, l'entraînement d'un modèle comme GPT-3, qui a des centaines de milliards de paramètres, a nécessité un warmup de plusieurs milliers d'étapes pour éviter l'instabilité.
Bénéfices et Effets
Le warmup offre plusieurs avantages concrets. Premièrement, il améliore la stabilité de l'entraînement, réduisant le risque de divergence de la perte ou de valeurs NaN dans les gradients. Deuxièmement, il conduit souvent à de meilleures performances finales, car le modèle peut explorer une région plus favorable du paysage de perte. Troisièmement, il permet l'utilisation de taux d'apprentissage maximaux plus élevés, ce qui peut accélérer la convergence après la phase de warmup.
Des études empiriques ont montré que le warmup peut être particulièrement efficace lorsqu'il est combiné avec l'optimiseur Adam et d'autres optimiseurs adaptatifs. Adam maintient des taux d'apprentissage par paramètre basés sur l'historique des gradients, et pendant les premières étapes, ces estimations sont bruitées. Le warmup atténue cela en limitant la taille des pas jusqu'à ce que les estimations de variance de l'optimiseur se stabilisent.
De plus, le warmup interagit avec d'autres techniques comme le écrêtage des gradients et l'initialisation des poids. Par exemple, lors de l'utilisation de réseaux résiduels, le warmup aide à garantir que les connexions résiduelles n'amplifient pas les erreurs initiales.
Variantes et Extensions
Plusieurs variantes du warmup ont été développées. Une variante courante est le « warmup graduel », où le taux d'apprentissage augmente linéairement sur un nombre fixe d'étapes. Une autre est le « warmup exponentiel », qui augmente le taux de manière exponentielle, bien que cela soit moins courant. Certains calendriers utilisent une approche « warmup puis décroissance », où le warmup est suivi d'une décroissance cosinusoïdale jusqu'à zéro, comme popularisé par le calendrier « Noam » dans l'article original sur le Transformer.
Une autre extension est le « warmup stochastique », où les étapes de warmup sont échantillonnées aléatoirement à chaque exécution d'entraînement, ce qui peut agir comme une forme de régularisation des hyperparamètres. Dans certains cas, le warmup est appliqué uniquement à certaines couches ou paramètres, comme la couche d'embedding dans un transformer, tandis que d'autres utilisent un warmup complet.
Le warmup est également lié à l'apprentissage par curriculum, où le modèle est d'abord entraîné sur des exemples plus faciles avant des plus difficiles. Les deux techniques visent à stabiliser l'entraînement précoce, mais elles opèrent sur des axes différents : le warmup ajuste le taux d'apprentissage, tandis que l'apprentissage par curriculum ajuste la distribution des données.
Implémentation Pratique
En pratique, l'implémentation du warmup nécessite de définir trois hyperparamètres : le taux d'apprentissage initial (souvent fixé à une petite fraction du taux maximal, comme 0,1 ou 0,01), le taux d'apprentissage maximal et le nombre d'étapes de warmup. Le taux d'apprentissage maximal est généralement déterminé par réglage des hyperparamètres, et les étapes de warmup sont souvent définies comme un pourcentage du nombre total d'étapes d'entraînement, comme 1 % à 10 %.
Par exemple, lors de l'entraînement d'un modèle Transformer (architecture) pour la traduction automatique, une configuration courante est un taux d'apprentissage maximal de 0,0001 avec 4000 étapes de warmup. En revanche, pour un grand modèle de langage avec des billions de jetons, le warmup pourrait durer 2000 étapes sur un total de 500 000 étapes.
Des cadres comme PyTorch et TensorFlow fournissent des planificateurs intégrés qui prennent en charge le warmup. Par exemple, la fonction get_cosine_schedule_with_warmup de PyTorch crée un calendrier avec un warmup linéaire suivi d'une décroissance cosinusoïdale. De même, get_scheduler de Hugging Face permet de spécifier un ratio de warmup.
Relation avec d'Autres Techniques
Le warmup est souvent utilisé en conjonction avec d'autres techniques de stabilisation. Le écrêtage des gradients limite l'amplitude des gradients, ce qui peut prévenir les explosions même sans warmup, mais la combinaison est plus robuste. La normalisation de couche et la normalisation par lots aident à stabiliser les activations, mais elles ne résolvent pas directement le problème des grands gradients initiaux.
Le warmup complète également le Dropout et d'autres méthodes de régularisation. En commençant avec un taux d'apprentissage faible, le modèle est moins susceptible de surajuster le bruit dans les premières étapes, ce qui permet au dropout d'être plus efficace plus tard.
De plus, le warmup est un composant clé de la recette d'entraînement pour de nombreuses architectures modernes, y compris le U-Net pour la segmentation d'images et divers modèles séquence à séquence. Il est également utilisé dans le fine-tuning de modèles pré-entraînés, où un court warmup (par exemple, 100 étapes) aide à adapter le modèle à une nouvelle tâche sans oubli catastrophique.
Défis et Limitations
Malgré ses avantages, le warmup introduit des hyperparamètres supplémentaires qui nécessitent un réglage. Choisir trop peu d'étapes de warmup peut conduire à l'instabilité, tandis qu'un nombre trop élevé peut ralentir la convergence. La durée optimale du warmup varie selon le modèle et le jeu de données, et il n'existe pas de règle universelle.
Certaines recherches ont exploré des alternatives au warmup, comme l'utilisation d'un taux d'apprentissage faible fixe ou l'emploi de méthodes adaptatives qui ajustent le taux en fonction des statistiques des gradients. Cependant, ces approches n'ont pas complètement remplacé le warmup en pratique, car le warmup reste une référence simple et efficace.
Une autre limitation est que le warmup est principalement conçu pour la descente de gradient stochastique et ses variantes. Pour d'autres algorithmes d'optimisation, comme les méthodes du second ordre, les bénéfices peuvent être moins prononcés, bien qu'ils soient encore utilisés dans certains cas.
Directions Futures
Alors que les modèles continuent de croître en taille, le warmup reste un domaine de recherche actif. Des travaux récents ont étudié des méthodes automatisées pour définir la durée du warmup, comme l'utilisation de la perte de validation pour déclencher la transition du warmup à la décroissance. Il y a aussi un intérêt pour comprendre les fondements théoriques du warmup, certaines études le reliant à la courbure du paysage de perte.
Dans le contexte de l'IA générative et de l'entraînement à grande échelle, le warmup est susceptible de rester une pratique standard. Des entreprises comme OpenAI, Anthropic et Google DeepMind ont toutes publié des détails d'entraînement incluant des phases de warmup, et c'est une valeur par défaut dans de nombreux scripts d'entraînement open source. À mesure que de nouvelles architectures émergent, le warmup continuera d'être adapté et affiné.
En résumé, le warmup est une technique simple mais puissante qui améliore la stabilité et les performances de l'entraînement des réseaux de neurones. En augmentant progressivement le taux d'apprentissage, il permet aux modèles de naviguer dans la phase chaotique initiale de l'optimisation, conduisant à de meilleurs résultats finaux. Son adoption généralisée dans le domaine souligne son importance dans l'apprentissage automatique moderne.