Tilde Research présente Aurora, un optimiseur qui corrige la mort neuronale cachée dans Muon

Des chercheurs de Tilde Research ont publié Aurora, un nouvel optimiseur pour l'entraînement de réseaux de neurones qui corrige un défaut structurel présent dans Muon, l'un des optimiseurs les plus adoptés par la communauté depuis 2024. Le problème découvert est frappant : dans les matrices de poids "tall", comme celles des couches MLP basées sur l'architecture SwiGLU, Muon tue silencieusement une fraction importante des neurones au fil de l'entraînement. Dès la 500e étape d'entraînement, plus d'un neurone sur quatre est effectivement mort et ne se réveille plus. Aurora propose une correction mathématique rigoureuse, accompagnée d'un benchmark validé sur un modèle de 1,1 milliard de paramètres, d'un nouveau record sur le classement modded-nanoGPT speedrun, et d'un code publié en accès libre.

Ce bug discret a des conséquences concrètes sur la qualité des modèles entraînés avec Muon. Le mécanisme central de Muon repose sur le calcul d'un "facteur polaire" à partir du gradient, ce qui produit une mise à jour quasi-orthogonale des poids. Mais pour les matrices tall, il est mathématiquement impossible de rester orthogonal tout en maintenant des mises à jour uniformes entre les neurones : l'optimiseur finit par sur-mettre à jour certains neurones et en ignorer d'autres. Les neurones sous-alimentés reçoivent de moins en moins de signal, entrent dans une spirale de mort, et cessent de contribuer au calcul. Cette inactivité se propage aux couches suivantes, privant tout le réseau d'une partie de sa capacité de représentation, sans qu'aucun signal d'erreur ne l'indique clairement.

Muon avait attiré l'attention de la communauté ML en surpassant AdamW en vitesse de convergence sur le benchmark nanoGPT speedrun, une compétition mesurant la rapidité d'entraînement d'un modèle de type GPT. Depuis, plusieurs groupes de recherche travaillant à l'échelle frontier l'ont adopté. Une variante intermédiaire, NorMuon, avait tenté d'améliorer Muon en normalisant les lignes du gradient, mais visait la mauvaise cible : elle imposait une norme de 1 à chaque ligne, alors que la valeur mathématiquement correcte pour une matrice tall est √(n/m). Tilde Research a d'abord formalisé cette correction sous le nom U-NorMuon, qui élimine complètement la mort neuronale à 340 millions de paramètres, avant de développer Aurora comme solution complète préservant à la fois l'uniformité des mises à jour et la précision du facteur polaire. La publication en code ouvert ouvre la voie à une adoption rapide dans les pipelines d'entraînement existants.