Résumé Technique : Stabiliser le Préentraînement Natif à Faible Rang des LLM

1. Énoncé du Problème

Les modèles de fondation ont connu un succès significatif, mais l'augmentation de leur nombre de paramètres crée des goulots d'étranglement substantiels en termes de calcul et de mémoire. Bien que la factorisation de rang faible (exprimant une matrice de poids $W = AB^\top$) se soit avérée efficace pour le réglage fin (par exemple, LoRA), l'application de cette technique au préentraînement à partir de zéro est restée insaisissable.

Les tentatives antérieures visant à entraîner des modèles nativement dans un régime de rang faible (sans poids de rang plein auxiliaires) souffrent d'une instabilité sévère. Le problème central identifié dans ce travail est la croissance incontrôlée de la norme spectrale (la plus grande valeur singulière) de la mise à jour de la matrice de poids. Contrairement à l'entraînement dense, où les optimiseurs comme AdamW maintiennent des normes spectrales bornées, la factorisation de rang faible permet des représentations équivalentes infinies ($W = (\lambda A)(\frac{1}{\lambda} B)$). Cette invariance d'échelle permet au facteur d'échelle $\lambda$ de croître sans limite, entraînant l'explosion des activations et la divergence de l'entraînement. Les méthodes existantes (par exemple, Wei et al., 2024a ; Huh et al., 2024) contournent cela en maintenant des poids de rang plein auxiliaires ou en utilisant des architectures hybrides, ce qui engendre un surcoût de mémoire important ou ne parvient pas à réaliser un véritable entraînement factorisé de bout en bout.

2. Méthodologie : Spectron

Les auteurs proposent Spectron, une méthode permettant un préentraînement natif à faible rang, stable et de bout en bout, sans composants de rang plein auxiliaires. Spectron traite l'instabilité spectrale par deux mécanismes combinés : la Renormalisation Spectrale et l'Orthogonalisation du Gradient.

Mécanisme Central

La méthode vise à borner la norme spectrale de la mise à jour de poids composée $\Delta W$ à un rayon prescrit $\eta$ (dérivé du taux d'apprentissage).

1. Orthogonalisation du Gradient : La méthode orthogonalise les mises à jour de momentum pour les facteurs $A$ et $B$ en utilisant des itérations de Newton-Schulz. Cela garantit que les mises à jour se situent dans un rayon de contrainte locale, empêissant une croissance directionnelle arbitraire.
2. Renormalisation Spectrale Adaptative : Pour satisfaire la contrainte globale $\|\Delta W\|_2 \leq \eta$, la méthode ajuste dynamiquement le taux d'apprentissage des facteurs en fonction des normes spectrales actuelles de $A$ et $B$.
   La règle de mise à jour pour les facteurs est :
   $$\Delta A \leftarrow \frac{\eta}{\|A\|_2 + \|B\|_2 + 1} \cdot \text{Ortho}(M_A)$$
   $$\Delta B \leftarrow \frac{\eta}{\|A\|_2 + \|B\|_2 + 1} \cdot \text{Ortho}(M_B)$$
   Ici, $\|A\|_2$ et $\|B\|_2$ sont estimés efficacement à l'aide d'une seule étape d'itération de puissance.

Efficacité de Mise en Œuvre

• Orthogonalisation : Implémentée via 5 itérations de Newton-Schulz.
• Estimation Spectrale : Implémentée via 1 étape d'itération de puissance.
• Surcoût : Le surcoût de calcul total est inférieur à 1 % des FLOPs, une amélioration significative par rapport aux méthodes d'entraînement "auto-guidées" qui nécessitent environ 25 % de FLOPs supplémentaires pour les poids denses auxiliaires.

3. Contributions Clés

1. Algorithme Spectron : Une méthode de mise à jour de facteur à faible rang adaptative qui combine l'orthogonalisation avec la renormalisation spectrale. Elle borne les normes spectrales des mises à jour de poids résultantes, permettant un entraînement stable à partir d'une initialisation aléatoire sans composants denses auxiliaires.
2. Validation Empirique : Démontre que les modèles factorisés entraînés avec Spectron atteignent une parité de performance, voire dépassent, les bases de référence denses sur trois échelles de modèles (94M, 297M, 454M de paramètres) sur le jeu de données FineWeb.
3. Lois d'Échelle Optimales en Calcul (Compute-Optimal Scaling Laws) : À travers l'analyse IsoFLOP, les auteurs ont dérivé des lois d'échelle spécifiques pour les transformeurs à faible rang natifs :
   • Taille de modèle optimale : $N_{opt} \propto C^{0.479}$
   • Tokens d'entraînement optimaux : $D_{opt} \propto C^{0.521}$
     Ces exposants diffèrent légèrement des lois de Chinchilla ($N \propto C^{0.49}$, $D \propto C^{0.51}$), indiquant que les modèles à faible rang optimales en calcul favorisent des architectures plus petites entraînées sur des ensembles de données proportionnellement plus vastes.

4. Résultats

• Stabilité : Spectron élimine les pics de perte et la divergence observés dans l'entraînement à faible rang naïf (AdamW) et l'entraînement auto-guidé. Il maintient des normes spectrales bornées tout au long de l'entraînement.
• Performance :
  • Sur le modèle de 454 millions de paramètres, Spectron a atteint une perplexité de validation de 12,11, surpassant la base de référence auto-guidée (13,70) et l'AdamW naïf (14,57).
  • Lorsqu'ils sont entraînés pour un nombre égal de FLOPs, un modèle factorisé de 454M égale la perte de validation d'un modèle dense de 780M, récupérant ainsi la performance d'un modèle possédant environ 42 % de paramètres supplémentaires.
  • Les tâches en aval (HellaSwag, PIQA, ARC-easy) montrent des améliorations constantes par rapport aux bases de référence et aux modèles denses de comptes de paramètres équivalents.
• Efficacité d'Inférence : En raison du nombre réduit de paramètres requis pour atteindre la même perplexité, les modèles Spectron offrent des réductions substantielles du coût d'inférence (estimées jusqu'à 50 % par rapport aux transformeurs denses optimaux selon Chinchilla sous les budgets de calcul modernes).

5. Signification et Revendications

Le papier affirme que l'entraînement à faible rang natif est une alternative viable et stable au préentraînement dense lorsque des contraintes spectrales appropriées sont appliquées.

• Intuition Scientifique : Ce travail révèle que les représentations de rang plein ne sont pas strictement nécessaires lors de l'optimisation ; l'instabilité observée précédemment était le résultat d'une croissance spectrale incontrôlée plutôt que d'une limitation inhérente aux structures de rang faible.
• Impact Pratique : En éliminant le besoin de poids de rang plein auxiliaires, Spectron démocratise le développement de modèles de fondation, réduisant potentiellement la consommation de mémoire et les exigences matérielles.
• Comportement de Mise à l'Échelle : Les lois d'échelle dérivées suggèrent que, pour un budget de calcul fixe, la stratégie optimale pour les modèles à faible rang est d'entraîner des modèles plus petits sur plus de données, produisant des modèles qui sont intrinsèquement plus efficaces en inférence.

Limites Reconnues :

• Les expériences ont été menées sur des modèles allant jusqu'à 1,5 milliard de paramètres, ce qui est inférieur aux régimes de trillions de tokens de la frontière actuelle.
• L'implémentation actuelle souffre d'un écart de temps d'exécution (wall-clock time) par rapport à l'entraînement dense en raison d'opérations de noyaux (kernels) non fusionnées, bien que les auteurs notent que cela peut être résolu avec des noyaux fusionnés personnalisés.
• La méthode se concentre sur les LLM, bien que les principes puissent s'étendre aux architectures multimodales.