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

Technische Samenvatting: Schalingswetten met Trainingsstappen en Batchgrootte

Probleemstelling
Bestaande schalingswetten voor Large Language Models (LLM's), zoals de Chinchilla-wet, modelleren de testverlies van een model als een functie van de modelgrootte ($N$) en de totale hoeveelheid trainingsdata ($D$). Hoewel deze effectief zijn voor het bepalen van de optimale toewijzing van rekenkracht tussen $N$ en $D$, houden deze wetten geen rekening met trainingshyperparameters, specifiek de batchgrootte ($b$) en het aantal trainingsstappen ($K$). Eerdere pogingen om hyperparameter-schaling te modelleren falen vaak in het integreren met verliesgebaseerde schalingswetten of vertrouwen op het concept van de "kritieke batchgrootte", wat in sommige theoretische formuleringen impliceert dat de optimale batchgrootte één is—een tegenstrijdigheid met empirische bevindingen waar optimale batchgroottes schalen met het tokenbudget. Bovendien vereist het fitten van schalingswetten voor optimale hyperparameters vaak uitputtend onderzoek naar verschillende batchgroottes voor diverse tokenbudgetten, wat enorme computationele kosten met zich meebrengt.

Methodologie
De auteurs stellen een drie-term schalingswet (3TL) voor die het trainingsdata-budget $D$ expliciet decomponeert in de batchgrootte in tokens ($M = b \times s$, waarbij $s$ de sequentielengte is) en het aantal trainingsstappen ($K$), zodanig dat $D = MK$. De voorgestelde functionele vorm voor het verlies $L$ is:

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

waarbij $(E, A, B, C, \alpha, \beta, \gamma)$ parameters zijn die gefit zijn vanuit trainingsruns.

De auteurs vergelijken deze 3TL met een twee-term wet (2TL) benadering, waarbij de modelgrootte $N$ vaststaat en de wet onafhankelijk voor elke $N$ wordt gefit met de vorm $L(M, K) = E + B/M^\beta + C/K^\gamma$.

Belangrijke methodologische stappen omvatten:

1. Afleiding van de Optimale Batchgrootte: Door de 3TL te minimaliseren met betrekking tot $M$ onder de restrictie $D=MK$, leiden de auteurs een optimale batchgrootte $M^\star$ af die schaalt als een machtswet van $D$: $M^\star \propto D^{\frac{\gamma}{\beta+\gamma}}$.
2. Data-utilisatie: In tegen tegenstelling tot eerdere methoden die vereisen dat de enkelvoudige optimale batchgrootte voor elke configuratie wordt geïdentificeerd (via volledige sweeps), wordt de 3TL gefit met behulp van data van suboptimale batchgroottes. Dit maakt het mogelijk om de wet robuust te fitten met aanzienlijk minder trainingsruns.
3. Modellering van het Suboptimale Regime: Om hardwarebeperkingen aan te pakken waarbij de theoretisch optimale batchgrootte onhaalbaar is, stellen de auteurs een tweestaps-fittingprocedure voor. Eerst gebruiken ze de 3TL om een vereenvoudigde functionele vorm $L(b) = \tilde{E} + \tilde{A}b^{-\tilde{\alpha}} + \tilde{B}b^{\tilde{\alpha}}$ te informeren (uitgaande van $\gamma \approx \beta$). Ten tweede fitten ze een machtswet op het interval van $\epsilon$-suboptimale batchgroottes (waarbij $\epsilon$ een drempel van 5% rekenverspilling vertegenwoordigt).
4. Analyse van de Kritieke Batchgrootte: De auteurs analyseren het gedrag van het aantal stappen $K_{\bar{L}}(b)$ dat nodig is om een doelverlies $\bar{L}$ te bereiken als functie van de batchgrootte, waarbij de 3TL-voorspellingen worden vergeleken met empirische bevindingen over de kritieke batchgrootte.

Belangrijkste Resultaten
De auteurs hebben hun benadering geëvalueerd op twee datasets: Li (Li et al., 2025a) en OpenEuroLLM.

• Schaal van Optimale Batchgrootte: De 3TL herstelt succesvol de schaling van de optimale batchgrootte $M^\star$. Op de Li-dataset komt de afgeleide schaling $M^\star \approx 0.667 \cdot D^{0.566}$ nauw overeen met de "Step-Law" ($M^\star \approx 0.58 \cdot D^{0.571}$) van Li et al. (2025a), die werd afgeleid van een volledige batchgrootte sweep.
• Computationele Efficiëntie: Een primaire bevinding is dat de 3TL kan worden gefit met slechts 28% van de trainingsruns die traditionele methoden vereisen. Door de dataset te maskeren om slechts 2 batchgroottes per $(N, D)$-combinatie te bevatten (in plaats van een volledige sweep), herstelt de 3TL nog steeds de correcte schalingswet voor $M^\star$. In contrast hiermee resulteert het direct fitten van een machtswet op de weinige beschikbare "beste" batchgroottes (zonder het verliesoppervlak te modelleren) in zeer instabiele en onnauwkeurige schalingswetten.
• Suboptimale Batchgroottes: De tweestaps-fittingprocedure modelleert succesvol het interval van suboptimale batchgroottes die maximaal 5% rekenverspilling veroorzaken. De auteurs vinden dat dit interval een breedte heeft van ongeveer $2^2$ (een factor 4) in batchgrootte, en dat het schalingsgedrag consistent is over verschillende modelgroottes en datasets.
• Kritieke Batchgrootte: De 3TL modelleert correct het fenomeen waarbij het aantal stappen om een doelverlies te bereiken minder gevoelig wordt voor toenames in batchgrootte voorbij een kritiek punt. Cruciaal is dat het model laat zien dat dit kritieke gedrag schaalt met het tokenbudget $D$, maar grotendeels invariant is voor de modelgrootte $N$, wat overeenkomt met de empirische observaties van Zhang et al. (2025).
• Vergelijking met Chinchilla: Wanneer de 3TL wordt teruggebracht naar een standaard Chinchilla-vorm (door de optimale $M^\star$ terug te substitueren in de verliesvergelijking), verschilt de geïmpliceerde reken-optimale schaling voor modelgrootte $N$ van een directe fit van de Chinchilla-wet. De 3TL suggereert een kleinere schalingsexponent voor $D$, wat benadrukt dat de keuze van de formulering van de schalingswet de afgeleide optimale configuraties beïnvloedt.

Betekenis en Claims
Het artikel claimt dat de drie-term wet een verenigd kader biedt dat de kloof overbrugt tussen Chinchilla-stijl schalingswetten (modelgrootte vs. data) en hyperparameter-schalingswetten (batchgrootte). De betekenis ligt in drie hoofdgebieden:

1. Robuustheid bij Sparse Data: Het vermogen om de wet te fitten met suboptimale batchgroottes vermindert de computationele kosten voor het afleiden van schalingswetten drastisch, wat het haalbaar maakt om optimale batchgroottes te bepalen zonder uitputtende sweeps.
2. Theoretische Consistentie: De formulering is structureel vergelijkbaar met convergentiegrenzen in stochastische optimalisatie, wat een theoretische brug slaat tussen empirische schalingsanalyse en optimalisatietheorie.
3. Praktische Toepasbaarheid: In tegen tegenstelling tot eerdere modellen die een optimale batchgrootte van één impliceren of falen in het beschrijven van suboptimale regimes, biedt de 3TL een niet-triviale optimale batchgrootte en modelleert het expliciet prestaties in suboptimale regimes, wat cruciaal is voor praktische training onder hardwarebeperkingen.

De auteurs merken beperkingen op, waaronder mogelijke inconsistenties in parameterwaarden over verschillende datasets (bijv. de onherleidbare verliesterm $E$) en het feit dat de wet nog steeds optimale leersnelheden vereist voor elke configuratie. Ze erkennen ook dat de door de 3TL geïmpliceerde reken-optimale schaling verschilt van directe Chinchilla-fits, wat suggereert dat het niet het enige instrument is voor het bepalen van optimale modelgroottes, maar superieur is voor het begrijpen van de interactie tussen stappen, batchgrootte en data.