🤖 machine learning

Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalization

Pro-KLShampoo est un nouvel optimiseur qui améliore KL-Shampoo en exploitant la structure « pic et plat » des valeurs propres observée de ses préconditionneurs pour combiner un suivi spectral complet sur un sous-espace de faible dimension avec une orthogonalisation sur les directions restantes, ce qui permet d'obtenir des performances supérieures en termes de perte de validation, d'efficacité mémoire et de vitesse d'entraînement à travers plusieurs échelles de LLM.

Auteurs originaux : Ruotong Sun, Ermin Wei

Publié 2026-05-08

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Ruotong Sun, Ermin Wei

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'enseigner à un robot géant et complexe (un Modèle de Langage à Grande Échelle) à parler la langue humaine. Pour ce faire, vous lui montrez des millions d'exemples et lui permettez de faire des erreurs. Chaque fois qu'il se trompe, vous lui donnez une « petite pousse » (un gradient) pour corriger sa trajectoire.

Le problème est que ces petites poussées sont désordonnées. Parfois, le robot a besoin d'une toute petite pousse précise dans une direction, et d'une énorme bousculade dans une autre. Si vous le poussez simplement au hasard, il apprend lentement. Pour apprendre vite, vous avez besoin d'un « préconditionneur » intelligent — un outil qui remodèle ces poussées afin qu'elles soient parfaitement équilibrées et efficaces.

Deux outils populaires pour ce travail sont KL-Shampoo et Muon.

KL-Shampoo est comme un maître sculpteur. Il tente de tailler la forme parfaite pour chaque petite poussée individuelle. Il est très précis, mais nécessite beaucoup d'efforts lourds (puissance de calcul et mémoire) pour calculer la forme de chaque pièce du puzzle.
Muon est comme un gymnaste. Il ne tente pas de remodeler chaque pièce ; au lieu de cela, il redresse simplement le momentum des poussées, les faisant se déplacer d'une manière propre et orthogonale (à angle droit). Il est rapide et léger, mais il pourrait manquer certains détails subtils que le sculpteur attrape.

La Grande Découverte : Le Motif « Pic et Plaine »
Les auteurs de cet article ont examiné de près le travail du « sculpteur » (le préconditionneur KL-Shampoo) et ont remarqué un motif étrange et magnifique. Ils ont découvert que les « formes de poussée » ne sont pas aléatoires. Au contraire, elles ressemblent à un paysage de pic et plaine :

Le Pic : Quelques directions ont des valeurs énormes et dominantes (comme quelques hautes montagnes).
La Plaine : Le reste des directions a toutes à peu près la même hauteur (comme une vaste plaine plate).

Cela se produit à travers toutes les couches du cerveau du robot, de la première à la dernière. C'est comme si le robot ne se souciait vraiment que de quelques directions spécifiques, et qu'ailleurs, ce n'était que du « bruit » plat.

La Solution : Pro-KLShampoo
Les auteurs ont construit un nouvel optimiseur appelé Pro-KLShampoo (KL-Shampoo Projeté). Imaginez-le comme un outil hybride qui obtient le meilleur des deux mondes en utilisant la découverte du « pic et plaine ».

Voici comment cela fonctionne, en utilisant une analogie simple :

Le « Salon VIP » (Le Sous-espace) :
L'algorithme identifie les directions du « pic » — les quelques montagnes importantes. Il les place dans un « Salon VIP » spécial (un sous-espace suivi). À l'intérieur de ce salon, il utilise l'outil de sculpture lourd et précis (KL-Shampoo) pour obtenir la forme parfaite de ces quelques directions critiques. Il ne perd pas de temps sur le reste.
Le « Champ Ouvert » (Le Complément) :
Pour les directions « plates » (le reste du paysage), il arrête d'essayer de sculpter chaque galet. Au lieu de cela, il utilise une astuce ingénieuse appelée Orthogonalisation (empruntée à l'outil Muon).
- L'Astuce Magique : Les auteurs ont prouvé mathématiquement que si vous vous contentez de « redresser » (orthogonaliser) les poussées dans cette zone plate, cela récupère magiquement le résultat algébrique exact comme si vous aviez effectué la sculpture lourde là-bas. C'est comme réaliser que pour un champ plat, vous n'avez pas besoin de carte ; vous avez juste besoin de marcher en ligne droite, et vous finirez exactement au même endroit que si vous aviez calculé chaque coordonnée.

Pourquoi est-ce mieux ?

Vitesse : En ignorant la sculpture lourde pour les parties « plates » et en se contentant de les redresser, l'algorithme fonctionne beaucoup plus vite. Il économise beaucoup de temps sur l'« heure murale » (temps réel) de l'ordinateur.
Mémoire : Il n'a pas besoin de stocker les formes massives et complexes pour les parties plates. Il stocke uniquement les petites formes « VIP » et un seul nombre pour le reste. Cela économise beaucoup de mémoire informatique.
Performance : Dans les tests sur différentes tailles de robots (GPT-2 et LLaMA), Pro-KLShampoo a appris plus vite et atteint un taux d'erreur plus bas que le KL-Shampoo original, tout en utilisant moins de mémoire.

En Résumé
L'article déclare : « Nous avons découvert que les mathématiques complexes derrière l'entraînement de l'IA moderne ont un motif simple : quelques gros pics et une queue plate. Nous avons construit un nouvel outil qui traite les pics avec un soin extrême et la queue plate avec une astuce simple et rapide. Cette astuce fonctionne aussi bien que les mathématiques difficiles, mais est beaucoup plus rapide et moins coûteuse à exécuter. »

Le résultat est une méthode plus intelligente et plus rapide pour entraîner des modèles d'IA, qui économise du temps et des ressources informatiques sans sacrifier la qualité.

Résumé Technique : Pro-KLShampoo

Énoncé du Problème

Le pré-entraînement des modèles de langage de grande taille (LLM) modernes repose fortement sur des optimiseurs exploitant la structure matricielle des gradients. Deux fronts distincts ont émergé :

Préconditionnement factorisé de Kronecker explicite : Représenté par KL-Shampoo, qui estime le préconditionneur en minimisant la divergence de Kullback-Leibler (KL) entre la vraie covariance des gradients et une approximation structurée de Kronecker.
Orthogonalisation de la quantité de mouvement des gradients : Illustrée par Muon, qui orthogonalise la quantité de mouvement pour réaliser la descente la plus raide sous la norme spectrale.

Ces approches sont généralement développées de manière isolée. Bien que KL-Shampoo offre un fondement théorique solide via la minimisation de la KL, il engendre des coûts computationnels et mémoire élevés dus au maintien de facteurs de Kronecker complets. À l'inverse, les méthodes d'orthogonalisation comme Muon sont efficaces mais opèrent sous un cadre géométrique différent (descente la plus raide sous norme spectrale) sans la perspective d'estimation KL. L'article répond au besoin d'unifier ces approches pour atteindre l'efficacité statistique de KL-Shampoo avec l'efficacité computationnelle des méthodes d'orthogonalisation.

Méthodologie : Pro-KLShampoo

Les auteurs proposent Pro-KLShampoo (KL-Shampoo projeté avec blanchiment récupéré par orthogonalisation). L'innovation centrale découle d'une observation structurelle des préconditionneurs de KL-Shampoo : leurs spectres de valeurs propres présentent une forme « pic-et-plat ». Cela consiste en quelques valeurs propres dominantes (le « pic ») suivies d'une queue approximativement uniforme (le « plat »). Cette structure est observée empiriquement à travers les couches et les étapes d'entraînement, et est théoriquement exacte sous un modèle de gradient signal-plus-bruit de rang $\rho$ .

1. Restriction Structurelle

Pro-KLShampoo restreint l'un des facteurs de Kronecker de KL-Shampoo (spécifiquement le facteur côté droit $R$ pour les couches rectangulaires où $m \le n$ ) à une famille paramétrique alignée sur la structure pic-et-plat :

Sous-espace projeté (dimension $r$ ) : Le préconditionneur conserve la structure spectrale complète sur un sous-espace suivi engendré par $U \in \text{St}(n, r)$ .
Sous-espace complémentaire (dimensions $n-r$ ) : Les directions restantes sont effondrées en une seule valeur propre partagée $\mu_\perp$ .

Cette restriction réduit le stockage du facteur côté droit de $O(n^2)$ à $O(nr)$ et le coût de la décomposition QR de $O(n^3)$ à $O(r^3)$ .

2. Récupération par Orthogonalisation

Une restriction naïve à une queue scalaire entraîne une perte d'information si la queue n'est pas parfaitement uniforme. Pour récupérer la forme algébrique du préconditionneur complet KL-Shampoo sur le sous-espace complémentaire, Pro-KLShampoo applique une orthogonalisation (décomposition polaire) à la composante de gradient complémentaire.

La règle de mise à jour est :
$\Delta W = -\alpha_{kl} L^{-1/2} G U S^{-1/2} U^\top - c_a \text{polar}\left(\mu_\perp^{-1/2} L^{-1/2} G P_\perp\right)$
Où :

Le premier terme représente la mise à jour complète KL-Shampoo au sein du sous-espace suivi.
Le second terme applique l'orthogonalisation au complémentaire.
Les auteurs prouvent une identité (Éq. 12) montrant que cette orthogonalisation récupère la forme algébrique du blanchiment par direction du KL-Shampoo complet sur le complémentaire, reliant efficacement l'orthogonalisation de style Muon au cadre d'estimation KL.

3. Suivi de Sous-espace

Le sous-espace optimal $U$ est identifié comme l'espace propre de la matrice $\Phi_L = E[G^\top L^{-1} G]$ . Sous la conjecture pic-et-plat, le sous-espace des $r$ premières valeurs propres de $\Phi_L$ est le choix optimal. L'algorithme suit ce sous-espace en utilisant une seule étape d'itération de puissance (via décomposition QR) sur une estimation courante de $\Phi_L$ .

Contributions Clés

Observation Structurelle et Restriction Paramétrique : L'article identifie la structure de valeurs propres « pic-et-plat » dans les préconditionneurs de KL-Shampoo (robuste empiriquement et théoriquement exacte sous des modèles de gradients de faible rang) et l'exploite en restreignant l'objectif KL à une famille paramétrique.
Récupération Algébrique par Orthogonalisation : Il démontre que l'application de l'orthogonalisation au sous-espace complémentaire récupère la forme algébrique du préconditionneur complet KL-Shampoo. Cela fournit un pont théorique entre l'orthogonalisation de style Muon et l'estimation KL.
Garantie de Convergence : Les auteurs fournissent une garantie de convergence non convexe à un taux de $O(1/\sqrt{T})$ pour l'algorithme idéalisé sous une régularité de norme d'opérateur.
Performance Empirique : Une évaluation extensive sur les modèles GPT-2 (124M/350M) et LLaMA (134M/450M) montre que Pro-KLShampoo surpasse constamment KL-Shampoo en perte de validation, en mémoire de pointe par GPU et en temps réel pour atteindre des niveaux de perte équivalents, sur tous les rangs de sous-espace testés ( $r \in \{32, 64, 128\}$ ).

Résultats Expérimentaux

Les auteurs ont évalué Pro-KLShampoo contre AdamW, Muon et KL-Shampoo sur quatre configurations de pré-entraînement.

Perte de Validation : Pro-KLShampoo a atteint la perte de validation la plus faible à chaque rang testé sur toutes les échelles de modèles.
Efficacité Mémoire : L'utilisation de mémoire de pointe par GPU était inférieure à celle de KL-Shampoo et comparable à celle de Muon.
Temps Réel : Pro-KLShampoo a atteint des niveaux spécifiques de perte de validation plus rapidement que KL-Shampoo.
- Sur GPT-2, l'accélération était principalement due à une réduction du temps par étape (remplacement de la QR $O(n^3)$ par $O(r^3)$ ).
- Sur LLaMA, l'accélération était due à une convergence plus rapide (courbes de perte d'entraînement plus basses), en plus de l'efficacité par étape.
Études d'Ablation :
- Le retrait de l'orthogonalisation (variante Smok-Hop) a entraîné une dégradation des performances, en particulier à plus grande échelle (LLaMA 450M), confirmant la nécessité de l'orthogonalisation pour récupérer le blanchiment par direction.
- Ni les variantes uniquement sous-espace ni uniquement complémentaire n'ont égalé les performances du Pro-KLShampoo complet, indiquant que les deux composants sont essentiels.

Importance et Revendications

L'article revendique que Pro-KLShampoo unifie avec succès deux fronts d'optimisation précédemment isolés. En exploitant la structure intrinsèque « pic-et-plat » des covariances de gradients, il atteint l'efficacité statistique de KL-Shampoo tout en utilisant l'orthogonalisation pour maintenir l'efficacité computationnelle.

Les auteurs soulignent que le remplacement de la mise à jour de queue scalaire par l'orthogonalisation n'est pas une heuristique mais une récupération algébrique de la forme complète de KL-Shampoo. Cela permet à la méthode de conserver les avantages du préconditionnement basé sur KL (tels que la robustesse du pas de taille) tout en réduisant considérablement la surcharge computationnelle associée au maintien complet des facteurs de Kronecker. Le travail suggère que la structure spectrale de faible dimension est omniprésente dans l'optimisation des réseaux profonds et peut être efficacement exploitée pour concevoir des optimiseurs plus efficaces pour le pré-entraînement des LLM.

L'article reste modeste concernant les limites, notant que le rang $r$ est défini empiriquement, et qu'un travail futur est nécessaire pour l'étalonnage adaptatif par couche du poids de mélange $\alpha_{kl}$ et une analyse de convergence de la variante pratique impliquant des itérations de Newton-Schulz et une quantité de mouvement de Nesterov.