Technische Samenvatting: Seesaw: Versnellen van training door het balanceren van de leercurve en de batchgrootte-planning

Probleemstelling

Het voortraineren van Large Language Models (LLM's) wordt steeds vaker beperkt door de wandkloktijd (wall-clock time), die vaak loopt tot meerdere maanden voor state-of-the-art modellen. Een primaire strategie om deze duur te verkorten is het vergroten van de batchgrootte, wat parallel rekenen benut om het totale aantal sequentiële optimalisatiestappen te verminderen. Het simpelweg vergroten van de batchgrootte leidt echter tot afnemende meeropbrengsten voorbij een "kritieke batchgrootte" (CBS), waarbij de steekproefefficiëntie daalt.

Hoewel recente grootschalige trainingsruns (bijv. LLaMA, Nemotron, OLMo) "batch ramp"-schema's hebben toegepast — waarbij de batchgrootte gedurende de training geleidelijk wordt vergroot — worden deze strategieën momenteel heuristisch afgestemd. Er is een gebrek aan theoretische rechtvaardiging voor hoe de afname van de leercurve (learning rate decay) en de toename van de batchgrootte (batch size ramp-up) optimaal gebalanceerd moeten worden, met name voor adaptieve optimizers zoals Adam. De centrale vraag die wordt behandeld is: Wat is de optimale batchgrootte-planning om de seriële runtime te minimaliseren zonder de modelprestaties op te offeren?

Methodologie

Theoretisch Kader

De auteurs ontwikkelen een principieel kader gebaseerd op ruisige lineaire regressie om een equivalentie vast te stellen tussen de afname van de leercurve en de toename van de batchgrootte.

1. SGD Equivalentie: Het artikel biedt het eerste niet-asymptotische eind-steekproefbewijs (finite-sample proof) dat aantoont dat voor Stochastic Gradient Descent (SGD) op ruisige lineaire regressie, een proces waarbij de batchgrootte elke fase verdubbelt terwijl de leercurve constant blijft, equivalent is (tot een constante factor in excess risk) aan een proces waarbij de batchgrootte constant is en de leercurve elke fase wordt gehalveerd, mits beide processen hetzelfde totaal aantal datapunten consumeren.
2. Extensie naar Genormaliseerde SGD (NSGD): Om de kloof naar adaptieve optimizers zoals Adam te overbruggen, analyseren de auteurs Genormaliseerde SGD, een hanteerbare analytische proxy voor Adam. Ze introduceren Aanname 2, die stelt dat in het door variantie gedomineerde regime dat in de praktijk wordt waargenomen, de verwachte genormaliseerde gradiëntnormen worden gedomineerd door additieve ruis (schalend als $O(1/B)$). Onder deze aanname reduceert NSGD effectief tot SGD met een geschaalde leercurve.
3. De Equivalentieconditie: De analyse leidt een specifieke relatie af die vereist is om de verliesdynamiek (loss dynamics) te behouden wanneer beide parameters worden aangepast. Voor een leercurve-afnamefactor $\alpha$ en een batchgrootte-toenamefactor $\beta$, blijven de processen equivalent als het product $\alpha\sqrt{\beta}$ constant wordt gehouden.

Het Seesaw-algoritme

Gebaseerd op de theoretische equivalentie, stellen de auteurs Seesaw voor, een scheduler die de leercurve en de batchgrootte dynamisch balanceert:

• Mechanisme: Telkens wanneer een standaard scheduler (bijv. cosine decay) de leercurve zou verlagen met een factor $\alpha$, verlaagt Seesaw in plaats daarvan de leercurve met $\sqrt{\alpha}$ en vergroot de batchg렸ote met $\alpha$.
• Restrictie: De auteurs identificeren een theoretische limiet om divergentie te voorkomen. Het meest agressieve ramp-up schema dat stabiliteit behoudt, wordt gedefinieerd door $\alpha = \sqrt{\beta}$. Het overschrijden hiervan (d.w.z. de batchgrootte te agressief vergroten ten opzichte van de afname van de leercurve) leidt tot divergentie.
• Implementatie: Seesaw is ontworpen als een "drop-in replacement" voor bestaande leercurve-schedulers. In de praktijk benadert het cosine decay via step decay, waarbij de batchgrootte wordt verhoogd op de specifieke token-aantallen waar de leercurve anders zou worden verlaagd.

Belangrijkste Bijdragen

1. Theoretische Equivalentie: Het artikel vestigt de eerste niet-asymptotische equivalentie tussen de afname van de leercurve en de toename van de batchgrootte voor SGD en breidt dit uit naar Genormaliseerde SGD onder een door variantie gedomineerd regime.
2. De Seesaw Scheduler: Introductie van een principieel algoritme dat de standaard leercurve-afname vervangt door een gekoppelde leercurve- en batchgrootte-planning, die theoretisch onderbouwd is om de verliesdynamiek te behouden.
3. Versnellingsanalyse: De auteurs leiden een theoretische bovengrens af voor de reductie van de seriële runtime. Door het meest agressieve stabiele ramp-up schema ($\alpha = \sqrt{\beta}$) te gebruiken tegenover een cosine decay baseline, wordt de maximale theoretische versnelling berekend als $(1 - 2/\pi) \approx 36,3\%$.

Empirische Resultaten

De auteurs evalueerden Seesaw op 150M, 300M en 600M parameter modellen getraind op Chinchilla-schaal ($D=20N$) met gebruik van de OLMo-codebase.

• Prestatie-gelijkwaardigheid: Seesaw komt overeen met de validatieverlies-dynamiek van standaard cosine annealing schema's voor alle modelgroottes wanneer getraind op de kritieke batchgrootte.
• Reductie in Runtime: Seesaw bereikt een significante reductie in seriële runtime, waarbij de wandkloktijd met ongeveer 36% wordt verminderd, wat de afgeleide theoretische limiet dicht nadert.
• Robuustheid: De methode werkt effectief met AdamW en diverse weight decay instellingen.
• Grenzen van Agressiviteit: Experimenten waarbij de $\alpha$- en $\beta$-parameters variëren, bevestigen de theoretische restrictie. Schema's die agressiever zijn dan $\alpha = \sqrt{\beta}$ (bijv. de batchgrootte verdubbelen terwijl de leercurve constant blijft of te traag afneemt) resulteren in divergentie of suboptimale prestaties.
• Kritieke Batchgrootte Limiet: De auteurs observeren dat de Seesaw-strategie effectief is tot de kritieke batchgrootte. Voorbij dit punt, waar de ruisterm niet langer de gradiëntnorm domineert, valt de equivalentie uiteen en presteert Seesaw slechter dan standaard cosine decay.

Betekenis en Claims

Het artikel claimt een rigoureuze theoretische basis te bieden voor de "batch ramp" praktijk, waardoor deze wordt verplaatst van heuristische afstemming naar een fundamentele optimalisatiestrategie. Door aan te tonen dat de afname van de leercurve en de toename van de batchgrootte onder specifieke condities uitwisselbaar zijn, biedt Seesaw een praktische methode om de pretraining van LLM's te versnellen zonder de modelkwaliteit in gevaar te brengen.

De auteurs benadrukken dat Seesaw dient als een drop-in replacement voor bestaande schedulers in adaptieve optimizers. De primaire betekenis ligt in het vermogen om de wandkloktijd van het trainen van grote modellen met ongeveer een derde te verminderen, waarmee de hardware-bottleneck wordt aangepakt waarbij de trainingsvereisten de hardwareverbeteringen inhalen. Het werk verheldert ook de grenzen van deze strategie, waarbij wordt opgemerkt dat het afhankelijk is van het door variantie gedomineerde regime en mogelijk niet van toepassing is zodra de batchgroottes de kritieke drempel overschrijden waar de gemiddelde gradiënttermen gaan domineren.