Technische Samenvatting: Training voor het Model dat u Teruggeeft

Probleemstelling

Moderne Large Language Model (LM) pipelines leveren vaak een gemiddeld model terug—meestal een Exponential Moving Average (EMA) van de trainingsiteraties—in plaats van de laatste iteratie zelf. Hoewel iteratie-gemiddelden bekend staan om het stabiliseren van optimalisatie en het verbeteren van generalisatie, blijft een fundamentele vraag bestaan: Gegeven dat de uiteindelijke output een iteratie-gemiddelde is, hoe moet het trainingsproces zelf worden aangepast om de prestaties van deze specifieke estimator te optimaliseren?

Eerder werk heeft dit impliciet geadresseerd door te observeren dat middeling agressievere leersnelheden toestaat. Dit paper neemt echter een directe aanpak: het streeft naar het ontwerpen van een optimalisatie-algoritme dat de trainings-traject expliciet controleert om de fout van de teruggegeven iteratie-gemiddelde estimator te minimaliseren.

Methodologie

Theoretische Formulering

De auteurs formuleren het probleem als een optimal control-probleem binnen een continue-tijd stochastisch kwadratisch model.

• Dynamiek: Het parametertraject $\theta_t$ volgt een stochastische differentiaalvergelijking (SDE) die ruizige gradiëntafdaling nabij een kwadratisch minimum representeert (Ornstein–Uhlenbeck proces), aangevuld met een control-input $u_t$.
• Doelstelling: Het doel is om een control-strategie $u_t$ te vinden die de gemiddelde kwadratische fout minimaliseert tussen de teruggegeven tijdgemiddelde estimator $\hat{\theta}_{EMA, T}$ en het ware optimum $\mu^*$, onderworpen aan een kwadratische straf op de grootte van de control-interventie (regularisatieparameter $\lambda$).
• Afleiding: Met behulp van Linear-Quadratic-Gaussian (LQG) control-theorie en het principe van certainty equivalence, leiden de auteurs een optimale controller af. De exacte oplossing hangt af van de posterior mean van het optimum gegeven het traject, wat onhandelbaar is voor grootschalige LMs.

Praktische Benadering: PACE

Om de oplossing implementeerbaar te maken, introduceren de auteurs een reeks heuristische benaderingen:

1. Estimator Substitutie: Vervang de onbekende posterior mean van het optimum door het huidige iteratie-gemiddelde ($\hat{\theta}_{EMA, t}$).
2. Tijd-benadering: Neem aan dat de huidige tijd $t$ veel kleiner is dan de totale trainingstijd $T$ ($t \ll T$), wat de tijdsafhankelijke coëfficiënten vereenvoudigt.
3. Discretisering en Preconditionering: Discretiseer de continue control en vervang de onbekende Hessiaan door de diagonale preconditioner die wordt gebruikt in AdamW (de tweede-moment schatting $v_t$).

Dit resulteert in PACE (Pullback Averaging Control for Efficient Optimization), een lichtgewicht wrapper rondom AdamW. De update-regel voegt een "pullback"-term toe die de live gewichten richting de EMA-gewichten beweegt:
$$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$$
waarbij $C$ een diagonale matrix is afgeleid van de Adam-preconditioner, en de kracht van de pullback wordt afgeknepen om te garanderen dat de update een convexe combinatie blijft van de standaard stap en de pullback.

Belangrijkste Bijdragen

1. Optimal Control Formulering: Het paper biedt de eerste principiële afleiding van een optimalisatie-algoritme dat specifiek is ontworpen om de fout van een iteratie-gemiddelde estimator te minimaliseren, in plaats van de laatste iteratie.
2. Theoretische Garanties:
   • Convex Setting: De auteurs bewijzen dat een gestileerde versie van PACE convergeert naar het minimum van elke convexe loss met unbiased stochastische gradiënten met de standaard $O(1/\sqrt{T})$ snelheid, tot een constante factor afhankelijk van de EMA-sterkte. Dit zorgt ervoor dat de methode nooit significant slechter is dan standaard SGD.
   • Quadratic Setting: In de kwadratische setting die gebruikt is voor de afleiding, verbeteren de gecontroleerde dynamieken de limiterende kwadratische fout van de iteratie-gemiddelde estimator strikt ten opzichte van ongecontroleerde dynamieken (standaard SGD met EMA). De verbetering kan arbitrair groot zijn op bepaalde kwadratische instanties.
3. Empirische Validatie: De methode wordt geëvalueerd op:
   • Fine-tuning: SmolLM2-1.7B, Qwen3-1.7B, en Gemma3-1B modellen op de smol-smoltalk dataset.
   • Pretraining: GPT-2 (124M) op FineWeb bij de Chinchilla-optimale token budget.
   • Baselines: PACE wordt vergeleken met vanilla AdamW, AdamW met EMA evaluatie, en het Schedule-Free algoritme.

Resultaten

• Prestaties: PACE verbetert strikt beide: vanilla AdamW en AdamW met EMA evaluatie over alle geteste modellen (1B–2B parameters) en datasets.
• Robuustheid: De verbetering is robuust over een breed scala aan leersnelheden, decay-schema's (constant, cosine, Warmup-Stable-Decay) en hyperparameters (pullback sterkte $c$, EMA power $\kappa$).
• Vergelijking met Schedule-Free: PACE is competitief met het Schedule-Free algoritme (Defazio et al., 2024), dat ook iteratie-gemiddelen gebruikt om de noodzaak voor leersnelheid-decay te verwijderen. Echter, PACE bereikt dit met een andere update-regel.
• Leersnelheid Decay: PACE presteert goed, zelfs met constante leersnelheden, wat suggereert dat het training kan stabiliseren in regimes waar leersnelheid-decay normaal gesproken vereist is. In token-budget vergelijkingen volgt de constante-LR traject van PACE de "river floor" van de loss landscape even effectief als WSD-takken, en presteert vaak zelfs beter.

Betekenis en Beperkingen

Betekenis:
Het paper demonstreert dat het optimaliseren voor het teruggegeven model (het gemiddelde) in plaats van de laatste iteratie leidt tot tastbare prestatiewinsten. Door het trainings-traject te behandelen als een control-probleem, biedt PACE een theoretisch onderbouwde methode om training te stabiliseren en generalisatie te verbeteren zonder complexe schema-tuning te vereisen. De auteurs merken op dat PACE een verbeterde theoretische convergentie biedt in de kwadratische setting vergeleken met standaard methoden, hoewel dit gepaard gaat met een iets tragere theoretische garantie in de algemene convexe setting door de EMA-factor.

Beperkingen:

• Theoretische Aannames: De afleiding rust op een vereenvoudigd kwadratisch model met additieve ruis. Hoewel convergentie bewezen is voor algemene convexe functies, blijft de rigoureuze incorporatie van momentum en adaptieve preconditionering (centraal in AdamW) in de theoretische analyse een openstaande richting.
• Schaal: Experimenten waren beperkt tot modellen tot 2B parameters voor fine-tuning en 124M voor pretraining. De effectiviteit op grotere schaal (bijv. 70B+) is niet geverifieerd.
• Geheugen Overhead: PACE vereist het onderhouden van een extra kopie van de modelgewichten (de EMA staat) tijdens de training, wat een niet-verwaarloosbare geheugen overhead introduceert, hoewel dit in de praktijk wordt gemitigeerd door grote batch sizes.
• Seed Sensitiviteit: Vanwege rekenkracht-beperkingen gebruikten de hoofdruns een enkele seed, hoewel studies naar multi-seed robuustheid de ordeningsstabiliteit van de methode bevestigden.

De auteurs concluderen dat PACE een veelbelovende stap vormt naar het afstemmen van optimalisatie-algoritmen op de werkelijke deployment-strategie van moderne LMs, waarbij het gemiddelde model het eindproduct is.