L'accumulation de gradients est une technique utilisée lors de l'entraînement des réseaux de neurones et d'autres modèles de Machine learning, notamment en Deep learning, pour approximer l'effet d'une grande taille de lot lorsque les contraintes de mémoire empêchent de traiter tous les échantillons en une seule fois. Au lieu de calculer le gradient sur un grand lot en une seule passe avant et arrière, l'optimiseur accumule les gradients sur plusieurs mini-lots plus petits et n'effectue une mise à jour des poids qu'après que le gradient accumulé atteint la taille de lot effective souhaitée. Cette approche permet aux praticiens d'entraîner des modèles avec des tailles de lot qui dépasseraient autrement la capacité mémoire d'un seul dispositif, comme un GPU ou un TPU.
La méthode est étroitement liée à la descente de gradient stochastique (SGD), un algorithme d'optimisation itératif introduit dans les années 1950 via l'algorithme de Robbins-Monro. Dans la SGD, le gradient réel de la fonction objectif, calculé sur l'ensemble du jeu de données, est approximé par le gradient sur un sous-ensemble de données sélectionné aléatoirement. Un compromis courant est l'approche par mini-lots, où le gradient est calculé sur un petit ensemble d'échantillons à chaque étape. L'accumulation de gradients étend cette idée en moyennant les gradients sur plusieurs mini-lots avant de mettre à jour les paramètres du modèle, découplant ainsi efficacement la taille de lot de l'empreinte mémoire.
Motivation et contraintes de mémoire
L'entraînement de grands modèles, tels que les grands modèles de langage et les transformeurs, nécessite une mémoire substantielle pour stocker les activations, les gradients et les états de l'optimiseur. La taille de lot influence directement la consommation mémoire : des lots plus grands exigent plus de mémoire pour les calculs intermédiaires. Sur du matériel à mémoire limitée, comme les GPU grand public ou les dispositifs de périphérie, la taille de lot maximale réalisable peut être petite. Cependant, de petites tailles de lot peuvent entraîner des estimations de gradient bruitées et des dynamiques d'entraînement instables. L'accumulation de gradients offre une solution en permettant l'utilisation de petits lots pour les passes avant et arrière tout en accumulant les gradients sur plusieurs étapes pour atteindre une taille de lot effective plus grande, améliorant ainsi la stabilité du gradient sans augmenter l'utilisation mémoire de pointe.
Fonctionnement de l'accumulation de gradients
Dans l'entraînement standard par mini-lots, le modèle traite un mini-lot de taille \(b\), calcule la perte, effectue la rétropropagation pour obtenir les gradients, puis met immédiatement à jour les poids à l'aide d'un optimiseur tel que SGD ou Adam. Avec l'accumulation de gradients, le processus est modifié : le modèle traite \(k\) mini-lots séquentiellement, accumulant les gradients (généralement en les sommant ou en les moyennant) après chaque passe arrière. Après \(k\) étapes, le gradient accumulé est utilisé pour mettre à jour les poids du modèle, et l'accumulateur de gradients est remis à zéro. La taille de lot effective est \(k \times b\). Cette approche est mathématiquement équivalente à un entraînement avec une taille de lot de \(k \times b\) lorsque la fonction de perte est une somme ou une moyenne sur les échantillons, à condition que les couches de normalisation par lots soient gérées avec soin.
Avantages et cas d'utilisation
L'accumulation de gradients est largement utilisée dans les scénarios où de grandes tailles de lot sont bénéfiques mais irréalisables en raison des limites de mémoire. Par exemple, l'entraînement de modèles génératifs ou le réglage fin de grands modèles de langage nécessite souvent des tailles de lot de centaines ou de milliers pour stabiliser l'entraînement et améliorer la convergence. En accumulant les gradients, les chercheurs peuvent simuler ces grands lots sur du matériel à mémoire limitée, comme un seul GPU ou même des systèmes basés sur CPU. De plus, l'accumulation de gradients permet l'entraînement sur plusieurs dispositifs sans synchroniser les gradients à chaque étape, ce qui peut réduire la surcharge de communication dans l'entraînement distribué. Cela est particulièrement pertinent pour les plateformes cloud comme Amazon Web Services, Microsoft Azure et Google Cloud, qui proposent des instances GPU avec des capacités mémoire variables.
Relation avec la taille de lot et le taux d'apprentissage
Lors de l'utilisation de l'accumulation de gradients, la taille de lot effective augmente, ce qui nécessite souvent d'ajuster le taux d'apprentissage. En pratique, une heuristique courante consiste à mettre à l'échelle le taux d'apprentissage linéairement avec la taille de lot, comme suggéré dans des recherches antérieures sur l'entraînement avec de grands lots. Cependant, la mise à l'échelle optimale dépend de l'architecture du modèle, de l'optimiseur et du jeu de données. L'accumulation de gradients ne change pas le nombre d'échantillons vus par mise à jour des poids ; elle modifie seulement la manière dont les gradients sont agrégés. Par conséquent, le nombre total d'étapes d'entraînement diminue à mesure que la taille de lot effective augmente, ce qui peut affecter la vitesse de convergence et les performances finales. Les praticiens doivent ajuster le taux d'apprentissage et les étapes d'accumulation pour obtenir des résultats comparables à un véritable entraînement avec de grands lots.
Considérations d'implémentation
L'implémentation de l'accumulation de gradients implique de modifier la boucle d'entraînement. Après chaque passe avant et arrière, les gradients sont accumulés dans les tampons de gradients du modèle (par exemple, en utilisant accumulate_grad de PyTorch ou GradientTape de TensorFlow avec des variables persistantes). L'étape de l'optimiseur n'est effectuée qu'après le nombre souhaité d'étapes d'accumulation. Il est crucial de mettre à l'échelle la perte de manière appropriée : si la perte est moyennée sur le mini-lot, les gradients accumulés doivent être divisés par le nombre d'étapes d'accumulation pour maintenir la magnitude correcte du gradient. Alternativement, si la perte est sommée, aucune mise à l'échelle n'est nécessaire. De plus, les couches de normalisation par lots nécessitent une attention particulière car elles calculent des statistiques sur le mini-lot actuel. Avec l'accumulation de gradients, la taille de lot effective pour la normalisation reste la taille du mini-lot, ce qui peut entraîner un comportement différent par rapport à un véritable entraînement avec de grands lots. Certains frameworks offrent une gestion automatique via une normalisation par lots synchronisée entre les dispositifs.
Comparaison avec d'autres techniques
L'accumulation de gradients est souvent comparée au checkpointing de gradients, qui réduit la mémoire en recalculant les activations pendant la rétropropagation, et à l'entraînement en précision mixte, qui réduit la mémoire en utilisant une arithmétique de précision inférieure. Ces techniques peuvent être combinées : par exemple, utiliser l'accumulation de gradients avec la précision mixte permet d'entraîner des tailles de lot effectives encore plus grandes. Contrairement à l'entraînement distribué sur plusieurs GPU, l'accumulation de gradients ne nécessite pas de communication entre les dispositifs, ce qui la rend plus simple à implémenter et moins sujette à la surcharge de synchronisation. Cependant, elle augmente le nombre d'étapes séquentielles, ce qui peut ralentir l'entraînement si la taille du mini-lot est très petite et que la surcharge de multiples passes arrière est significative.
Limites et alternatives
L'accumulation de gradients n'est pas un substitut parfait à un véritable entraînement avec de grands lots. Les estimations de gradient sont calculées à partir de différents mini-lots, ce qui peut introduire de légères différences en raison du mélange des données et des statistiques de normalisation par lots. Dans certains cas, cela peut entraîner un comportement de convergence différent. Les alternatives incluent l'utilisation de dispositifs à plus grande mémoire, comme les systèmes Cerebras ou les accélérateurs Groq, qui offrent une bande passante et une capacité mémoire élevées, ou l'utilisation du parallélisme de modèle et du parallélisme de pipeline pour distribuer le lot entre les dispositifs. De plus, certains optimiseurs, comme LAMB ou LARS, sont conçus pour gérer directement de grandes tailles de lot, réduisant potentiellement le besoin d'accumulation.
Contexte historique et adoption
Le concept d'accumulation de gradients sur plusieurs mini-lots précède l'apprentissage profond moderne, avec des racines dans l'"algorithme de rétropropagation en mode groupe" des années 1980. Il a gagné en importance dans les années 2010 à mesure que les modèles d'apprentissage profond ont grandi en taille et que les contraintes de mémoire sont devenues un goulot d'étranglement. Aujourd'hui, l'accumulation de gradients est une fonctionnalité standard dans les principaux frameworks d'apprentissage profond, notamment PyTorch, TensorFlow et JAX, et est largement utilisée pour entraîner des modèles de pointe dans des organisations telles que OpenAI, Anthropic et Google DeepMind. C'est également une technique courante dans les environnements de recherche et de production, en particulier pour le réglage fin de grands modèles sur du matériel limité.
Conclusion
L'accumulation de gradients est une technique pratique et largement adoptée pour entraîner de grands modèles sous contraintes de mémoire. En accumulant les gradients sur plusieurs mini-lots, elle permet de simuler de grandes tailles de lot sans nécessiter des augmentations mémoire proportionnelles. Bien qu'elle introduise une certaine surcharge et exige un réglage minutieux des taux d'apprentissage et des couches de normalisation, elle reste un outil essentiel dans la boîte à outils du praticien en apprentissage profond, en particulier pour l'entraînement à grande échelle et les tâches de réglage fin.