Résumé technique : Transfert d'hyperparamètres avec des couches à mélange d'experts

1. Énoncé du problème

Les couches à mélange d'experts (MoE) sont devenues un mécanisme crucial pour mettre à l'échelle les réseaux de neurones modernes en découplant le nombre total de paramètres entraînables des paramètres actifs lors du passage avant. Cependant, l'entraînement de modèles MoE épars introduit une complexité significative concernant la sélection des hyperparamètres (HP). Contrairement aux modèles denses, les architectures MoE introduisent de nouveaux paramètres entraînables (poids du routeur) et de nouvelles dimensions architecturales (nombre d'experts, taille de l'expert) qui nécessitent un réglage minutieux.

Le réglage direct des hyperparamètres (tels que le taux d'apprentissage, l'échelle d'initialisation et la décroissance des poids) à grande échelle est prohibitif en termes de calcul. Bien que des techniques de transfert d'HP existent pour les transformateurs denses — permettant d'extrapoler les hyperparamètres optimaux trouvés sur de petits modèles vers des modèles plus grands —, ces méthodes n'ont pas été rigoureusement adaptées aux dynamiques d'échelle spécifiques des couches MoE éparses. Le défi central consiste à déterminer comment mettre à l'échelle les hyperparamètres lors de l'augmentation simultanée de la largeur, de la profondeur, du nombre d'experts et de la taille de l'expert, sans réajuster à chaque échelle, tout en garantissant la stabilité et les performances de l'entraînement.

2. Méthodologie

2.1. Paramétrisation proposée

Les auteurs proposent une nouvelle paramétrisation pour les modèles transformateurs avec des couches MoE, étendant la paramétrisation CompleteP (développée précédemment pour les transformateurs denses) afin d'inclure des règles d'échelle spécifiques aux MoE. L'objectif est de définir un ensemble de règles prédisant comment les valeurs brutes des hyperparamètres (écart-type d'initialisation $\sigma$ et taux d'apprentissage $\eta$) doivent évoluer à mesure que les dimensions du modèle augmentent, afin que les dynamiques d'entraînement restent cohérentes.

La dérivation repose sur le principe de la paramétrisation de la mise à jour maximale ($\mu$P), qui exige que les composants du réseau (pré-activations et mises à jour résiduelles) restent en $O(1)$ à l'initialisation et reçoivent des mises à jour de $\Theta(1)$ par étape d'entraînement. Les auteurs étendent cela aux MoE en exigeant que les conditions de mise à jour maximale s'appliquent non seulement à la sortie de la couche, mais également aux composants individuels des experts (coefficients de mélange et sorties d'experts).

Les règles d'échelle clés dérivées pour le module MoE incluent :

• Poids du routeur : Le taux d'apprentissage évolue comme $\eta \propto n_{embd}^{-1}$. L'initialisation évolue comme $n_{embd}^{-\gamma}$ (avec $\gamma \geq 0.5$).
• Biais d'experts : Initialisés à zéro avec un taux d'apprentissage constant $\eta \propto 1$ (indépendant du nombre d'experts), à condition que l'éparcité soit fixe.
• Poids MLP d'experts :
  • Projection ascendante ($W_{up}$) : $\sigma_{init} \propto n_{embd}^{-1/2}$, $\eta \propto n_{embd}^{-1}$.
  • Projection descendante ($W_{down}$) : $\sigma_{init} \propto \alpha_{ffn}^{-1} n_{embd}^{-1/2}$, $\eta \propto \alpha_{ffn}^{-1} n_{embd}^{-1}$.
  • Ici, $\alpha_{ffn}$ représente le multiplicateur de la taille cachée de l'expert par rapport à la dimension d'embedding.

Crucialement, les auteurs fixent le ratio d'éparcité $\kappa = n_{act}/n_{exp}$ (la fraction d'experts actifs) comme une constante tout en augmentant le nombre total d'experts ($n_{exp}$) et le nombre d'experts actifs ($n_{act}$). Cela contraste avec les approches qui fixent le nombre d'experts actifs tout en augmentant le pool total.

2.2. Fondement théorique : Théorie du champ moyen dynamique (DMFT)

Pour justifier ces règles d'échelle heuristiques, les auteurs emploient la Théorie du champ moyen dynamique (DMFT). Ils analysent les dynamiques d'entraînement des réseaux résiduels avec des couches MoE dans la limite simultanée d'une largeur infinie ($n_{embd}$), d'une profondeur infinie ($L$), d'une taille d'expert infinie ($n_{hid}$) et d'un nombre d'experts infini ($n_{exp}$), tout en maintenant un éparcité d'activation constante $\kappa$.

L'analyse révèle une hiérarchie de champ moyen à trois niveaux novatrice :

1. Représentations du flux résiduel : Champ moyen sur les sorties d'experts.
2. Sorties d'experts : Champ moyen sur les neurones individuels de chaque expert.
3. Neurones individuels : Champ moyen au sein de chaque expert.

L'analyse DMFT démontre que, sous la paramétrisation proposée, les dynamiques d'entraînement limites sont :

• Indépendantes du ratio FFN ($\alpha_{ffn}$) : Les dynamiques ne dépendent pas de la taille spécifique des experts par rapport à la dimension d'embedding, à condition que la limite d'échelle conjointe soit prise.
• Dépendantes uniquement de l'éparcité $\kappa$ : Les dynamiques sont cohérentes pour tous les paramètres d'échelle tant que le ratio d'éparcité est fixe.
• Invariants d'échelle : L'évolution des statistiques de résumé du réseau (par exemple, les noyaux de caractéristiques par couche) est cohérente à travers les échelles, garantissant théoriquement que les hyperparamètres se transfèrent de manière fiable.

2.3. Configuration expérimentale

Les auteurs valident leur paramétrisation empiriquement en utilisant des modèles de langage transformateurs uniquement décodeur sur les ensembles de données FineWeb et C4.

• Modèles de base : Réglés sur de petits modèles avec environ 38 M de paramètres activés.
• Mise à l'échelle : Étendus jusqu'à des modèles comportant jusqu'à 2 M de paramètres au total, en variant la largeur, la profondeur, le nombre d'experts et la taille de l'expert.
• Budget de tokens : Les expériences ont été menées sur un budget de tokens fixe de 1 milliard de tokens (2000 étapes) pour isoler les dynamiques d'entraînement précoce, ainsi que sur des horizons plus longs (jusqu'à 7,5 milliards de tokens).
• Optimiseur : Optimiseur Adam standard.
• Équilibrage de charge : Une stratégie sans perte auxiliaire a été utilisée, mettant à jour directement les biais des experts pour encourager l'équilibrage de charge sans ajouter de terme de régularisation à la fonction de perte.

3. Contributions clés

1. Paramétrisation MoE : L'article étend la paramétrisation CompleteP aux modèles MoE épars, fournissant des règles d'échelle explicites pour les poids du routeur, les biais d'experts et les poids MLP d'experts à travers la largeur, la profondeur, le nombre d'experts et la taille de l'expert.
2. Justification théorique via DMFT : Les auteurs fournissent une fondation théorique rigoureuse pour leur paramétrisation en utilisant la DMFT. Ils dérivent une description explicite des dynamiques d'entraînement dans la limite de largeur/profondeur infinie, prouvant que les dynamiques convergent vers un système stable et invariant d'échelle qui dépend uniquement de l'éparcité, et non de l'échelle spécifique des dimensions des experts.
3. Validation empirique du transfert d'HP : L'étude démontre que les hyperparamètres optimaux (taux d'apprentissage et échelle d'initialisation) identifiés sur de petits modèles de base (38 M de paramètres actifs) se transfèrent de manière fiable vers des modèles beaucoup plus grands (jusqu'à 2 M de paramètres au total) à travers diverses dimensions architecturales.
4. Insights architecturaux : Les auteurs vérifient empiriquement que, sous leur paramétrisation, augmenter le nombre d'experts (tout en maintenant le nombre total de paramètres fixe) produit de meilleures performances qu'augmenter la taille des experts individuels. Cette découverte s'aligne sur la littérature récente mais est obtenue ici sans nécessiter de balayages coûteux d'hyperparamètres à chaque échelle.

4. Résultats

• Transfert fiable d'HP : Selon les règles d'échelle proposées, les taux d'apprentissage de base optimaux et les écarts-types d'initialisation se transfèrent efficacement à travers des modèles allant de 51 M à 2 M de paramètres au total. Les courbes de perte des modèles mis à l'échelle se superposent à celles du modèle de base lors des premières itérations avant de diverger (les modèles plus grands atteignant une perte plus faible).
• Stabilité : La paramétrisation assure des dynamiques d'entraînement stables, y compris un équilibrage de charge uniforme des experts, même lors de l'augmentation du nombre d'experts. Les auteurs notent que le pré-entraînement MoE est particulièrement sensible aux hyperparamètres d'échelle constante (multiplicateurs traités comme $\Theta(1)$), et que leur réglage est crucial pour la stabilité.
• Performances : Les modèles entraînés avec des hyperparamètres zero-shot (transférés depuis de petits modèles) atteignent des performances compétitives par rapport aux bases denses (par exemple, GPT-2 small/medium) lorsqu'ils sont appariés sur le nombre de paramètres actifs.
• Nombre d'experts vs Taille : Les expériences confirment qu'augmenter le nombre d'experts est plus efficace en termes de paramètres qu'augmenter la taille des experts à un nombre de paramètres fixe. Cet avantage persiste même sur des horizons d'entraînement plus longs (5 milliards de tokens).
• Sensibilité à l'éparcité : L'étude met en évidence que le transfert d'hyperparamètres n'est valide que lorsque le ratio d'éparcité $\kappa$ est fixe. Augmenter le nombre d'experts tout en fixant le nombre d'experts actifs (envoyant $\kappa \to 0$) brise la transférabilité des hyperparamètres optimaux.

5. Importance et revendications

L'article revendique fournir un cadre pratique mais rigoureux pour la mise à l'échelle des modèles MoE. En combinant une paramétrisation heuristique avec une analyse DMFT, les auteurs offrent une méthode pour :

• Réduire les coûts d'entraînement : Permettre la sélection d'hyperparamètres optimaux pour des modèles MoE à grande échelle en réglant uniquement de petits modèles de base peu coûteux.
• Garantir la stabilité : Fournir des règles qui préviennent l'instabilité de l'entraînement (par exemple, effondrement ou divergence des experts) qui est courante dans l'entraînement MoE épars.
• Guider la conception architecturale : Offrir des preuves théoriques et empiriques qu'augmenter le nombre d'experts est préférable à l'augmentation de la taille des experts sous une éparcité fixe, aidant ainsi à la conception de modèles à grande échelle efficaces.

Les auteurs reconnaissent des limites, notant que le travail actuel se concentre sur le taux d'apprentissage et l'initialisation, laissant d'autres hyperparamètres (taille de lot, décroissance des poids, calendriers de LR) pour une investigation future. Ils notent également que, bien que l'analyse DMFT soutienne la limite de largeur infinie, le comportement des transferts de faible largeur nécessite une étude théorique supplémentaire. De plus, l'article ne revendique pas résoudre les lois d'échelle « optimales en calcul » pour les MoE (par exemple, les exposants Chinchilla), car le compromis FLOP-performance dans les MoE diffère considérablement de celui des modèles denses en raison des contraintes matérielles induites par l'éparcité.

En résumé, ce travail établit une paramétrisation fondamentale permettant l'extrapolation fiable des hyperparamètres des petits aux grands modèles MoE, fondée sur une théorie novatrice de champ moyen à trois niveaux des dynamiques d'entraînement.