Résumé Technique : Exploiter les symétries de l'espace des poids pour approximer la courbure

Énoncé du Problème

L'estimation précise de la courbure (la matrice Hessienne) d'une fonction de perte est fondamentale pour de nombreux sous-domaines de l'apprentissage automatique, notamment l'optimisation de second ordre, l'apprentissage profond bayésien, l'apprentissage continu et la compression de modèles. Cependant, la construction, le stockage et l'inversion de la matrice Hessienne complète est trop coûteux en termes de calcul et de mémoire pour les réseaux profonds modernes. Bien que les techniques d'approximation existantes telles que (E)KFAC, Shampoo et Soap aient atteint une certaine scalabilité en supposant des structures spécifiques de factorisation de Kronecker ou de blocs diagonaux, le domaine manque d'un cadre théorique rigoureux expliquant pourquoi ces structures spécifiques fonctionnent ou comment les dériver systématiquement en fonction de l'architecture du modèle. De plus, les travaux précédents ont largement négligé les contraintes structurelles sur la courbure imposées par les symétries bien connues de l'espace des poids (par exemple, l'invariance par permutation des unités cachées).

Méthodologie

Les auteurs proposent un nouveau cadre pour approximer la Hessienne en effectuant une moyenne analytique sur l'« orbite de groupe » générée par les symétries de l'espace des poids. L'idée centrale repose sur le fait que de nombreuses fonctions de perte sont invariantes sous des transformations de paramètres spécifiques (actions de groupe), ce qui implique que le gradient de la perte est équivariant sous ces mêmes transformations.

1. Moyennage d'orbite et équivariance du gradient

Étant donné un ensemble actuel de paramètres $w$ et un gradient $g$, les auteurs considèrent l'ensemble de tous les couples $(w^{(i)}, g^{(i)})$ atteignables en appliquant des transformations de symétrie issues d'un groupe $G$. Parce que la perte est invariante, le gradient à n'importe quel point de l'orbite est lié au gradient original par l'action du groupe. En moyennant ces gradients sur l'ensemble de l'orbite, la méthode agrège l'information de courbure sans nécessiter de multiples calculs de gradient.

2. Approximation structurée de la Hessienne

La méthode dérive une approximation structurée de la Hessienne, $\hat{H}$, en résolvant une condition de sécante basée sur les moyennes d'orbite des différences de paramètres et des différences de gradients.

• Algèbre du commutant : L'approximation de la Hessienne résultante appartient à l'algèbre du commutant du groupe de symétrie. Cela signifie que l'approximation peut être représentée comme une superposition linéaire d'un petit nombre d'éléments de base (tenseurs creux et binaires formés par des produits de Kronecker de matrices identité et de matrices à un seul type de valeur, ici des matrices remplies de uns).
• Factorisation : Au lieu de stocker la matrice complète, la méthode ne stocke que les « facteurs » scalaires ou matriciels qui pondèrent ces éléments de base. Cela permet une estimation, un stockage et une inversion tractables.
• Simplification : Les auteurs proposent une simplification pratique où la moyenne d'orbite des différences de paramètres est supposée proportionnelle à la matrice identité. Cela conduit à une estimation de la courbure $\hat{H}_g = S_g^{1/2}$, où $S_g$ est la moyenne d'orbite du produit extérieur du gradient.

3. L'optimiseur « Symo »

Le cadre est instancié sous la forme d'un nouvel optimiseur appelé Symo (optimiseur basé sur la symétrie). La règle de mise à jour est $w_{t+1} = w_t - \eta (\hat{H}_g)^{-1} g_t$. Les auteurs fournissent un compilateur « juste-à-temps » (just-in-time) qui prend les spécifications symboliques des symétries d'un réseau et génère des routines PyTorch efficaces pour calculer les facteurs nécessaires.

Principales Contributions

1. Nouvelle approximation de la Hessienne : Une méthode pour dériver une estimation structurée de la Hessienne à partir d'un seul calcul de gradient en moyennant sur l'orbite du groupe de symétrie. L'approximation est prouvée hautement structurée et réside dans l'algèbre du commutant.
2. Unification théorique de Shampoo et Muon : Le papier prouve que des choix spécifiques de groupes de symétrie au sein de ce cadre récupèrent exactement les règles de mise à jour des optimiseurs Shampoo et Muon.
   • En restreignant le groupe de symétrie à des sous-groupes de permutations signées spécifiques (par exemple, en appliquant une symétrie d'identité aux couches d'entrée/sortie et des permutations signées aux couches cachées), la matrice de courbure dérivée devient un bloc diagonal avec des facteurs de Kronecker, ce qui est mathématiquement équivalent à Shampoo/Muon.
   • Cela fournit une explication fondée sur des principes de l'efficacité de ces optimiseurs : ils exploitent implicitement les symétries de l'espace des poids, trouvant un « point d'équilibre » entre la précision de l'approximation et le coût computationnel.
3. Contrôle de la complexité : Le cadre démontre que la complexité de l'approximation de la Hessienne est directement régie par la taille du groupe de symétrie choisi. Des groupes plus grands produisent des algèbres du commutant de plus faible dimension (moins de facteurs, calcul moins coûteux) mais des orbites plus grandes (potentiellement une précision locale moindre). Des groupes plus petits produisent plus de termes mais une meilleure fidélité locale.
4. Validation empirique :
   • Qualité de la Hessienne : Des expériences sur un MLP simple montrent que les approximations proposées (particulièrement la version simplifiée $\hat{H}_g$) atteignent une similitude cosinus élevée avec la véritable Hessienne le long de la direction du gradient négatif, surpassant les estimations de paramètres moyennées par orbite et égalant les performances de Shampoo sur cette métrique.
   • Équivalence des optimiseurs : Sur des auto-encodeurs MNIST et nanoGPT (un petit modèle de langage), l'optimiseur Symo produit des courbes d'entraînement presque identiques à celles de Muon, validant l'équivalence théorique sous des contraintes de groupe spécifiques. Les deux surpassent significativement Adam.

Résultats

• Compromis Précision vs Coût : Le papier explore empiriquement différents choix de groupes. Un « groupe plus large » (autorisant les permutations de toutes les couches, y compris l'entrée/sortie) surpasse toujours Adam mais est moins efficace qu'un « groupe plus petit » (restreignant les permutations aux couches internes et utilisant des approximations en blocs diagonaux), ce qui s'aligne sur la performance de Shampoo. Cela suggère que restreindre le groupe pour correspondre aux invariances spécifiques de l'architecture est crucial pour une performance optimale.
• Déficience de rang : Les auteurs notent que la matrice de moyenne d'orbite complète $S_w$ peut être déficiente en rang, ce qui déstabilise l'inversion. La simplification proposée ($S_w \approx I$) évite ce problème et produit empiriquement de meilleures estimations de courbure que la solution complète semi-définie positive.
• Scalabilité : La méthode parvient à passer à l'échelle sur des architectures Transformer (nanoGPT), démontant que l'approche sensible à la symétrie est applicable au-delà des simples MLP.

Signification et Revendications

Le papier prétend fournir un prisme théorique unificateur pour visualiser les optimiseurs de second ordre existants. En présentant Shampoo et Muon comme des instances spécifiques de l'approximation de la Hessienne basée sur la symétrie, ce travail fait passer ces méthodes de l'ingénierie heuristique à la dérivation mathématique fondée sur des principes.

Les auteurs affirment que leur cadre :

• Offre une recette systématique pour construire des approximations de courbure adaptées à des budgets computationnels spécifiques en sélectionnant les groupes de symétrie appropriés.
• Fournit une voie structurée allant des symétries du modèle à une courature structurée, capable d'inspirer de nouveaux optimiseurs pour différentes classes de modèles.
• Suggère des applications au-delà de l'optimisation, incluant l'estimation de l'incertitude, l'apprentissage continu et la compression de modèles, bien que celles-ci soient présentées comme des directions futures potentielles plutôt que comme des résultats démontrés.

Le papier reste modeste quant aux limites, reconnaissant que l'exactitude de l'expansion de Taylor autour de la moyenne d'orbite n'est pas entièrement bornée pour les poids entraînés, et que la garantie formelle de l'alignement des vecteurs propres entre le préconditionneur et la véritable Hessienne reste un problème ouvert. De plus, bien que le cadre automatise la dérivation des facteurs, la sélection du groupe optimal pour de nouvelles architectures arbitraires reste une tâche dépendante de l'utilisateur.