Résumé technique : Momentum de mémoire adaptative via un cadre basé sur modèle pour l'optimisation en apprentissage profond

1. Énoncé du problème

Les modèles d'apprentissage profond modernes reposent principalement sur des optimiseurs du premier ordre basés sur le momentum (par exemple, SGD avec Heavy Ball, AdamW, Lion). Ces méthodes utilisent un coefficient de momentum, $\beta$, pour régir la « mémoire » de l'optimiseur, déterminant dans quelle mesure les gradients passés influencent la direction de la mise à jour actuelle. En pratique, $\beta$ est presque universellement fixé à une valeur constante (généralement $\beta = 0,9$) et reste inchangé tout au long de l'entraînement.

Les auteurs identifient deux limitations critiques de cette approche statique :

1. Sous-optimalité : Les preuves empiriques suggèrent qu'un $\beta$ fixe unique est rarement optimal sur l'ensemble de la trajectoire d'entraînement ou à travers différents jeux de données et architectures de modèles.
2. Instabilité : Le paysage de performance autour du $\beta$ fixe optimal est hautement instable ; de légères déviations par rapport à la valeur optimale peuvent dégrader drastiquement la convergence.
3. Écart théorique : Bien que le momentum accélère la convergence dans des settings déterministes convexes, les garanties théoriques pour des settings stochastiques non convexes d'apprentissage profond montrent souvent que le Heavy Ball stochastique n'atteint pas de taux de convergence meilleurs que le SGD simple, malgré son succès empirique.

La question centrale posée est la suivante : Le coefficient de momentum peut-il s'adapter au fil du temps pour mieux correspondre au paysage d'optimisation ?

2. Méthodologie

L'article propose la Mémoire Adaptative (AM), un cadre qui remplace le coefficient de momentum constant par un coefficient dynamique, variant dans le temps $\beta_t$, calculé en ligne durant l'optimisation.

Cadre central : Approximation basée sur modèle

La méthode est dérivée de l'optimisation proximale basée sur modèle. Au lieu de minimiser directement la fonction de perte réelle $f(x)$, les auteurs l'approximent en utilisant un modèle de substitution $f^m_t(x)$ construit à partir de deux plans :

1. Le plan du gradient : Basé sur le gradient actuel $\nabla f(x_t)$.
2. Le plan du momentum : Basé sur la direction de descente précédente $\frac{1}{\eta}(x_{t-1} - x_t)$, qui encode le momentum accumulé.

Le modèle de substitution est défini comme suit :
$$f^m_t(x) = \max \left\{ f(x_t) + g_t^\top(x - x_t), \quad \hat{f}(x_t) + \frac{1}{\eta}(x_{t-1} - x_t)^\top(x - x_t) \right\}$$
où $g_t = \nabla f(x_t)$ et $\hat{f}(x_t)$ est un terme de biais représentant le plan du momentum.

Déduction de $\beta_t$ adaptatif

L'étape de mise à jour est formulée comme un problème de minimisation proximale avec un terme de régularisation supplémentaire pour contrôler l'alignement entre les directions de descente actuelle et précédente :
$$x_{t+1} \in \arg\min_x \left( f^m_t(x) + \frac{1}{2\eta}\|x - x_t\|^2 + \frac{\lambda}{\eta}\langle x_{t-1} - x_t, x - x_t \rangle \right)$$

La résolution de ce problème d'optimisation fournit une expression en forme close pour le coefficient de momentum adaptatif $\beta_t^*$. La solution implique un programme quadratique qui équilibre le gradient actuel et le momentum accumulé. La règle de mise à jour résultante ressemble à la méthode Heavy Ball mais avec un $\beta_t$ calculé dynamiquement :
$$d_{t+1} = \frac{\beta_t^* + \lambda}{1 + \lambda} d_t + \frac{1 - \beta_t^*}{1 + \lambda} g_t$$
$$x_{t+1} = x_t - \eta d_{t+1}$$

Détails de mise en œuvre pratique

Pour assurer la stabilité et l'efficacité dans des settings stochastiques, les auteurs introduisent plusieurs modifications pratiques :

• Recadrage (Clipping) : $\beta_t$ est recadré dans l'intervalle $[0, \beta_{\max}]$, où $\beta_{\max}$ est généralement fixé à la valeur fixe standard (par exemple, 0,9).
• Approximation de la perte : Calculer la différence exacte de perte $\Delta f = f(x_{t-1}) - f(x_t)$ sur le même lot est inefficace, tandis que l'utilisation de lots différents introduit du bruit. Les auteurs approximent cela en utilisant un développement du premier ordre : $\Delta f \approx \eta g_t^\top d_t$.
• Stratégie de surestimation : Contrairement aux méthodes de bundle traditionnelles qui utilisent des plans sous-estimants pour garantir la convergence, les auteurs constatent que l'utilisation d'un plan surestimant ($\hat{f}(x_t) = f(x_{t-1})$) conduit à une convergence plus rapide en pratique.
• Extension à AdamW : Le cadre est étendu à AdamW en incorporant des préconditionneurs diagonaux ($P_t$) et une décroissance de poids découplée. Le $\beta_t$ adaptatif est calculé par couche pour tenir compte de l'hétérogénéité des couches dans les grands modèles de langage (LLM).

3. Contributions clés

1. Motivation : L'article démontre empiriquement que le momentum fixe est sous-optimal et que la recherche du $\beta$ statique optimal est fastidieuse et non robuste.
2. Nouvelle méthodologie : Il introduit un cadre proximale utilisant une approximation à deux plans de la fonction de perte pour dériver un coefficient de momentum adaptatif en forme close. C'est la première application d'un tel cadre proximale spécifiquement pour l'optimisation basée sur le momentum.
3. Variantes algorithmiques : Les auteurs implémentent des variantes de Mémoire Adaptative pour le SGD (AM-MGD) et AdamW (AM-AdamW).
4. Garanties théoriques : Des taux de convergence de $O(1/\sqrt{T})$ sont établis pour les settings convexes et non convexes sous des hypothèses standard de régularité et de gradient borné, plaçant la méthode dans le paysage standard des méthodes stochastiques du premier ordre.
5. Efficacité des hyperparamètres : La méthode ne nécessite aucun réglage supplémentaire d'hyperparamètres au-delà des paramètres de base standards (par exemple, $\lambda$ et $\beta_{\max}$ sont fixes à travers les tâches).

4. Résultats expérimentaux

Les auteurs ont évalué AM sur un spectre de tâches :

• Problèmes convexes : Sur des tâches de régression logistique (jeux de données LIBSVM), AM-MGD a constamment surpassé les baselines à momentum fixe (y compris le $\beta^*$ optimal recherché par grille) sans nécessiter de réglage.
• Classification d'images : Sur CIFAR-10/100 et ImageNet utilisant les architectures VGG, ResNet et Wide-ResNet, AM-MGD a atteint une optimisation plus rapide et une meilleure généralisation.
  • Dynamiques : Le $\beta_t$ adaptatif a naturellement présenté trois phases : un momentum faible (similaire au SGD) durant l'entraînement initial bruyant, un momentum élevé durant la phase d'accélération, et des « redémarrages doux » fréquents (chutes de $\beta_t$) près des points d'interpolation où les gradients s'affaiblissent.
• Pré-entraînement de grands modèles de langage : Sur les modèles LLaMA (de 20M à 1B de paramètres) entraînés sur C4 :
  • AM-AdamW a surpassé AdamW standard à toutes les échelles.
  • Élimination du warm-up : Notamment, AM-AdamW a stabilisé l'entraînement précoce même sans un calendrier de warm-up du taux d'apprentissage, un régime où AdamW standard échoue souvent ou nécessite un réglage minutieux. Cela suggère que AM peut servir d'alternative sans hyperparamètres au warm-up.
• Études d'ablation :
  • La méthode est robuste au choix de $\lambda$ (dans $[0,01, 1]$).
  • Les gains de performance persistent à travers des tailles de lots et des taux d'apprentissage variables, AM permettant une convergence à des taux d'apprentissage plus élevés où les baselines statiques échouent.
  • La surcharge computationnelle est négligeable (augmentation du temps d'environ 0,2 % pour les LLM) sans empreinte mémoire supplémentaire.

5. Signification et affirmations

L'article affirme que la Mémoire Adaptative offre une alternative raisonnée et sans réglage au momentum statique. Sa signification réside dans :

• Le pont entre théorie et pratique : Elle fournit un mécanisme qui s'adapte à la géométrie locale du paysage de perte, comblant le fossé entre les taux de convergence théoriques et les performances empiriques.
• Simplification des pipelines d'entraînement : En stabilisant l'entraînement précoce et en éliminant potentiellement le besoin de calendriers de warm-up du taux d'apprentissage, AM simplifie l'entraînement des modèles à grande échelle.
• Robustesse : La méthode démontre une robustesse aux taux d'apprentissage élevés et aux tailles de lots variables, réduisant la sensibilité de l'entraînement en apprentissage profond au choix des hyperparamètres.

Les auteurs concluent que ce travail ouvre la voie à de nouvelles façons d'induire l'adaptativité dans l'optimisation, spécifiquement en allant au-delà des taux d'apprentissage adaptatifs pour atteindre des momentum adaptatifs, et suggère un travail futur sur les coefficients de momentum par paramètre et une investigation plus approfondie des redémarrages de momentum.