Résumé technique : Entraînement pour le modèle que vous restituez

Énoncé du problème

Les pipelines de modèles de langage (LM) modernes restituent fréquemment un modèle moyenné — typiquement une moyenne exponentielle mobile (EMA) des itérations d'entraînement — plutôt que l'itération finale elle-même. Bien que la moyenne des itérations soit connue pour stabiliser l'optimisation et améliorer la généralisation, une question fondamentale demeure : étant donné que la sortie finale sera une moyenne d'itérations, comment le processus d'entraînement lui-même doit-il être modifié pour optimiser la performance de cet estimateur spécifique ?

Les travaux antérieurs ont abordé cela de manière implicite en observant que la moyenne permet des taux d'apprentissage plus agressifs. Cependant, ce papier adopte une approche directe : il cherche à concevoir un algorithme d'optimisation qui contrôle explicitement la trajectoire d'entraînement pour minimiser l'erreur de l'estimateur de moyenne d'itérations restitué.

Méthodologie

Formulation théorique

Les auteurs cadrent le problème comme un problème de contrôle optimal au sein d'un modèle quadratique stochastique en temps continu.

• Dynamique : La trajectoire des paramètres $\theta_t$ suit une équation différentielle stochastique (SDE) représentant une descente de gradient bruitée près d'un minimum quadratique (processus d'Ornstein–Uhlenbeck), augmentée par une entrée de contrôle $u_t$.
• Objectif : Le but est de trouver une stratégie de contrôle $u_t$ qui minimise l'erreur quadratique moyenne entre l'estimateur moyenné temporellement restitué $\hat{\theta}_{EMA, T}$ et l'optimum réel $\mu^*$, sous réserve d'une pénalité quadratique sur l'amplitude de l'intervention de contrôle (paramètre de régularisation $\lambda$).
• Dérivation : En utilisant la théorie du contrôle LQG (Linear-Quadratic-Gaussian) et le principe d'équivalence de certitude, les auteurs dérivent un contrôleur optimal. La solution exacte dépend de la moyenne a posteriori de l'optimum étant donné la trajectoire, ce qui est intraitable pour les LM à grande échelle.

Approximation pratique : PACE

Pour rendre la solution implémentable, les auteurs introduisent une série d'approximations heuristiques :

1. Substitution de l'estimateur : Remplacer la moyenne a posteriori inconnue de l'optimum par la moyenne actuelle des itérations ($\hat{\theta}_{EMA, t}$).
2. Approximation temporelle : Supposer que le temps actuel $t$ est très inférieur au temps total d'entraînement $T$ ($t \ll T$), simplifiant ainsi les coefficients dépendant du temps.
3. Discrétisation et préconditionnement : Discrétiser le contrôle continu et remplacer la matrice Hessienne inconnue par le préconditionneur diagonal utilisé dans AdamW (l'estimation du second moment $v_t$).

Cela produit PACE (Pullback Averaging Control for Efficient Optimization), un wrapper léger autour d'AdamW. La règle de mise à jour ajoute un terme de "rappel" (pullback) qui déplace les poids vivants vers les poids EMA :
$$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$$
où $C$ est une matrice diagonale dérivée du préconditionneur d'Adam, et l'intensité du rappel est plafonnée pour garantir que la mise à jour reste une combinaison convexe de l'étape standard et du rappel.

Contributions clés

1. Formulation de contrôle optimal : Le papier fournit la première dérivation fondée sur des principes d'un algorithme d'optimisation spécifiquement conçu pour minimiser l'erreur d'un estimateur de moyenne d'itérations, plutôt que de l'itération finale.
2. Garanties théoriques :
   • Cadre convexe : Les auteurs prouvent qu'une version stylisée de PACE converge vers le minimum de toute fonction de perte convexe avec des gradients stochastiques non biaisés au taux standard de $O(1/\sqrt{T})$, à un facteur constant près (dépendant de la force de l'EMA). Cela garantit que la méthode n'est jamais significativement moins bonne qu'un SGD standard.
   • Cadre quadratique : Dans le cadre quadratique utilisé pour la dérivation, les dynamiques contrôlées améliorent strictement l'erreur quadratique limite de l'estimateur de moyenne d'itérations par rapport aux dynamiques non contrôlées (SGD standard avec EMA). L'amélioration peut être arbitrairement grande sur certaines instances quadratiques.
3. Validation empirique : La méthode est évaluée sur :
   • Fine-tuning : Modèles SmolLM2-1.7B, Qwen3-1.7B et Gemma3-1B sur le jeu de données smol-smoltalk.
   • Pré-entraînement : GPT-2 (124M) sur FineWeb avec un budget de tokens Chinchilla-optimal.
   • Baselines : PACE est comparé à AdamW vanilla, AdamW avec évaluation EMA, et l'algorithme Schedule-Free.

Résultats

• Performance : PACE améliore strictement tant l'AdamW vanilla que l'AdamW avec évaluation EMA sur tous les modèles testés (1B–2B paramètres) et jeux de données.
• Robustesse : L'amélioration est robuste sur une large gamme de taux d'apprentissage, de schémas de décroissance (constant, cosinus, Warmup-Stable-Decay) et d'hyperparamètres (force de rappel $c$, puissance $\kappa$ de l'EMA).
• Comparaison avec Schedule-Free : PACE est compétitif avec l'algorithme Schedule-Free (Defazio et al., 2024), qui utilise également la moyenne d'itérations pour supprimer le besoin de décroissance du taux d'apprentissage. Cependant, PACE atteint cela avec une règle de mise à jour différente.
• Décroissance du taux d'apprentissage : PACE performe bien même avec des taux d'apprentissage constants, suggérant qu'il peut stabiliser l'entraînement dans des régimes où la décroissance du taux d'apprentissage est typiquement requise. Dans les comparaisons de budget de tokens, la trajectoire à taux constant de PACE suit le "fond du fleuve" (river floor) du paysage de perte aussi efficacement que les branches WSD, et les surpasse souvent.

Signification et limites

Signification :
Le papier démontre qu'optimiser pour le modèle restitué (la moyenne) plutôt que pour l'itération finale conduit à des gains de performance tangibles. En traitant la trajectoire d'entraînement comme un problème de contrôle, PACE fournit une méthode théoriquement fondée pour stabiliser l'entraînement et améliorer la généralisation sans nécessiter un réglage complexe des schémas. Les auteurs notent que PACE offre une convergence théorique améliorée dans le cadre quadratique par rapport aux méthodes standards, bien que cela s'accompagne d'une garantie de convergence théorique légèrement plus lente dans le cadre convexe général dû au facteur EMA.

Limites :

• Hypothèses théoriques : La dérivation repose sur un modèle quadratique simplifié avec un bruit additif. Bien que la convergence soit prouvée pour les fonctions convexes générales, l'incorporation rigoureuse du momentum et du préconditionnement adaptatif (central à AdamW) dans l'analyse théorique reste une direction ouverte.
• Échelle : Les expériences ont été limitées à des modèles allant jusqu'à 2B paramètres pour le fine-tuning et 124M pour le pré-entraînement. L'efficacité à des échelles plus grandes (ex: 70B+) n'est pas vérifiée.
• Surcharge mémoire : PACE nécessite de maintenir une copie supplémentaire des poids du modèle (l'état EMA) pendant l'entraînement, ce qui introduit une surcharge mémoire non négligeable, bien que cela soit atténué par de grands batchs en pratique.
• Sensibilité à la graine (seed) : En raison des contraintes de calcul, les tests principaux ont utilisé une seule graine, bien que des études de robustesse multi-graines aient confirmé la stabilité de l'ordre des résultats.

Les auteurs concluent que PACE représente une étape prometteuse vers l'alignement des algorithmes d'optimisation avec la stratégie de déploiement réelle des LM modernes, où le modèle moyenné est le produit final.