Résumé technique : Théorie des calendriers de taux d'apprentissage optimaux et lois d'échelle pour un modèle à caractéristiques aléatoires

Énoncé du problème

L'entraînement des modèles d'apprentissage profond nécessite la sélection d'hyperparamètres tels que le taux d'apprentissage (LR), la taille de lot et l'horizon d'entraînement. Bien que l'approche empirique par « essais et erreurs » soit courante, il existe un manque de compréhension théorique concernant la manière dont les hyperparamètres optimaux évoluent avec l'horizon d'entraînement $T$. Les stratégies existantes pour transférer les hyperparamètres entre différentes tailles de modèles échouent souvent à se transférer entre différents horizons d'entraînement. De plus, bien que divers calendriers de LR (par exemple, décroissance linéaire, échauffement-stabilité-décroissance, recuit cosinusoïdal) soient utilisés en pratique, les conditions théoriques dictant quel calendrier est optimal pour une tâche donnée restent floues. Cet article répond au besoin d'un cadre rigoureux pour dériver des calendriers de LR optimaux dépendant de l'horizon d'entraînement et de la structure spécifique de la tâche d'apprentissage.

Méthodologie

Les auteurs emploient un modèle à caractéristiques aléatoires en loi de puissance entraîné par descente de gradient stochastique (SGD) pour dériver des calendriers d'hyperparamètres optimaux. Le modèle suppose des données d'entrée $x$ avec une fonction cible $y(x) = w^* \cdot \psi(x) + \sigma_0 z$, où les caractéristiques $\psi(x)$ présentent une décroissance spectrale en loi de puissance ($\lambda_k \sim k^{-b}$) et les poids de l'enseignant évoluent comme $(w^*_k)^2 \lambda_k \sim k^{-a}$.

La méthodologie centrale fait appel à la théorie du contrôle optimal :

1. Formulation de la dynamique : La perte de test attendue $L_t$ est dérivée comme une fonction du calendrier de taux d'apprentissage $\{\eta_t\}$ et du calendrier de taille de lot $\{m_t\}$. Les dynamiques de la perte sont régies par une relation de récurrence linéaire pour les erreurs de mode $c_{t,k}$.
2. Objectif d'optimisation : Le but est de minimiser la perte de test finale $L_T$ à un horizon fixe $T$. Les auteurs formulent trois problèmes d'optimisation :
   • LR constant optimal.
   • Calendrier de LR optimal (variable dans le temps $\eta_t$).
   • Calendrier conjoint optimal de LR et de taille de lot.
3. Approches de résolution :
   • Numérique : Utilisation d'une méthode de tir unique (implémentée dans CasADi) pour résoudre le problème de contrôle optimal en temps discret pour un $T$ fini.
   • Analytique : Dérivation d'approximations en temps continu pour un grand $T$ en utilisant le théorème de Noether et l'analyse asymptotique du fonctionnel de perte. Cela implique la minimisation d'un fonctionnel du taux d'apprentissage intégré $\chi(t) = \int_t^T \eta(s) ds$.

Contributions et résultats clés

1. Identification de deux phases distinctes

L'analyse révèle deux régimes distincts pour les calendriers de taux d'apprentissage optimaux, déterminés par la relation entre les exposants spectraux $a$ (difficulté de la tâche) et $b$ (complexité des caractéristiques) :

• La phase facile ($a > b$) :

  • Le calendrier optimal suit une forme de décroissance polynomiale : $\eta^*_T(t) \sim T^{b/a - 1} f(t/T)$, où $f(z) \propto (1-z)^{2b-1}$.
  • Le taux d'apprentissage décroît progressivement pour équilibrer la réduction du biais et de la variance tout au long de la trajectoire d'entraînement.
  • L'exposant de mise à l'échelle de la perte est $\zeta_T = (a-1)/a$, correspondant au taux minimax-optimal pour des tailles de lot fixes.

• La phase difficile ($a < b$) :

  • Le calendrier optimal ressemble au WSD (Warmup-Stable-Decay). Il maintient un taux d'apprentissage maximal constant ($\eta_{max}$) pour la majorité des étapes d'entraînement, suivi d'une phase de recuit rapide sur une fraction négligeable du temps total ($1 - t_s/T \sim T^{-(b-a)/(2b-1)}$).
  • Cette stratégie privilégie la suppression du biais via un LR constant maximal, en compressant la réduction de la variance dans les dernières étapes.
  • L'exposant de mise à l'échelle de la perte s'améliore pour atteindre $\zeta_T = (a-1)/b$, surpassant les stratégies de LR constant.

2. Supériorité par rapport aux références

Les calendriers optimaux dérivés surpassent strictement les références standard :

• LR constant : Le calendrier optimal atteint de meilleurs exposants de mise à l'échelle dans les deux phases.
• Calendriers en loi de puissance : Même en optimisant les exposants des calendriers en loi de puissance ($\eta \sim T^{-\xi} t^{-\delta}$), ils récupèrent asymptotiquement le même comportement de mise à l'échelle que la référence de LR constant optimale, échouant à égaler la performance des calendriers optimaux dépendant de l'horizon dérivés via la théorie du contrôle.

3. Optimisation conjointe du taux d'apprentissage et de la taille de lot

Les auteurs étendent l'analyse pour optimiser conjointement $\eta(t)$ et $m(t)$ sous un budget de données fixe $B_{tot}$ :

• Phase facile : Un « ramp de lot » (augmentation de la taille de lot au fil du temps) combiné à un taux d'apprentissage maximal constant réduit le temps réel. La perte évolue comme $T^{-(a-1)/b}$, améliorant l'exposant du temps réel par rapport aux tailles de lot fixes, tandis que l'exposant du budget d'échantillons reste $(a-1)/a$.
• Phase difficile : L'optimisation conjointe produit les mêmes exposants de mise à l'échelle que l'optimum à taille de lot fixe, la taille de lot restant bornée.

4. Calendriers de momentum optimaux

Pour les tâches avec momentum, les auteurs optimisent le paramètre de momentum $\beta(t)$ parallèlement à $\eta(t)$.

• Dans la phase facile, le calendrier de momentum optimal produit des lois d'échelle similaires à celles du SGD optimal.
• Dans la phase difficile, l'optimisation du calendrier de momentum fournit une amélioration de l'exposant de mise à l'échelle au-delà de ce qui est réalisable avec le SGD seul. Les dynamiques de momentum optimales montrent une diminution transitoire initiale suivie d'une augmentation lorsque $t/T \to 1$.

5. Validation empirique dans les réseaux profonds

Les phases théoriques sont validées sur des références d'apprentissage profond :

• ResNet sur CIFAR-5M : Présente un comportement de phase difficile. Le taux d'apprentissage de base optimal pour les calendriers de recuit est approximativement constant à travers les horizons, cohérent avec le profil WSD.
• Transformateurs de style GPT-2 sur C4 : Présentent un comportement de phase facile. Le taux d'apprentissage de base optimal se décale (diminue) avec l'horizon $T$, même sous recuit, cohérent avec le profil de décroissance polynomiale.

Importance et affirmations

L'article prétend fournir une base théorique pour comprendre pourquoi les calendriers de taux d'apprentissage varient selon les tâches et les horizons d'entraînement. En utilisant la théorie du contrôle optimal sur un modèle à caractéristiques aléatoires soluble, les auteurs démontrent que :

1. Dépendance à l'horizon : Les taux d'apprentissage optimaux ne sont pas automatiquement transférables entre les horizons d'entraînement ; le calendrier doit être explicitement conçu en fonction de $T$.
2. Structure de la tâche : La « difficulté » d'une tâche (caractérisée par les exposants spectraux $a$ et $b$) dicte si une tâche tombe dans la phase facile ou difficile, déterminant ainsi la forme du calendrier optimal (décroissance polynomiale vs WSD).
3. Efficacité : Les calendriers dépendant de l'horizon et l'optimisation conjointe de la taille de lot peuvent considérablement améliorer les lois d'échelle et réduire le temps d'entraînement réel par rapport aux heuristiques standard.

Les auteurs reconnaissent les limites, notant que leur théorie s'applique actuellement aux modèles linéaires dans le régime paresseux/noyau et ne dérive pas encore analytiquement l'exposant amélioré spécifique pour le momentum dans la phase difficile, ni ne prend en compte des optimiseurs sophistiqués comme Adam ou les méthodes de préconditionnement. Cependant, ce travail offre une explication rigoureuse de l'efficacité de calendriers comme le WSD et suggère que le réglage empirique des taux d'apprentissage de base à travers les horizons devrait tenir compte des propriétés spectrales sous-jacentes du modèle et de la tâche.