Résumé technique : Nora (Alignement Orthogonal Normalisé des Lignes)

Énoncé du problème

L'entraînement des grands modèles de langage (LLM) repose fortement sur des optimiseurs adaptatifs comme Adam, qui traitent les paramètres comme des vecteurs plats, ignorant ainsi les interactions structurelles riches au sein du paysage de perte. Bien que des optimiseurs basés sur des matrices comme Muon aient émergé pour y remédier en appliquant des mises à jour orthogonales et des itérations de Newton-Schulz, ils entraînent des coûts de calcul prohibitifs de $O(m^2n)$, limitant l'évolutivité.

À l'inverse, des approches simplifiées comme RMNP (Préconditionnement Normalisé par Momentum par Ligne) réduisent la complexité à $O(mn)$ en exploitant la dominance diagonale par blocs de lignes du Hessien du Transformer. Cependant, RMNP ne prend pas en compte l'invariance d'échelle inhérente aux réseaux de neurones modernes (due à BatchNorm, RMSNorm, LayerNorm). Il autorise des « mises à jour radiales » (mises à jour alignées avec le vecteur de poids) qui ne modifient pas la sortie fonctionnelle mais introduisent du bruit, déstabilisent le taux d'apprentissage effectif et provoquent des oscillations chaotiques dans les normes des poids.

Les méthodes existantes échouent généralement à unifier trois desiderata fondamentaux : l'efficacité (préconditionnement de type Muon), la stabilité (respect de l'invariance d'échelle) et la vitesse (surcoût de calcul minimal).

Méthodologie : L'optimiseur Nora

Les auteurs proposent Nora (Alignement Orthogonal Normalisé des Lignes), un optimiseur conçu pour satisfaire simultanément ces trois desiderata. La philosophie de conception centrale repose sur l'observation géométrique que, pour des paramètres invariants d'échelle, un apprentissage efficace se produit le long de directions angulaires (tangentes) plutôt que radiales.

Étapes algorithmiques

Pour une matrice de poids $w \in \mathbb{R}^{m \times n}$ et un momentum $v_t$ :

1. Projection Orthogonale par Ligne : Nora projette explicitement le momentum $v_t$ sur le complément orthogonal par ligne des poids $w$. Pour chaque ligne $i$, la composante parallèle à la ligne de poids est supprimée :
   $$v_{t,i:}^{\perp} = v_{t,i:} - \frac{\langle v_{t,i:}, w_{t,i:} \rangle}{\|w_{t,i:}\|^2} w_{t,i:}$$
   Cela garantit $\langle v_{t,i:}^{\perp}, w_{t,i:} \rangle = 0$, filtrant strictement le bruit radial et stabilisant les normes des poids.

2. Normalisation par Ligne (Préconditionnement) : En exploitant la dominance diagonale par blocs de lignes du Hessien du Transformer, Nora approxime le préconditionnement structuré de Muon. Au lieu d'itérations coûteuses de Newton-Schulz, il applique un préconditionneur diagonal au momentum projeté. Algébriquement, cela se simplifie en une simple normalisation $\ell_2$ par ligne :
   $$d_t = \text{RN}(v_t^{\perp}) \implies d_{t,i:} = \frac{v_{t,i:}^{\perp}}{\|v_{t,i:}^{\perp}\|_2}$$
   Cette étape atteint l'efficacité du préconditionnement de type Muon tout en maintenant l'orthogonalité établie lors de l'étape de projection.

3. Mise à jour : Les poids sont mis à jour en utilisant la direction normalisée $d_t$ avec un taux d'apprentissage $\eta_t$.

Fondements théoriques

• Évolutivité : En utilisant le cadre de la Paramétrisation de Mise à jour Maximale ($\mu$P), les auteurs prouvent que pour que Nora réalise un apprentissage de caractéristiques stable lorsque la largeur du réseau $n \to \infty$, le taux d'apprentissage doit évoluer comme $\eta_t \propto 1/\sqrt{n}$.
• Convergence : L'article fournit des garanties de convergence rigoureuses pour des paysages d'optimisation non convexes. Il établit que Nora atteint un point $\epsilon$-stationnaire avec une complexité de $O(m L \sigma^2 \Delta \epsilon^{-4})$ sous des hypothèses de régularité $(\infty, 2)$ adaptées.
• Implémentation : La logique centrale ne nécessite que deux lignes de code basé sur des matrices, en faisant un module plug-and-play compatible avec les frameworks existants.

Contributions clés

1. Conception d'un nouvel optimiseur : Introduction de Nora, qui unifie l'estimation efficace du préconditionnement, la stabilité via le respect de l'invariance d'échelle, et une vitesse de calcul élevée.
2. Rigueur théorique : Preuve que Nora est un optimiseur évolutif avec des lois d'échelle dérivées ($\eta \propto 1/\sqrt{n}$) et des garanties de convergence non convexe.
3. Validation empirique : Des expériences extensives entraînant des modèles de style LLaMA (60M et 135M paramètres) démontrent que Nora surpasse Muon, Mano et RMNP tant en vitesse de convergence qu'en temps d'exécution réel.

Résultats expérimentaux

Les auteurs ont évalué Nora sur des tâches de modélisation linguistique autorégressive à travers divers taux d'apprentissage et échelles de modèles :

• Performance : Sur le modèle de 135M, Nora a obtenu la perte de validation la plus faible (3,079) et la perplexité la plus basse (21,74) parmi tous les optimiseurs comparés (Muon : 3,142/23,17 ; Mano : 3,097/22,13 ; RMNP : 3,112/22,46).
• Dynamique d'entraînement : Bien que RMNP ait montré de solides performances au stade précoce, il a atteint un plateau. Nora a commencé avec une perte initiale plus élevée mais a continué à s'améliorer tardivement dans l'entraînement, dépassant les autres méthodes pour atteindre les meilleures métriques finales.
• Efficacité : Les tests de performance sur des formes de matrices représentatives (par exemple, matrices MLP à l'échelle de 1 milliard) ont montré que la normalisation par ligne de Nora est environ 70 fois plus rapide que l'orthogonalisation de Newton-Schulz en 5 étapes utilisée par Muon.
• Ablation : Les comparaisons avec Mano (une méthode similaire de projection par ligne/colonne) ont confirmé que la supériorité de Nora découle de sa dépendance spécifique à la dominance diagonale par blocs de lignes et de son architecture rationalisée, plutôt que d'un réglage des hyperparamètres ou de paramètres de décroissance des poids.

Signification et revendications

L'article affirme que Nora démontre que l'efficacité, la stabilité et la vitesse peuvent être unifiées par un simple principe géométrique par ligne. En stabilisant rigoureusement les normes de poids et les vitesses angulaires tout en approximatant le préconditionnement structuré, Nora offre une « direction pratique et fondée pour la conception d'optimiseurs évolutifs ». Les auteurs soulignent que Nora atteint une efficacité des données à la pointe de l'état de l'art sans le surcoût de calcul prohibitif de l'orthogonalisation complète de matrices, ce qui en fait un candidat très prometteur pour l'entraînement de LLM à grande échelle.