Résumé Technique : Couche PC — Préconditionnement par Poids Polynomial pour l'Amélioration du Pré-entraînement des LLM

1. Énoncé du Problème

L'entraînement de modèles de langage de grande taille (LLM) à grande échelle repose largement sur des techniques de normalisation pour stabiliser l'optimisation. Bien que des méthodes comme RMSNorm et QK-Norm contrôlent les sorties des couches, et que la normalisation des poids (WN) ou la normalisation spectrale (SN) agissent dans l'espace des poids, il est de plus en plus reconnu que la géométrie des matrices de poids — spécifiquement leur spectre de valeurs singulières — est un axe de conception critique, bien qu'encore sous-exploré.

Le problème central abordé est le conditionnement des matrices de poids. Dans les réseaux profonds, le rapport entre les plus grandes et les plus petites valeurs singulières (le nombre de conditionnement) dicte la stabilité de la propagation du signal. Un effondrement de la queue inférieure des valeurs singulières dégrade la propagation du signal à travers la profondeur, rendant le réseau plus difficile à optimiser. Si l'imposition d'une orthogonalité exacte (réduisant toutes les valeurs singulières à 1) préserve la norme, elle supprime la nécessaire anisotropie spectrale, ce qui peut affaiblir le pouvoir de représentation du modèle. Le défi consiste à concevoir un mécanisme qui préconditionne doucement le spectre des poids — en réduisant l'écart des valeurs singulières pour améliorer le conditionnement sans les forcer toutes à l'unité, préservant ainsi l'expressivité du modèle. De plus, un tel mécanisme doit être efficace sur le plan computationnel, évitant les opérations coûteuses comme la Décomposition en Valeurs Singulières (SVD) explicite pendant l'entraînement.

2. Méthodologie : La Couche PC

Les auteurs proposent la couche de Préconditionnement (PC), un module intégré qui remodèle le spectre des valeurs singulières de matrices de poids sélectionnées via un préconditionnement polynomial.

Mécanisme Central

Au lieu de calculer une SVD explicite ($W = U \Sigma V^\top$) pour modifier directement les valeurs singulières, la couche PC applique un polynôme de matrice de bas degré $g(W)$ à la matrice de poids $W$.

• Cartographie Spectrale : Si $W$ possède des valeurs singulières $\sigma_i$, la matrice transformée $g(W)$ possède des valeurs singulières $|g(\sigma_i)|$. En choisissant un polynôme $g$ qui amplifie les petites valeurs singulières et sature les grandes, la méthode réduit efficacement le nombre de conditionnement.
• Implémentation pour les Matrices Rectangulaires : Puisque les matrices de poids sont rectangulaires, le polynôme est appliqué via la matrice de Gram. Pour un poids $W$, la transformation est définie par $g(W) = p(WW^\top)W$ (ou $Wp(W^\top W)$), où $p$ est un polynôme de bas degré. Cela préserve la forme de la matrice et ne nécessite que des multiplications de matrices répétées.
• Conception du Polynôme : Le polynôme $g(\sigma) = p(\sigma^2)\sigma$ est ajusté par la méthode des moindres carrés pondérés pour approximer une fonction cible linéaire par morceaux $PL_b(\sigma) = \min(\sigma/b, 1)$. Cette cible amplifie les valeurs singulières inférieures à un seuil $b$ et plafonne celles supérieures à $b$ à 1. Le « niveau PC » ($k$) contrôle le degré du polynôme et l'agressivité du façonnage (un $b$ plus petit implique un façonnage plus fort).

Procédure d'Entraînement (Algorithme 1)

La couche PC opère comme une reparamétrisation pendant l'entraînement :

1. Normalisation Spectrale : Le poids $W$ est normalisé par une estimation de sa norme spectrale $s(W) \approx \|W\|_2$, calculée via une itération de puissance en flux tendu (10 étapes avec un démarrage à chaud à partir des étapes précédentes) pour éviter une SVD complète. Cela produit $\tilde{W} = W/s(W)$, confinant les valeurs singulières approximativement dans $[0, 1]$.
2. Préconditionnement Polynomial : Le polynôme $g$ est appliqué à $\tilde{W}$ pour remodeler le spectre.
3. Récupération de la Norme et Mise à l'Échelle : Le résultat est remis à l'échelle par la norme spectrale estimée $[s(W)]_{\text{stop-grad}}$ pour restaurer la magnitude, suivi d'un scalaire apprenable $\gamma$ (initialisé à 1) pour permettre un ajustement adaptatif des métriques de propagation du signal.
4. Inférence : Après l'entraînement, les poids préconditionnés sont fusionnés de nouveau dans l'architecture originale. Aucun calcul polynomial n'a lieu lors de l'inférence, ce qui entraîne un coût d'inférence nul.

Blocs Cibles

Dans les architectures de type Llama, la couche PC est appliquée aux projections FFN ($W_{\text{gate}}, W_{\text{up}}, W_{\text{down}}$) et à la projection de sortie de l'attention ($W_O$).

3. Justification Théorique

L'article fournit un fondement théorique au principe de contrôle du spectre en utilisant des réseaux linéaires profonds.

• Garantie de Convergence : Les auteurs prouvent que si les matrices de poids dans un réseau linéaire profond maintiennent des nombres de conditionnement uniformément bornés (c'est-à-dire que les valeurs singulières sont éloignées de 0 et de $\infty$), l'algorithme de descente de gradient converge vers un minimum global à un taux géométrique.
• Complexité d'Itération : Le taux de convergence dépend de la borne du nombre de conditionnement global $\kappa_R$. Plus précisément, la complexité d'itération pour atteindre une erreur $\epsilon$ évolue en $O(\kappa_R^{2L} \log(1/\epsilon))$. Cela implique qu'un meilleur conditionnement des poids (un $\kappa_R$ plus petit) se traduit directement par une convergence plus rapide, motivant le contrôle du spectre des poids pendant l'entraînement.

4. Résultats Expérimentaux

La couche PC a été évaluée sur des modèles Llama-271M et Llama-1B pré-entraînés sur le jeu de données FineWeb en utilisant deux optimiseurs : AdamW et Muon.

Performance d'Optimisation

• AdamW : PC a considérablement amélioré l'efficacité des jetons (tokens).
  • Llama-271M : A atteint la perte de référence avec 39 % de jetons en moins (accélération de 1,63×).
  • Llama-1B : A atteint la perte de référence avec 50 % de jetons en moins (accélération de 2×).
• Muon : PC a également apporté des gains constants, bien que plus faibles qu'avec AdamW (probablement parce que Muon exerce déjà un contrôle spectral implicite).
  • Llama-271M : 1,07× d'accélération.
  • Llama-1B : 1,13× d'accélération.
• Tâches Aval : Sur les évaluations zero-shot (LAMBADA, HellaSwag, etc.), PC a amélioré la précision moyenne de 0,0206 point (AdamW) et de 0,0125 point (Muon) sur le modèle 1B.

Conditionnement Spectral

• Nombre de Conditionnement : Le Nombre de Conditionnement Modifié Global (GMCN) a diminué de manière significative. Pour le modèle 1B sous AdamW, le GMCN global est passé d'environ 42,4 (référence) à 25,0 (PC), soit une réduction d'environ 41 %.
• Distribution des Valeurs Singulières : Les histogrammes montrent que PC déplace la masse du spectre des valeurs singulières loin de l'extrémité inférieure (proche de zéro) vers une plage intermédiaire plus régulière, rétrécissant l'étalement relatif sans faire s'effondrer le spectre en une valeur unique.

Études d'Ablation

• Niveau PC ($k$) : Sous AdamW, les degrés plus élevés ($k=4$) étaient optimaux. Sous Muon, les degrés plus faibles ($k=2$) étaient optimaux, suggérant que le contrôle implicite de Muon nécessite un façonnage externe moins agressif.
• Récupération de la Norme : Essentielle pour la performance ; sans la restauration de l'échelle spectrale après le façonnage polynomial, la performance chutait en dessous de la référence.
• $\gamma$ Apprenable : Crucial pour stabiliser la dynamique d'entraînement (réduction des pics de l'RMS des activations), mais avait un effet mineur sur la perte de validation finale par rapport à la récupération de la norme.

Coût Computationnel

• FLOPs : Le surcoût est minimal (~0,39 % pour AdamW, ~0,24 % pour Muon sur Llama-1B) grâce à l'utilisation de polynômes de bas degré et de la méthode de Horner efficace.
• Mémoire : La mémoire de pointe a augmenté d'environ 9,56 % (AdamW) et 8,73 % (Muon) en raison des tenseurs intermédiaires pour la rétropropagation, mais cela reste gérable.

5. Signification et Revendications

L'article affirme que la couche PC offre une solution légère et prête à l'emploi pour améliorer le pré-entraînement des LLM en contrôlant explicitement le conditionnement du spectre des poids.

• Insight Théorique : Elle comble le fossé entre la nécessité théorique de nombres de conditionnement bornés pour la convergence géométrique et la conception d'algorithmes pratiques.
• Efficacité Pratique : Contrairement aux méthodes nécessitant une SVD explicite ou des contraintes complexes, PC n'utilise que des multiplications de matrices, ce qui la rend scalable et compatible avec les optimiseurs et architectures existants (AdamW, Muon).
• Coût d'Inférence Nul : En fusionnant les poids préconditionnés après l'entraînement, la méthode n'ajoute aucun coût supplémentaire lors de l'inférence.
• Façonnage Doux : Les auteurs soulignent que leur approche réalise un « façonnage doux du spectre » plutôt que de forcer une quasi-orthogonalité, préservant la capacité du modèle à apprendre des transformations anisotropes nécessaires à l'expressivité.

Ce travail se présente comme une étape vers le traitement de la géométrie des poids comme un levier de contrôle direct pour la propagation du signal, offrant des améliorations constantes à travers différentes échelles de modèles et optimiseurs.