Résumé technique : Sur la non-linéarité de la mise à l'échelle du taux d'apprentissage pour l'entraînement des LLM

Énoncé du problème

L'entraînement efficace des grands modèles de langage (LLM) nécessite un réglage minutieux des hyperparamètres, pourtant, effectuer directement des balayages de taux d'apprentissage aux échelles cibles est souvent prohibitif en termes de coûts. Par conséquent, les praticiens s'appuient sur le transfert d'hyperparamètres, prédisant les réglages optimaux pour des modèles à grande échelle à partir d'exécutions à plus petite échelle. Une hypothèse prédominante dans ce domaine est que le taux d'apprentissage optimal ($\eta^*$) suit une relation de mise à l'échelle log-linéaire (loi de puissance) par rapport à l'échelle des données ($D$) et à la taille du modèle ($N$) :
$$\log \eta^*(D, N) = a \log D + b \log N + c$$
Cependant, les études existantes échouent souvent à évaluer rigoureusement si cette hypothèse log-linéaire est globalement valide ou seulement locale, et elles omettent fréquemment le coût computationnel requis pour ajuster ces lois de mise à l'échelle. Cet article étudie la validité de cette loi de mise à l'échelle ainsi que l'efficacité de différentes stratégies d'extrapolation (axe des données vs axe du modèle) et de paramétrisation (taux d'apprentissage vs taux d'apprentissage effectif).

Méthodologie

Les auteurs ont mené une étude empirique systématique utilisant des architectures de type GPT-2 (allant de 22M à 707M de paramètres) entraînées sur le jeu de données FineWeb-100B avec des budgets de tokens allant de 5B à 100B.

1. Conception expérimentale :

   • Les modèles ont été entraînés avec AdamW avec une taille de lot fixe (0,52M de tokens) et un programme de type warmup-steady-decay (WSD).
   • Les taux d'apprentissage ont été balayés sur une grille logarithmique ($\log_2 \eta \in [-16, -7]$).
   • Taux d'apprentissage optimal ($\eta^*$) : Estimé en ajustant un polynôme cubique à la perte de validation en fonction de $\log \eta$.
   • Taux d'apprentissage effectif ($\eta_{eff}$) : Défini comme la taille du pas dans l'espace des poids normalisés : $\eta_{eff}(t) = \|\hat{w}_{t+1} - \hat{w}_t\|_2$, où $\hat{w} = w/\|w\|_2$. Le taux d'apprentissage effectif optimal ($\eta^*_{eff}$) a été estimé de manière similaire.

2. Métriques d'évaluation :

   • $R^2$ hors distribution ($R^2_{OOD}$) : Mesure la précision prédictive des lois de mise à l'échelle sur des échelles non vues.
   • Ratio de calcul supplémentaire (ECR - Extra Compute Ratio) : Quantifie le coût de calcul additionnel encouru lors de l'utilisation d'un taux d'apprentissage prédit par rapport au véritable taux optimal pour atteindre la même perte de validation.
   • Extrapolation sous contrainte de calcul : Les auteurs ont comparé l'extrapolation le long de l'axe des données (variation de $D$ avec un $N$ large et fixe) par rapport à l'axe du modèle (variation de $N$ avec un $D$ large et fixe) sous des budgets de calcul identiques.

3. Analyse théorique :

   • L'article analyse la dynamique implicite d'AdamW, spécifiquement la relation entre le taux d'apprentissage brut $\eta$ et le taux d'apprentissage effectif $\eta_{eff}$, en considérant le comportement transitoire des normes de poids.
   • Les expériences ont été répétées en utilisant AdamH, un optimiseur qui contrôle explicitement $\eta_{eff}$, afin d'isoler les effets de la dynamique des normes de poids.

Résultats clés

1. Non-linéarité du taux d'apprentissage optimal

L'étude démontre que le taux d'apprentissage optimal $\eta^*$ ne suit pas une loi de mise à l'échelle log-linéaire globale.

• Courbure ascendante : À mesure que l'échelle des données $D$ et la taille du modèle $N$ augmentent, $\eta^*$ présente une courbure ascendante systématique. L'extrapolation log-linéaire à partir de petites échelles sous-estime systématiquement le taux d'apprentissage optimal pour les échelles plus grandes.
• Linéarité du taux d'apprentissage effectif : En revanche, le taux d'apprentissage effectif ($\eta^*_{eff}$) présente un comportement de mise à l'échelle nettement plus linéaire sur les mêmes plages, avec des coefficients de corrélation de Pearson plus élevés.

2. Supériorité de l'extrapolation sur l'axe des données

Lors de la comparaison des stratégies d'extrapolation sous des budgets de calcul fixes :

• Axe des données vs Axe du modèle : L'ajustement des lois de mise à l'échelle le long de l'axe de la taille des données ($D$) produit des prédictions nettement plus précises que le long de l'axe de la taille du modèle ($N$).
• Efficacité de calcul : L'utilisation de l'extrapolation sur l'axe des données avec le taux d'apprentissage effectif n'entraîne qu'environ 2 % de calcul supplémentaire par rapport à l'utilisation du véritable taux d'apprentissage optimal. C'est moins de 20 % du calcul requis pour un balayage complet des hyperparamètres à l'échelle cible.

3. Explication mécaniste

L'article attribue la non-linéarité de $\eta^*$ à la dynamique transitoire des normes de poids dans AdamW :

• Deux régimes : La relation entre $\eta$ et $\eta_{eff}$ dépend de la question de savoir si la norme de poids $\|w_t\|$ a atteint l'équilibre.
  • Régime d'équilibre : $\eta_{eff} \propto \sqrt{\eta}$ (pente de 2 dans l'espace log-log).
  • Régime de pré-équilibre : $\eta_{eff} \propto \eta$ (pente de 1 dans l'espace log-log).
• Impact de l'échelle : À mesure que le taux d'apprentissage optimal diminue pour les modèles/données plus grands, l'horizon d'entraînement reste souvent dans le régime de pré-équilibre (où les normes de poids ne se sont pas encore relaxées). Cela déplace la dynamique de la taille du pas effectif, provoquant la courbure ascendante observée dans $\eta^*$.
• Validation via AdamH : En utilisant AdamH pour contrôler explicitement $\eta_{eff}$, la loi de mise à l'échelle devient strictement log-linéaire, confirmant que la non-linéarité dans l'AdamW standard provient de la cartographie implicite et variable dans le temps entre $\eta$ et $\eta_{eff}$.

Signification et affirmations

L'article revendique trois contributions principales :

1. Correction empirique : Il démontre que la loi de mise à l'échelle log-linéaire standard pour les taux d'apprentissage n'est qu'une propriété locale et échoue aux échelles supérieures en raison d'une courbure ascendante, conduisant à une extrapolation inexacte.
2. Stratégies de transfert pratiques : Il identifie deux choix spécifiques qui améliorent considérablement le transfert d'hyperparamètres :
   • Paramétrer la loi de mise à l'échelle en utilisant le taux d'apprentissage effectif ($\eta_{eff}$) plutôt que le taux d'apprentissage brut.
   • Extrapoler le long de l'axe des données ($D$) plutôt que l'axe de la taille du modèle ($N$).
3. Insight théorique : Il fournit une explication mécaniste de la non-linéarité observée, en la liant au comportement transitoire des normes de poids dans AdamW. Les auteurs soutiennent que $\eta_{eff}$ est la quantité la plus fondamentale régissant la dynamique d'optimisation dans les architectures invariantes d'échelle.

Les auteurs notent des limites, notamment l'accent mis sur les architectures GPT-2 et AdamW, le recours à l'ajustement cubique pour l'estimation de $\eta^*$, et le fait que l'analyse est limitée à des échelles allant jusqu'à 100B de tokens et 707M de paramètres. Ils ne prétendent pas que ces résultats se généralisent automatiquement à toutes les architectures ou optimiseurs sans vérification supplémentaire.