How to Allocate Your Tokens? Scaling Laws with Training Steps and Batch Size

Résumé Technique : Lois d'Échelle avec les Étapes d'Entraînement et la Taille de Batch

Énoncé du Problème
Les lois d'échelle existantes pour les Grands Modèles de Langage (LLM), telles que la loi Chinchilla, modélisent la perte de test du modèle en fonction de la taille du modèle ($N$) et de la quantité totale de données d'entraînement ($D$). Bien qu'efficaces pour déterminer l'allocation optimale de calcul entre $N$ et $D$, ces lois ne prennent pas explicitement en compte les hyperparamètres d'entraînement, spécifiquement la taille de batch ($b$) et le nombre d'étapes d'entraînement ($K$). Les tentatives précédentes de modélisation de l'échelle des hyperparamètres échouent souvent à s'intégrer aux lois d'échelle basées sur la perte ou reposent sur le concept de « taille de batch critique » qui, dans certaines formulations théoriques, implique une taille de batch optimale de un — une contradiction avec les découvertes empiriques où les tailles de batch optimales évoluent avec le budget de tokens. De plus, l'ajustement de lois d'échelle pour les hyperparamètres optimaux nécessite généralement des balayages exhaustifs de tailles de batch pour diverses quantités de tokens, ce qui engendre des coûts de calcul massifs.

Méthodologie
Les auteurs proposent une loi d'échelle à trois termes (3TL) qui décompose explicitement le budget de données d'entraînement $D$ en une taille de batch en tokens ($M = b \times s$, où $s$ est la longueur de séquence) et en un nombre d'étapes ($K$), de telle sorte que $D = MK$. La forme fonctionnelle proposée pour la perte $L$ est :

$$L(N, M, K) = E + \frac{A}{N^\alpha} + \frac{B}{M^\beta} + \frac{C}{K^\gamma}$$

où $(E, A, B, C, \alpha, \beta, \gamma)$ sont des paramètres ajustés à partir des cycles d'entraînement.

Les auteurs comparent cette 3TL à une approche de loi à deux termes (2TL), où la taille du modèle $N$ est fixée, et la loi est ajustée indépendamment pour chaque $N$ en utilisant la forme $L(M, K) = E + B/M^\beta + C/K^\gamma$.

Les principales étapes méthodologiques incluent :

1. Dérivation de la Taille de Batch Optimale : En minimisant la 3TL par rapport à $M$ sous la contrainte $D=MK$, les auteurs dérivent une taille de batch optimale $M^\star$ qui évolue selon une loi de puissance de $D$ : $M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$.
2. Utilisation des Données : Contrairement aux méthodes précédentes qui nécessitent d'identifier la taille de batch unique optimale pour chaque configuration (via des balayages complets), la 3TL est ajustée à partir de données provenant de tailles de batch sous-optimales. Cela permet d'ajuster la loi de manière robuste avec nettement moins de cycles d'entraînement.
3. Modélisation du Régime Sous-Optimal : Pour répondre aux contraintes matérielles où la taille de batch théorique optimale est irréalisable, les auteurs proposent une procédure d'ajustement en deux étapes. Premièrement, ils utilisent la 3TL pour informer une forme fonctionnelle simplifiée $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$ (en supposant $\gamma \approx \beta$). Deuxièmement, ils ajustent une loi de puissance sur l'intervalle des tailles de batch $\epsilon$-sous-optimales (où $\epsilon$ représente un seuil de gaspillage de calcul de 5 %).
4. Analyse de la Taille de Batch Critique : Les auteurs analysent le comportement du nombre d'étapes $K_{\bar{L}}(b)$ nécessaires pour atteindre une perte cible $\bar{L}$ en fonction de la taille de batch, comparant les prédictions de la 3TL aux observations empiriques concernant la taille de batch critique.

Résultats Clés
Les auteurs ont évalué leur approche sur deux ensembles de données : Li (Li et al., 2025a) et OpenEuroLLM.

• Évolution de la Taille de Batch Optimale : La 3TL récupère avec succès l'évolution de la taille de batch optimale $M^\star$. Sur le jeu de données Li, la dérivation de l'évolution $M^\star \approx 0.667 \cdot D^{0.566}$ s'aligne étroitement avec la « Step-Law » ($M^\star \approx 0.58 \cdot D^{0.571}$) de Li et al. (2025a), qui a été dérivée d'un balayage complet des tailles de batch.
• Efficacité de Calcul : Un résultat principal est que la 3TL peut être ajustée en utilisant seulement 28 % des cycles d'entraînement requis par les méthodes traditionnelles. En masquant le jeu de données pour n'inclure que 2 tailles de batch par combinaison $(N, D)$ (au lieu d'un balayage complet), la 3TL récupère tout de même la bonne loi d'échelle pour $M^\star$. En revanche, l'ajustement direct d'une loi de puissance sur les quelques « meilleures » tailles de batch disponibles (sans modéliser la surface de perte) produit des lois d'échelle hautement instables et inexactes.
• Tailles de Batch Sous-Optimales : La procédure d'ajustement en deux étapes parvient à modéliser l'intervalle des tailles de batch sous-optimales qui engendrent au plus 5 % de gaspillage de calcul. Les auteurs constatent que cet intervalle possède une largeur d'environ $2^2$ (un facteur de 4) en taille de batch, et que son comportement d'évolution est cohérent à travers différentes tailles de modèles et différents jeux de données.
• Taille de Batch Critique : La 3TL modélise correctement le phénomène où le nombre d'étapes pour atteindre une perte cible devient moins sensible aux augmentations de la taille de batch au-delà d'un point critique. Crucialement, le modèle montre que ce comportement critique évolue avec le budget de tokens $D$ mais est largement invariant à la taille du modèle $N$, correspondant aux observations empiriques de Zhang et al. (2025).
• Comparaison avec Chinchilla : Lorsque la 3TL est réduite à une forme Chinchilla standard (en substituant la $M^\star$ optimale dans l'équation de la perte), l'évolution de calcul optimale pour la taille du modèle $N$ impliquée diffère d'un ajustement direct de la loi Chinchilla. La 3TL suggère un exposant d'échelle plus petit pour $D$, soulignant que le choix de la formulation de la loi d'échelle impacte les configurations optimales dérivées.

Signification et Revendications
L'article affirme que la loi à trois termes offre un cadre unifié qui comble le fossé entre les lois d'échelle de type Chinchilla (taille du modèle vs données) et les lois d'échelle des hyperparamètres (taille de batch). Sa signification réside principalement dans trois domaines :

1. Robustesse avec des Données Éparses : La capacité d'ajuster la loi en utilisant des tailles de batch sous-optimales réduit considérablement le coût de calcul de la dérivation des lois d'échelle, rendant possible la détermination des tailles de batch optimales sans balayages exhaustifs.
2. Cohérence Théorique : La formulation est structurellement similaire aux bornes de convergence en optimisation stochastique, fournissant un pont théorique entre l'analyse empirique des échelles et la théorie de l'optimisation.
3. Applicabilité Pratique : Contrairement aux modèles précédents qui impliquent une taille de batch optimale de un ou échouent à décrire les régimes sous-optimaux, la 3TL fournit une taille de batch optimale non triviale et modélise explicitement les performances dans les régimes sous-optimaux, ce qui est crucial pour l'entraînement pratique sous contraintes matérielles.

Les auteurs notent des limites, notamment des incohérences potentielles dans les valeurs des paramètres entre différents jeux de données (par exemple, le terme de perte irréductible $E$) et le fait que la loi nécessite toujours des taux d'apprentissage optimaux pour chaque configuration. Ils reconnaissent également que l'évolution impliquée par la 3TL diffère des ajustements directs de Chinchilla, suggérant qu'elle n'est peut-être pas le seul instrument pour déterminer les tailles de modèles optimales, mais qu'elle est supérieure pour comprendre l'interaction entre les étapes, la taille de batch et les données.