Résumé Technique : Zeta – Blanchiment Dual pour l'Optimisation de Matrice via un Préconditionnement Adaptatif aux Coordonnées

Énoncé du Problème

L'entraînement des réseaux de neurones à grande échelle s'est de plus en plus tourné vers des optimiseurs conscients des matrices qui exploitent les propriétés structurelles des paramètres de poids, dépassant l'adaptation élément par élément. Un exemple de premier plan est Muon, qui impose des contraintes d'orthogonalité sur les mises à jour de poids via un blanchiment spectral (spécifiquement, l'itération de Newton–Schulz) afin d'égaliser les valeurs singulières et d'empêcher qu'une seule direction ne domine la mise à jour.

Cependant, cet article identifie une vulnérabilité critique et sous-estimée de Muon et des méthodes spectrales similaires : l'itération de Newton–Schulz est hautement sensible au conditionnement de l'entrée. L'itération fonctionne bien lorsque les valeurs singulières sont bien distribuées, mais les matrices de momentum brutes dans l'entraînement des réseaux profonds présentent souvent une hétérogénéité d'échelle par coordonnée sévère. Dans ces matrices, quelques coordonnées possèdent des magnitudes disproportionnellement grandes, faussant le spectre singulier. Lorsque de telles matrices non conditionnées entrent dans l'étape de Newton–Schulz, le processus d'orthogonalisation se dégrade, conduisant à des directions de mise à jour déformées qui ne reflètent plus fidèlement la géométrie du gradient. Les auteurs soutiennent que le blanchiment spectral nécessite implicitement une entrée normalisée en échelle, une condition que le momentum brut viole systématiquement.

Méthodologie : L'Optimiseur Zeta

Pour combler cette lacune structurelle, les auteurs proposent Zeta, un optimiseur à blanchiment dual qui organise le blanchiment des coordonnées et le blanchiment spectral en un pipeline strictement ordonné et non ajustable. L'idée centrale est que le blanchiment spectral doit s'appuyer sur une stabilisation statistique ; le blanchiment des coordonnées établit l'isotropie statistique nécessaire pour que le blanchiment spectral fonctionne de manière fiable.

Le pipeline Zeta opère comme suit :

1. Blanchiment des Coordonnées (Préconditionnement) :
   Avant d'entrer dans l'étape spectrale, Zeta applique une normalisation par coordonnée à la matrice de momentum. Il maintient des moyennes mobiles exponentielles du premier moment ($M_t$) et du second moment ($V_t$), de manière similaire à Adam. La matrice de momentum est ensuite normalisée élément par élément par le second moment courant :
   $$\tilde{G}_t = \frac{M_t}{\sqrt{V_t} + \epsilon}$$
   Cette étape homogénéise la disparité d'échelle entre les coordonnées, transformant la matrice en un état statistiquement isotrope. Théoriquement, cela agit comme un préconditionneur diagonal approximatif, traitant le déséquilibre d'échelle qui fausserait autrement le spectre singulier.

2. Blanchiment Spectral (Orthogonalisation) :
   La matrice homogénéisée $\tilde{G}_t$ est ensuite injectée dans l'itération de Newton–Schulz. Parce que l'entrée est désormais exempte de disparités extrêmes d'échelle par entrée, son spectre singulier est mieux conditionné. L'itération converge vers une approximation orthogonale fidèle, égalisant les valeurs singulières pour produire une direction de mise à jour équilibrée.

3. Construction de la Mise à Jour :
   La matrice orthogonale résultante $U_t$ est mise à l'échelle à l'aide d'un facteur normalisé par RMS pour garantir que la magnitude de la mise à jour est compatible avec les configurations d'optimiseurs standards (par exemple, les branches AdamW) sans altérer la direction spectrale.

4. Exécution par Double Voie :
   Pour les paramètres de type matrice (par exemple, les projections d'attention et de feed-forward), le pipeline dual complet est appliqué. Pour les paramètres non matriciels (par exemple, les biais, les échelles de LayerNorm), l'étape spectrale est sautée, et l'optimiseur revient à une mise à jour de type AdamW.

Contributions Théoriques

L'article fournit trois garanties théoriques clés soutenant la conception :

• Blanchiment Approximatif : L'étape de normalisation des coordonnées est prouvée être une approximation diagonale efficace de la transformée de blanchiment idéale (Théorème 1), passant d'une mise à l'échelle heuristique à une étape de préconditionnement fondée sur des principes.
• Réduction de l'Erreur d'Orthogonalisation : Les auteurs prouvent que le blanchiment des coordonnées améliore strictement le nombre de conditionnement de l'entrée de l'itération de Newton–Schulz. Par conséquent, l'erreur d'orthogonalisation de l'étape spectrale est strictement réduite par rapport à une opération sur le momentum brut (Théorème 2).
• Préservation de l'Échelle : La mise à l'échelle RMS finale préserve la direction spectrale de la mise à jour tout en synchronisant sa magnitude globale avec le reste du réseau (Théorème 3).

Résultats Empiriques

Zeta a été évalué sur des tâches de modélisation de langage (0,6B à 8B paramètres), des architectures de mélange d'experts (MoE) et des tâches de vision, égalant ou surpassant systématiquement les bases solides (AdamW, Muon et AdaMuon).

• Vitesse de Convergence : Dans les tâches de pré-entraînement, Zeta a démontré une convergence plus rapide que Muon et AdamW. Par exemple, sur un modèle de 1,7 milliard de paramètres, Zeta a atteint une accélération de 1,67× (par rapport à AdamW) pour atteindre le même niveau de perte, surpassant Muon (1,52×) et AdaMuon (1,61×). Sur un modèle de 8 milliards de paramètres, Zeta a obtenu une accélération de 1,25× contre 1,19× pour Muon.
• Efficacité : Le coût de temps par itération de Zeta est resté comparable à celui d'AdamW (par exemple, 25,47s contre 24,46s pour le modèle de 8B sur 16 GPU), indiquant que le processus de blanchiment dual n'entraîne pas de pénalités de temps d'exécution significatives.
• Performance MoE : Dans les architectures MoE, où l'activation parcimonieuse des experts crée une hétérogénéité importante des gradients, Zeta a maintenu un avantage clair, convergeant vers une perte plus faible avec moins de fluctuations que ses concurrents.
• Tâches Aval : Sur 12 benchmarks de compréhension du langage (ex: MMLU, GSM8k), les modèles entraînés avec Zeta ont atteint la précision moyenne la plus élevée sur les échelles 1,7B et 8B. De même, sur les tâches de vision CIFAR-100, Zeta a obtenu la précision de validation Top-1 la plus élevée pour ViT-Tiny et ViT-Base.
• Robustesse des Hyperparamètres : Une recherche par grille sur les coefficients de momentum ($\beta_1, \beta_2$) a montré que Zeta est robuste aux choix d'hyperparamètres, la perte d'entraînement finale variant de moins de 1,2 % sur la plage testée.

Signification et Revendications

L'article affirme que Zeta résout un décalage fondamental dans les optimiseurs actuels conscients des matrices : l'application d'opérations géométriques puissantes (blanchiment spectral) à des signaux présentant des échelles par coordonnées incontrôlées. En vérifiant la prévalence du déséquilibre d'échelle intra-matrice via des tests d'uniformité du chi carré et en démontrant que le blanchiment des coordonnées le corrige efficacement, les auteurs établissent que résoudre le déséquilibre d'échelle avant l'orthogonalisation est un prérequis pour des méthodes spectrales stables.

La signification de Zeta réside dans son unification de deux paradigmes d'optimisation auparavant concurrents (mise à l'échelle adaptative par coordonnée et orthogonalisation spectrale) en un pipeline unique et logiquement nécessaire. L'ordre des opérations n'est pas un hyperparamètre ajustable mais une dépendance mathématique : le blanchiment des coordonnées doit précéder le blanchiment spectral pour garantir que l'entrée est bien conditionnée. Cette approche conduit à une convergence plus rapide et une meilleure généralisation sans sacrifier les avantages géométriques de l'optimisation au niveau de la matrice.