Résumé Technique : Apprentissage dans un Espace de Poids Courbe (SymExpLin)

1. Énoncé du Problème

L'entraînement standard des réseaux de neurones repose sur des mises à jour additives dans un espace de paramètres linéaire. Cependant, de nombreuses opérations internes des réseaux sont intrinsèquement multiplicatives (par exemple, le passage à l'échelle des amplitudes de caractéristiques, les gains de normalisation ou les valeurs singulières). Dans un espace linéaire, des changements relatifs d'une même magnitude (par exemple, doubler vs diviser par deux) nécessitent des distances d'optimisation inégales. Par exemple, passer d'un poids de 1,0 à 2,0 nécessite un pas de taille 1,0, tandis que passer de 1,0 à 0,5 nécessite un pas de 0,5.

Ce décalage crée des défis d'optimisation :

• Distributions à queue lourde : Les modèles préentraînés présentent souvent des distributions de poids à queue lourde où une petite fraction de paramètres croît de manière significative par rapport à leur échelle d'initialisation (par exemple, 30 fois plus grande), tandis que la plupart restent proches de l'initialisation.
• Invariance de l'optimiseur : Les optimiseurs adaptatifs comme Adam normalisent les mises à jour par coordonnée, rendant les tailles de pas invariantes à la magnitude du paramètre. Par conséquent, un petit paramètre (ex. 0,004) et un grand paramètre (ex. 0,4) reçoivent des mises à jour absolues similaires. Cela entraîne une différence de perturbation relative de 100×, ce qui peut sur-perturber les petits paramètres tout en sous-perturbant les grands paramètres critiques.
• Échelle Statique vs Dynamique : Les solutions existantes comme la Paramétrisation de Mise à l'Échelle Maximale (µP) traitent les décalages d'échelle de manière statique en fonction de la largeur du tenseur lors de l'initialisation, mais ne s'adaptent pas à la croissance dynamique de chaque paramètre individuel pendant l'entraînement.

2. Méthodologie : SymExpLin (SEL)

Les auteurs proposent SymExpLin (SEL), une reparamétrisation non linéaire des poids qui décompose les poids effectifs en une voie exponentielle symétrique et une voie linéaire. Cela crée une géométrie d'espace de paramètres courbe où les mises à jour additives dans l'espace brut se traduisent par des changements proportionnels à la magnitude dans l'espace des poids effectifs.

2.1 La Transformation

Le poids effectif $w_{eff}$ est calculé à partir d'un poids brut $w$ en utilisant les composants suivants :

1. Voie Symétrique-Exponentielle ($E_\theta$) : Une fonction apprenable et préservant le signe basée sur l'exponentielle symétrique (symexp). Pour de petits $|w|$, elle se comporte de manière linéaire ; pour de grands $|w|$, elle croît de manière exponentielle. Cela permet aux grands pas bruts de produire des déplacements effectifs plus importants.
   • $E_\theta(|w|) = e_w (\exp(\beta |w| \cdot m - n) - \exp(-n))$
   • Paramètres : $\beta$ (courbure), $m$ (multiplicateur interne), $n$ (décalage de la pente à l'origine), $e_w$ (échelle de sortie).
2. Voie Linéaire ($L_\theta$) : Une voie de type identité $L_\theta(w) = w \cdot l_w$ qui fournit une route directe pour le gradient, supposée stabiliser l'optimisation.
3. Enveloppement du Signe : Le signe du poids brut est appliqué à la somme des voies : $w_{eff} = \text{sign}(w) \cdot (E_\theta(|w|) + L_\theta(w)) + \Delta_\phi$.
4. Résiduel ($\Delta_\phi$) : Une correction additive optionnelle de faible rang (de type LoRA).

2.2 Choix de Conception Clés

• Initialisation Décalée (Mismatched) : Les auteurs identifient une asymétrie bénéfique où les poids bruts sont initialisés en inversant une règle de combinaison « congruente » (où les parties linéaire et exponentielle sont additionnées à l'extérieur de l'enveloppement du signe), mais où le passage vers l'avant utilise une règle « décalée » (où les deux sont enveloppées à l'intérieur du signe). Cela fait en sorte que les poids négatifs sont initialisés à une magnitude plus petite que les poids positifs, agissant comme un mécanisme de rupture de symétrie précoce.
• Paramètres d'Échelle Structurés : Les échelles apprenables ($e_w, l_w, m, n$) sont paramétrées à l'aide de motifs multiplicatifs ligne-colonne (vecteurs de rang 1) pour équilibrer l'expressivité et le coût de calcul.
• Recul de Taux d'Apprentissage (Annealing) : Un calendrier ciblé est utilisé pour les paramètres d'échelle. L'échelle externe $e_w$ commence avec un taux d'apprentissage élevé pour s'engager rapidement sur un signe, tandis que le multiplicateur interne $m$ commence bas et augmente progressivement (annealing) pour affiner la courbure une fois que l'échelle est stabilisée.
• Calibration de la Courbure : Le paramètre de courbure $\beta$ est défini de manière heuristique en fonction de la largeur du modèle pour garantir que le régime exponentiel s'active près des magnitudes typiques d'initialisation (qui diminuent à mesure que la largeur augmente).

3. Contributions Clés

• SymExpLin (SEL) : Une reparamétrisation des poids applicable à toute matrice ou biais apprenable, décomposant les poids en composantes symétriques-exponentielles et linéaires avec des échelles structurées indépendantes et apprenables.
• Initialisation Décalée : La découverte que l'utilisation de règles de combinaison différentes pour l'inversion de l'initialisation par rapport au passage vers l'avant crée un biais asymétrique bénéfique qui améliore l'optimisation précoce.
• Échelle Structurée & Annealing : Introduction de paramètres d'échelle multiplicatifs ligne-colonne et d'un calendrier spécifique de recul du taux d'apprentissage pour gérer les contrôles exponentiels redondants.
• Praticité : La méthode est facile à implémenter. La transformation d'entraînement est élément par élément, et après l'entraînement, les poids effectifs sont précalculés et intégrés dans des couches linéaires standards, ce qui résulte en un coût d'inférence nul.
• Complémentarité : SEL opère sur un axe dynamique par paramètre, complétant les méthodes statiques par tenseur comme µP.

4. Résultats Expérimentaux

La méthode a été évaluée sur la modélisation de langage autorégressive en utilisant OpenWebText sur neuf configurations de largeur $\{1024, 2048, 3072\}$ et de profondeur $\{12, 24, 36\}$.

• Descente de la Perte : SEL a systématiquement atteint une perte de validation plus basse que la référence (Adam avec paramétrage standard) sur toutes les configurations.
• Efficacité des Étapes : SEL a atteint la perte de validation équivalente en 1,32× à 1,49× moins d'étapes d'entraînement. Les gains ont été les plus prononcés dans les modèles de plus grande largeur (ex. accélération de 1,49× pour 3072×24).
• Accélération Temps Réel (Wall-Clock) : Malgré un surcoût par étape de 1,7 % (petit modèle) à 5,5 % (grand modèle) dû à la transformation, la réduction du nombre total d'étapes a entraîné une accélération du temps réel estimée entre 1,29× et 1,37×.
• Ablations :
  • Voie Linéaire : Le retrait de la voie linéaire a conduit à un cas de « échec d'apprentissage », confirmant sa nécessité pour la stabilité.
  • Décalé vs Congruent : L'initialisation décalée a systématiquement surpassé l'initialisation congruente, particulièrement dans les petits modèles.
  • Annealing : Le calendrier de taux d'apprentissage pour les paramètres d'échelle a fourni des améliorations constantes, bien que faibles, par rapport aux multiplicateurs fixes.

5. Signification et Revendications

L'article affirme que SEL traite le décalage fondamental entre la nature multiplicative de nombreux calculs neuronaux et la nature additive de l'optimisation standard. En courbant l'espace des paramètres, SEL permet à l'optimiseur d'effectuer des pas qui sont naturellement proportionnels à la magnitude des poids, fournissant ainsi une mise à l'échelle relative sans nécessiter de réglage du taux d'apprentissage par paramètre.

Les auteurs soulignent que SEL est un remplacement pratique et direct pour les couches linéaires standards qui :

1. Accélère la convergence, particulièrement pour les modèles ayant des distributions de poids à queue lourde.
2. Ne nécessite aucun changement architectural lors de l'inférence.
3. Est inversible, permettant aux poids préentraînés standards d'être convertis en espace SEL pour la poursuite de l'entraînement ou le fine-tuning.

Ce travail suggère que la paramétrisation en « espace log doux » aide les poids à atteindre les cibles de grande magnitude requises par le réseau, lesquelles sont souvent désavantagées par les mises à jour additives linéaires.