A Proof of Learning Rate Transfer under $μ$P

Titel: Een Bewijs van Learning Rate Transfer onder µP

Auteur: Soufiane Hayou (Johns Hopkins University)
Context: AISTATS 2026 (voorgesteld)

1. Het Probleem

De prestaties van kunstmatige intelligentie worden momenteel gedreven door het schalen van neurale netwerken (vergroting van breedte en diepte). Een praktische uitdaging bij het schalen is het aanpassen van hyperparameters (zoals de leersnelheid, $\eta$) om numerieke overflows te voorkomen en optimale prestaties te behalen.

• Learning Rate Transfer: Dit is het fenomeen waarbij de optimale leersnelheid stabiel blijft naarmate de breedte ($n$) van het netwerk toeneemt. Dit zou het mogelijk maken om hyperparameters te tunen op kleine modellen en ze direct toe te passen op veel grotere modellen zonder extra tuningkosten.
• De Huidige Status: Het is empirisch waargenomen dat deze transfer werkt onder de Maximal Update Parametrization (µP), maar er bestond tot nu toe geen strikt wiskundig bewijs. Onder andere parametriseringen, zoals de Standard Parametrization (SP) of Neural Tangent Parametrization (NTP), schuift de optimale leersnelheid vaak naar nul of divergeert deze, waardoor transfer niet mogelijk is.
• De Vraag: Kunnen we wiskundig bewijzen dat de optimale leersnelheid convergeert naar een niet-nul constante onder µP, en waarom faalt dit onder andere methoden?

2. Methodologie

De auteur analyseert een lineaire Multi-Layer Perceptron (MLP) met oneindige breedte.

• Model: Een lineair netwerk $f(x) = V^\top W_L \dots W_1 W_0 x$ getraind met Gradient Descent (GD) op een kwadratische loss-functie.
• Parametrisatie:
  • µP: De initiatie-variabiliteit en leersnelheid worden geschaald om "maximale feature learning" te garanderen in de oneindige breedte-limiet. Specifiek: $W_\ell \sim \mathcal{N}(0, n^{-1})$, $V \sim \mathcal{N}(0, n^{-2})$, en $\eta$ is constant (voor GD).
  • SP: Standaard initiatie waarbij $V \sim \mathcal{N}(0, n^{-1})$ en $\eta$ constant is.
• Analyse-aanpak:
  1. Polynoomrepresentatie: De auteur observeert dat de output van het netwerk na $t$ trainingsstappen, $f^{(t)}(x)$, kan worden uitgedrukt als een polynoom in de leersnelheid $\eta$.
  2. Asymptotische Analyse: Er wordt onderzocht hoe de coëfficiënten van dit polynoom gedragen naarmate de breedte $n \to \infty$.
  3. Convergentie van het Argmin: De optimale leersnelheid $\eta^*_n$ is het punt dat de loss-functie minimaliseert. Het bewijs toont aan dat de loss-functie convergeert naar een deterministische limietfunctie, en dat het minimum van deze limietfunctie een niet-nul constante is.

3. Belangrijkste Bijdragen

1. Eerste Rigoureuze Bewijzen: Dit paper levert het eerste strikte wiskundige bewijs voor learning rate transfer in diepe lineaire netwerken onder µP.
2. Convergentie naar een Niet-Nul Constante: Het wordt bewezen dat onder µP de optimale leersnelheid $\eta^*_n$ convergeert in kans (en bijna zeker) naar een deterministische constante $\eta^*_\infty > 0$ naarmate $n \to \infty$.
3. Wiskundige Karakterisering:
   • Voor $t=1$ (één trainingsstap) wordt de optimale leersnelheid expliciet afgeleid als $\eta^*_\infty = \frac{m}{L} \frac{y^\top K y}{\|Ky\|^2}$, waarbij $K$ de genormaliseerde input Gram-matrix is.
   • De convergentiesnelheid wordt bepaald als $O(n^{-1/2})$.
4. Contrast met SP en NTP: Het paper bewijst dat onder Standard Parametrization (SP) de optimale leersnelheid convergeert naar 0 ($\eta^*_n \to 0$), wat betekent dat er geen transfer plaatsvindt. Onder NTP divergeert de leersnelheid of schuift deze ook weg, wat leidt tot het "lazy training" regime zonder feature learning.
5. Generalisatie naar Stap $t$: Hoewel de analyse complexer wordt voor $t \geq 2$ (waarbij hogere-orde termen in $\eta$ niet langer verwaarloosbaar zijn), wordt bewezen dat de optimale leersnelheid onder milde voorwaarden toch convergeert naar een niet-nul constante.

4. Resultaten

• Theoretische Resultaten:
  • Theorema 1: Bewijst LR transfer bij $t=1$ onder µP met een convergentiesnelheid van $O(n^{-1/2})$.
  • Theorema 2: Bewijst het falen van LR transfer onder SP (convergentie naar 0).
  • Theorema 3: Generaliseert het bewijs naar willekeurige trainingsstappen $t$, mits de limiet-loss-functie een uniek minimum heeft.
• Empirische Validatie:
  • Simulaties met lineaire MLP's van verschillende breedtes ($n=128$ tot $8192$) tonen aan dat de trainingsloss-curves onder µP voor verschillende breedtes bijna identiek zijn en een gemeenschappelijk optimum hebben.
  • Onder SP verschuift het optimum duidelijk naar lagere leersnelheden naarmate het netwerk breder wordt.
  • Experimenten met niet-lineaire netwerken (ReLU), diepere netwerken ($L=3, 9, 27$) en de Adam-optimizer bevestigen dat het fenomeen ook buiten de strikte lineaire theorie empirisch geldt.

5. Significantie en Implicaties

• Theoretisch Fundament: Het paper sluit een belangrijke theoretische lacune door het intuïtieve concept van "feature learning" onder µP te koppelen aan een wiskundig bewijs van hyperparameter-stabiliteit.
• Praktische Toepassing: Het bevestigt de waarde van µP voor het schalen van modellen. Practici kunnen nu met vertrouwen hyperparameters op kleine modellen bepalen en deze toepassen op grote modellen (bijv. Large Language Models), wat aanzienlijke rekenkosten bespaart.
• Verduidelijking van Parametrisatie: Het onderstreept waarom specifieke initiatie- en schalingsregels (zoals $V \sim n^{-2}$ in µP) cruciaal zijn voor het behoud van leerbaarheid in de oneindige breedte-limiet, in tegenstelling tot traditionele methoden.
• Toekomstig Werk: De auteur suggereert dat de bewijsmethode (polynoom-expansie en Tensor Programs) kan worden uitgebreid naar niet-lineaire netwerken en complexere optimalisatoren, hoewel dit nieuwe technische uitdagingen met zich meebrengt.

Conclusie:
Dit paper levert het eerste rigoureuze bewijs dat learning rate transfer een fundamenteel eigenschap is van de µP-parametrisatie in de oneindige breedte-limiet. Het verklaart waarom µP superieur is voor het schalen van modellen en biedt een wiskundige basis voor het gebruik van "zero-shot" hyperparameter transfer in de praktijk.