Un nouveau package Python, AutoMuon, permet d’utiliser l’optimiseur Muon comme un remplacement direct d’AdamW dans les pipelines d’entraînement PyTorch. Cette initiative vise à simplifier l’intégration de Muon, un optimiseur dont l’application était jusqu’à présent plus complexe.
L’optimiseur Muon est principalement conçu pour les matrices de poids 2D, telles que celles des projections linéaires et des couches de convolution agissant sur les états cachés. Cependant, d’autres composants des modèles, comme les embeddings, les normalisations et les biais, nécessitent toujours l’utilisation d’optimiseurs comme AdamW.
AutoMuon résout cette problématique en scannant automatiquement le modèle lors de son initialisation. Il identifie alors les paramètres appropriés pour chaque optimiseur, assignant Muon ou AdamW selon les besoins spécifiques de chaque couche. Le développeur se déclare ouvert aux contributions via des « pull requests ».
Cette approche « drop-in » rend l’optimiseur Muon plus accessible aux développeurs PyTorch, ouvrant la voie à son adoption dans de nouvelles architectures de réseaux de neurones.
Source : Reddit r/MachineLearning