L'« Stochastic Weight Averaging » (SWA) est une technique de moyennage de modèles en Deep learning qui améliore la généralisation en moyennant les poids d'un Neural network sur plusieurs points de sa trajectoire d'entraînement. Contrairement à l'entraînement traditionnel qui ne conserve que les poids finaux, le SWA collecte les poids à intervalles réguliers pendant la phase tardive de l'entraînement et calcule leur moyenne arithmétique. Cette procédure simple conduit souvent à des solutions plus plates dans le paysage de la fonction de perte, lesquelles sont associées à de meilleures performances de test et à une meilleure robustesse. Le SWA a été introduit par Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov et Andrew Gordon Wilson dans leur article de 2018 intitulé « Averaging Weights Leads to Wider Optima and Better Generalization ». La méthode ne nécessite aucune modification de l'architecture du modèle sous-jacent ni de la fonction de perte, ce qui la rend facile à intégrer dans les pipelines d'entraînement existants.
L'idée centrale derrière le SWA est que pendant l'entraînement avec la descente de gradient stochastique (SGD) ou ses variantes, les paramètres du modèle oscillent autour d'une région de faible perte. La moyenne de ces paramètres peut lisser les oscillations et produire une solution qui se situe dans un minimum large et plat. De tels minima plats sont supposés généraliser mieux que les minima pointus car ils sont moins sensibles aux perturbations et aux décalages entre les distributions d'entraînement et de test. Il a été démontré que le SWA améliore la précision et l'étalonnage sur diverses architectures et tâches, notamment la classification d'images, la détection d'objets et la modélisation du langage.
Mécanisme et algorithme
Le SWA fonctionne en deux phases. Dans la première phase, le modèle est entraîné avec un calendrier de taux d'apprentissage standard (par exemple, un recuit cosinus ou une décroissance par paliers) pour atteindre une région proche d'un minimum. Dans la seconde phase, l'entraînement se poursuit avec un taux d'apprentissage constant ou cyclique, et les poids sont collectés à la fin de chaque époque ou après un nombre fixe d'itérations. Les poids SWA sont mis à jour comme une moyenne mobile : n_swa = n_swa + 1 ; w_swa = (w_swa * (n_swa - 1) + w) / n_swa, où w représente les poids actuels et w_swa les poids moyennés. Le modèle final utilise w_swa.
Un détail clé est que les statistiques de normalisation par lots doivent être recalculées après la moyenne, car la moyenne mobile et la variance des couches de normalisation par lots ne sont pas moyennées avec les poids. En pratique, après le calcul des poids SWA, une passe avant sur les données d'entraînement est effectuée pour mettre à jour les statistiques de normalisation par lots. Cette étape est cruciale pour les modèles qui utilisent la normalisation par lots, tels que les ResNets.
Motivation théorique
Le succès du SWA est souvent expliqué à travers le prisme de la géométrie du paysage de la fonction de perte. Les recherches d'Izmailov et al. et les travaux ultérieurs ont montré que la SGD tend à converger vers des points proches de la frontière d'une région de faible perte, tandis que la moyenne sur plusieurs points déplace la solution vers le centre de cette région. Ce point central se situe généralement dans un bassin plus plat, ce qui est associé à une meilleure généralisation. Le lien entre les minima plats et la généralisation a été étudié de manière approfondie, avec des connexions aux bornes PAC-Bayes et à la netteté de la fonction de perte.
Une autre perspective relie le SWA à l'inférence bayésienne. Sous certaines hypothèses, la trajectoire de la SGD peut être considérée comme une chaîne de Markov qui échantillonne une distribution a posteriori. La moyenne le long de cette trajectoire approxime la moyenne a posteriori, qui est la prédiction optimale au sens de Bayes pour une erreur quadratique. Cette interprétation relie le SWA aux méthodes de Monte Carlo par chaîne de Markov à gradient stochastique, bien que le SWA soit plus simple et ne nécessite pas un réglage minutieux des échelles de bruit.
Variantes et extensions
Plusieurs variantes du SWA ont été proposées pour améliorer encore les performances. Une extension notable est le « Stochastic Weight Averaging Gaussian » (SWAG), introduit par Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov et Andrew Gordon Wilson en 2019. Le SWAG approxime la distribution a posteriori des poids en ajustant une gaussienne à l'aide des premiers et seconds moments des poids collectés. Cela permet l'estimation de l'incertitude et un meilleur étalonnage, et peut être utilisé pour l'apprentissage profond bayésien.
Une autre variante est le « Fast Geometric Ensembling » (FGE), qui utilise un calendrier de taux d'apprentissage cyclique pour collecter des poids à partir de différents modes du paysage de la fonction de perte. Le SWA peut être combiné avec le FGE pour obtenir des performances encore meilleures. De plus, le SWA a été intégré à des techniques telles que les calendriers de taux d'apprentissage, la augmentation de données et l'élagage pour améliorer la robustesse et l'efficacité.
Applications
Le SWA a été appliqué avec succès dans un large éventail de domaines. En vision par ordinateur, il améliore la précision des classificateurs d'images sur des ensembles de données comme CIFAR-10, CIFAR-100 et ImageNet. Par exemple, il a été démontré que le SWA augmente la précision de test d'un PreAct ResNet-164 sur CIFAR-100 de 81,15 % à 82,90 % sans aucune modification architecturale. En détection d'objets, le SWA aide à stabiliser l'entraînement et à améliorer la précision moyenne.
En traitement du langage naturel, le SWA a été utilisé pour entraîner plus efficacement les Transformers et les grands modèles de langage. Il a été démontré qu'il réduit le surapprentissage et améliore la perplexité dans les tâches de modélisation du langage. Le SWA bénéficie également aux modèles génératifs tels que les GAN et les modèles de diffusion en fournissant des mises à jour de poids plus lisses, ce qui peut conduire à un entraînement plus stable et à des échantillons de meilleure qualité.
Au-delà de l'apprentissage supervisé standard, le SWA a été adapté pour l'apprentissage semi-supervisé, l'apprentissage par transfert et l'adaptation de domaine. Sa simplicité et sa généralité en font un outil précieux dans la boîte à outils du praticien en Machine learning.
Comparaison avec d'autres techniques
Le SWA est souvent comparé à d'autres méthodes de régularisation et d'ensachage. Contrairement aux ensembles traditionnels qui moyennent les prédictions de plusieurs modèles entraînés indépendamment, le SWA moyenne les poids, ce qui est moins coûteux en calcul et ne nécessite qu'une seule session d'entraînement. Cela le rend particulièrement attrayant lorsque le stockage du modèle ou le coût d'inférence est une préoccupation.
Comparé à des techniques comme le Dropout ou la normalisation par lots, le SWA ne modifie pas la dynamique d'entraînement mais traite plutôt les poids a posteriori. Il peut être utilisé en conjonction avec ces méthodes. Le SWA diffère également de l'Adam et d'autres optimiseurs adaptatifs, qui maintiennent des taux d'apprentissage par paramètre ; le SWA est généralement appliqué sur les trajectoires de la SGD ou d'Adam.
Une limitation du SWA est qu'il nécessite un calendrier pour déterminer quand commencer à collecter les poids. Si la collecte commence trop tôt, la moyenne peut être dominée par des poids sous-optimaux ; si elle commence trop tard, les avantages diminuent. Cependant, en pratique, une règle empirique simple consiste à commencer la collecte après que le taux d'apprentissage a diminué à une valeur faible.
Considérations pratiques
La mise en œuvre du SWA est simple. La plupart des frameworks d'apprentissage profond, tels que PyTorch et TensorFlow, fournissent des utilitaires ou des exemples pour le SWA. Les principales étapes sont les suivantes : (1) entraîner le modèle avec un calendrier standard, (2) poursuivre l'entraînement avec un taux d'apprentissage constant ou cyclique, (3) collecter les poids à intervalles réguliers, (4) calculer la moyenne mobile, et (5) recalculer les statistiques de normalisation par lots.
Il a été démontré que le SWA est robuste aux choix des hyperparamètres. Le taux d'apprentissage pour la seconde phase est généralement défini sur la valeur minimale du calendrier initial ou sur une petite constante. La fréquence de collecte des poids peut être chaque époque ou toutes les quelques centaines d'itérations. En pratique, collecter les poids pendant 5 à 20 époques est souvent suffisant.
Impact et réception
L'introduction du SWA a eu un impact significatif sur la communauté de l'apprentissage profond. Il a fourni une méthode simple mais efficace pour améliorer la généralisation sans coût de calcul supplémentaire lors de l'inférence. L'article a été largement cité et a inspiré de nombreux travaux de suivi sur la moyenne des poids et les minima plats. Le SWA est désormais une technique standard dans de nombreux pipelines d'entraînement, en particulier dans la recherche académique et les compétitions.
Le SWA a également influencé le développement d'autres méthodes de moyennage, telles que la moyenne mobile exponentielle (EMA), qui est couramment utilisée dans l'entraînement des modèles génératifs et des grands modèles de langage. L'EMA est une variante qui accorde plus de poids aux poids récents, et elle est souvent utilisée en pratique en raison de sa simplicité et de son efficacité. Le SWA et l'EMA sont complémentaires, et certains frameworks permettent de passer de l'un à l'autre.
Limites et orientations futures
Malgré ses avantages, le SWA présente des limites. Il suppose que le paysage de la fonction de perte est relativement lisse et que la moyenne des poids est significative, ce qui peut ne pas être le cas pour toutes les architectures ou toutes les tâches. Par exemple, les modèles avec normalisation par lots nécessitent une manipulation minutieuse, et les modèles avec des espaces de paramètres complexes (par exemple, les réseaux récurrents) peuvent ne pas bénéficier autant. De plus, le SWA ne fournit pas d'estimations d'incertitude à moins d'être étendu (par exemple, avec le SWAG).
Les orientations de recherche futures incluent le développement de calendriers de moyennage adaptatifs, la combinaison du SWA avec des méthodes bayésiennes et l'application du SWA à de nouveaux domaines tels que l'apprentissage par renforcement et l'apprentissage fédéré. Alors que l'apprentissage profond continue d'évoluer, les techniques de moyennage des poids comme le SWA restent un outil fondamental pour améliorer les performances et la fiabilité des modèles.