Résumé technique : Seesaw : Accélérer l'entraînement en équilibrant la planification du taux d'apprentissage et de la taille des lots

Énoncé du problème

Le pré-entraînement des grands modèles de langage (LLM) est de plus en plus contraint par le temps de calcul réel (wall-clock time), s'étendant souvent sur plusieurs mois pour les modèles de pointe. Une stratégie primaire pour réduire cette durée consiste à augmenter la taille des lots (batch size), ce qui exploite le calcul parallèle pour réduire le nombre total d'étapes d'optimisation séquentielles. Cependant, augmenter simplement la taille des lots de manière indéfinie mène à des rendements décroissants au-delà d'une « taille de lot critique » (Critical Batch Size, CBS), où l'efficacité de l'échantillonnage chute.

Bien que des campagnes d'entraînement récentes à grande échelle (ex. LLaMA, Nemotron, OLMo) aient employé des programmes de « rampe de lot » (batch ramp) — augmentant progressivement la taille des lots au cours de l'entraînement — ces stratégies sont actuellement ajustées de manière heuristique. Il existe un manque de justification théorique sur la manière d'équilibrer de façon optimale la décroissance du taux d'apprentissage avec la montée en puissance de la taille des lots, particulièrement pour les optimiseurs adaptatifs comme Adam. La question centrale abordée est la suivante : Quel est le programme de taille de lot optimal pour minimiser le temps d'exécution sériel sans sacrifier la performance du modèle ?

Méthodologie

Cadre théorique

Les auteurs développent un cadre principiel basé sur la régression linéaire bruitée pour établir une équivalence entre la décroissance du taux d'apprentissage et la montée en puissance de la taille des lots.

1. Équivalence SGD : Le papier fournit la première preuve (non asymptotique) sur échantillon fini démontrant que, pour la descente de gradient stochastique (SGD) sur la régression linéaire bruitée, un processus où la taille du lot double à chaque phase alors que le taux d'apprentissage reste fixe est équivalent (à un facteur constant près dans l'excès de risque) à un processus où la taille du lot est fixe et le taux d'apprentissage est divisé par deux à chaque phase, à condition que les deux processus consomment le même nombre total de points de données.
2. Extension à la SGD normalisée (NSGD) : Pour combler l'écart avec les optimiseurs adaptatifs comme Adam, les auteurs analysent la SGD normalisée, un substitut analytique traçable pour Adam. Ils introduisent l'Hypothèse 2, qui postule que dans le régime dominé par la variance observé en pratique, les normes de carrés de gradient attendues sont dominées par un bruit additif (évoluant en $O(1/B)$). Sous cette hypothèse, la NSGD se réduit effectivement à une SGD avec un taux d'apprentissage rescalé.
3. La condition d'équivalence : L'analyse dérive une relation spécifique requise pour maintenir la dynamique de perte lors de l'ajustement de ces deux paramètres. Pour un facteur de décroissance du taux d'apprentissage $\alpha$ et un facteur d'augmentation de la taille du lot $\beta$, les processus restent équivalents si le produit $\alpha\sqrt{\beta}$ est maintenu constant.

L'algorithme Seesaw

Sur la base de l'équivalence théorique, les auteurs proposent Seesaw, un planificateur qui équilibre dynamiquement le taux d'apprentissage et la taille des lots :

• Mécanisme : Chaque fois qu'un planificateur standard (ex. décroissance cosinus) réduirait le taux d'apprentissage d'un facteur $\alpha$, Seesaw réduit plutôt le taux d'apprentissage de $\sqrt{\alpha}$ et augmente la taille du lot de $\alpha$.
• Contrainte : Les auteurs identifient une limite théorique pour prévenir la divergence. Le schéma de montée en puissance le plus agressif qui maintient la stabilité est défini par $\alpha = \sqrt{\beta}$. Dépasser cela (c'est-à-dire augmenter la taille du lot trop agressivement par rapport à la décroissance du taux d'apprentissage) entraîne une divergence.
• Implémentation : Seesaw est conçu comme un remplacement direct (drop-in replacement) des planificateurs existants. En pratique, il approxime la décroissance cosinus par une décroissance par paliers (step decay), déclenchant des augmentations de taille de lot aux nombres de tokens spécifiques où le taux d'apprentissage serait autrement coupé.

Contributions clés

1. Équivalence théorique : Le papier établit la première équivalence non asymptotique entre la décroissance du taux d'apprentissage et la montée en puissance de la taille des lots pour la SGD, et étend cela à la SGD normalisée sous un régime dominé par la variance.
2. Le planificateur Seesaw : Introduction d'un algorithme principiel qui remplace la décroissance standard du taux d'apprentissage par un programme couplé de taux d'apprentissage et de taille de lot, théoriquement fondé pour préserver la dynamique de perte.
3. Analyse de l'accélération : Les auteurs dérivent une borne supérieure théorique pour la réduction du temps d'exécution sériel. En utilisant la montée en puissance la plus agressive et stable ($\alpha = \sqrt{\beta}$) par rapport à une référence de décroissance cosinus, la vitesse maximale théorique est calculée comme étant de $(1 - 2/\pi) \approx 36,3\%$.

Résultats empiriques

Les auteurs ont évalué Seesaw sur des modèles de 150M, 300M et 600M de paramètres entraînés à l'échelle Chinchilla ($D=20N$) en utilisant le code OLMo.

• Parité de performance : Seesaw correspond à la dynamique de perte de validation des programmes de lissage cosinus (cosine annealing) standards pour toutes les tailles de modèles lorsqu'ils sont entraînés à la taille de lot critique.
• Réduction du temps d'exécution : Seesaw permet une réduction significative du temps d'exécution sériel, réduisant le temps de calcul réel (wall-clock time) d'environ 36 %, s'approchant étroitement de la limite théorique dérivée dans l'analyse.
• Robustesse : La méthode fonctionne efficacement avec AdamW et diverses configurations de décroissance de poids (weight decay).
• Limites de l'agressivité : Les expériences variant les paramètres $\alpha$ et $\beta$ confirment la contrainte théorique. Les programmes plus agressifs que $\alpha = \sqrt{\beta}$ (par exemple, doubler la taille du lot tout en gardant le taux d'apprentissage constant ou en le faisant décroître trop lentement) entraînent une divergence ou une performance sous-optimale.
• Limite de la taille de lot critique : Les auteurs observent que la stratégie Seesaw est efficace jusqu'à la taille de lot critique. Au-delà de ce point, où le terme de bruit ne domine plus la norme du gradient, l'équivalence se brise, et Seesaw est moins performant qu'une décroissance cosinus standard.

Signification et affirmations

Le papier affirme fournir une base théorique rigoureuse pour la pratique de la « rampe de lot », la faisant passer d'un ajustement heuristique à une stratégie d'optimisation principielle. En démontrant que la décroissance du taux d'apprentissage et la montée en puissance de la taille des lots sont interchangeables sous certaines conditions, Seesaw offre une méthode pratique pour accélérer le pré-entraînement des LLM sans compromettre la qualité du modèle.

Les auteurs soulignent que Seesaw sert de remplacement direct pour les planificateurs existants dans les optimiseurs adaptatifs. La principale importance réside dans la capacité de réduire le temps de calcul réel de l'entraînement des grands modèles d'environ un tiers, répondant ainsi au goulot d'étranglement matériel où les besoins d'entraînement dépassent les améliorations matérielles. Le travail clarifie également les limites de cette stratégie, notant qu'elle repose sur le régime dominé par la variance et peut ne pas s'appliquer une fois que les tailles de lots dépassent le seuil critique où les termes de gradient moyen deviennent dominants.