Technische Samenvatting: Over de Niet-lineariteit van de Schaling van de Leersnelheid bij LLM-training

Probleemstelling

Het efficiënt trainen van grote taalmodellen (LLM's) vereist zorgvuldige afstemming van hyperparameters, maar het direct doorzoeken van leersnelheden bij doelschalen is vaak even prohibitief duur als onmogelijk. Daarom vertrouwen beoefenaars op hyperparameter-transfer, waarbij optimale instellingen voor grootschalige modellen worden voorspeld op basis van kleinschaligere runs. Een gangbare aanname in dit domein is dat de optimale leersnelheid ($\eta^*$) een log-lineaire (machtswet) schaleringsrelatie volgt met betrekking tot de datagrootte ($D$) en de modelgrootte ($N$):
$$\log \eta^*(D, N) = a \log D + b \log N + c$$
Echter, bestaande studies slagen er vaak niet in om rigoureus te evalueren of deze log-lineaire aanname globaal geldt of slechts lokaal, en zij laten frequent de computationele kosten weg die nodig zijn om deze schalingswetten aan te passen. Dit artikel onderzoekt de geldigheid van deze schalingswet en de efficiëntie van verschillende extrapolatiestrategieën (data-as versus model-as) en parametrisaties (leersnelheid versus effectieve leersnelheid).

Methodologie

De auteurs voerden een systematische empirische studie uit met GPT-2-achtige architecturen (variërend van 22M tot 707M parameters) getraind op de FineWeb-100B dataset met token-budgetten van 5B tot 100B.

1. Experimenteel Ontwerp:

   • Modellen werden getraind met AdamW met een vaste batchgrootte (0,52M tokens) en een warmup-steady-decay (WSD) schema.
   • Leersnelheden werden gescand over een logaritmisch rooster ($\log_2 \eta \in [-16, -7]$).
   • Optimale Leersnelheid ($\eta^*$): Geschat door een derdegraads polynoom te fitten op de validatieverliesfunctie als functie van $\log \eta$.
   • Effectieve Leersnelheid ($\eta_{eff}$): Gedefinieerd als de stapgrootte in genormaliseerde gewichtsruimte: $\eta_{eff}(t) = \|\hat{w}_{t+1} - \hat{w}_t\|_2$, waarbij $\hat{w} = w/\|w\|_2$. De optimale effectieve leersnelheid ($\eta^*_{eff}$) werd eveneens geschat.

2. Evaluatiemetrieken:

   • Out-of-Distribution $R^2$ ($R^2_{OOD}$): Meet de voorspellende nauwkeurigheid van schalingswetten op ongeziene schalen.
   • Extra Compute Ratio (ECR): Kwantificeert de extra computationele kosten die worden gemaakt wanneer een voorspelde leersnelheid wordt gebruikt in plaats van de werkelijke optimale leersnelheid om hetzelfde validatieverlies te bereiken.
   • Compute-Constrained Extrapolation: De auteurs vergeleken extrapolatie langs de data-as (variërend van $D$ met een vaste grote $N$) versus de model-as (variërend van $N$ met een vaste grote $D$) onder identieke compute-budgetten.

3. Theoretische Analyse:

   • Het artikel analyseert de impliciete dynamiek van AdamW, specifiek de relatie tussen de ruwe leersnelheid $\eta$ en de effectieve leersnelheid $\eta_{eff}$, rekening houdend met het transiënte gedrag van gewichtsnormen.
   • Experimenten werden herhaald met AdamH, een optimizer die $\eta_{eff}$ expliciet controleert, om de effecten van gewichtsnorm-dynamiek te isoleren.

Belangrijkste Resultaten

1. Niet-lineariteit van de Optimale Leersnelheid

De studie concludeert dat de optimale leersnelheid $\eta^*$ niet een globale log-lineaire schalingswet volgt.

• Opwaartse Kromming: Naarmate de datagrootte $D$ en de modelgrootte $N$ toenemen, vertoont $\eta^*$ een systematische opwaartse kromming. Log-lineaire extrapolatie vanuit kleinere schalen onderschat consequent de optimale leersnelheid voor grotere schalen.
• Lineariteit van de Effectieve Leersnelheid: In tegen plaats vertoont de optimale effectieve leersnelheid ($\eta^*_{eff}$) een aanzienlijk meer lineair schalingsgedrag over dezelfde bereiken, met hogere Pearson-correlatiecoëfficiënten.

2. Superieuriteit van Data-as Extrapolatie

Bij het vergelijken van extrapolatiestrategieën onder vaste compute-budgetten:

• Data-as versus Model-as: Het fitten van schalingswetten langs de datagrootte-as ($D$) levert aanzienlijk nauwkeurigere voorspellingen op dan langs de modelgrootte-as ($N$).
• Compute Efficiëntie: Het gebruik van data-as extrapolatie met de effectieve leersnelheid brengt slechts ~2% extra compute met zich mee ten opzichte van het gebruik van de werkelijke optimale leersnelheid. Dit is minder dan 20% van de compute die nodig is voor een volledige hyperparameter-sweep op de doelschaal.

3. Mechanistische Verklaring

Het artikel schrijft de niet-lineariteit in $\eta^*$ toe aan de transiënte dynamiek van gewichtsnormen in AdamW:

• Twee Regimes: De relatie tussen $\eta$ en $\eta_{eff}$ hangt af van de vraag of de gewichtsnorm $\|w_t\|$ het evenwicht heeft bereikt.
  • Evenwichtsregime: $\eta_{eff} \propto \sqrt{\eta}$ (helling 2 in log-log ruimte).
  • Pre-evenwichtsregime: $\eta_{eff} \propto \eta$ (helling 1 in log-log ruimte).
• Impact van Schaal: Naarmate de optimale leersnelheid afneemt voor grotere modellen/data, blijft de trainingshorizon vaak binnen het pre-evenwichtsregime (waar de gewichtsnormen nog niet tot rust zijn gekomen). Dit verschuift de effectieve stapgrootte-dynamiek, wat de geobserveerde opwaartse kromming in $\eta^*$ veroorzaakt.
• Validatie via AdamH: Wanneer AdamH wordt gebruikt om $\eta_{eff}$ expliciet te controleren, wordt de schalingswet strikt log-lineair, wat bevestigt dat de niet-lineariteit in standaard AdamW voortkomt uit de impliciete, tijdvariërende mapping tussen $\eta$ en $\eta_{eff}$.

Betekenis en Claims

Het artikel claimt drie primaire bijdragen:

1. Empirische Correctie: Het toont aan dat de standaard log-lineaire schalingswet voor leersnelheden slechts een lokale eigenschap is en faalt bij grotere schalen door opwaartse kromming, wat leidt tot onnauwkeurige extrapolatie.
2. Praktische Transferstrategieën: Het identificeert twee specifieke keuzes die hyperparameter-transfer aanzienlijk verbeteren:
   • Het parametriseren van de schalingswet met behulp van de effectieve leersnelheid ($\eta_{eff}$) in plaats van de ruwe leersnelheid.
   • Het extrapoleren langs de data-as ($D$) in plaats van de modelgrootte-as ($N$).
3. Theoretisch Inzicht: Het biedt een mechanistische verklaring voor de geobserveerde niet-lineariteit, door deze te koppelen aan het transiënte gedrag van gewichtsnormen in AdamW. De auteurs stellen dat $\eta_{eff}$ de meer fundamentele grootheid is die de optimalisatiedynamiek beheerst in schaal-invariante architecturen.

De auteurs vermelden beperkingen, waaronder de focus op GPT-2-architecturen en AdamW, de afhankelijkheid van cubic fitting voor de schatting van $\eta^*$, en het feit dat de analyse beperkt is tot schalen tot 100B tokens en 707M parameters. Zij beweren niet dat deze bevindingen automatisch generaliseren naar alle architecturen of optimizers zonder verdere verificatie.