$μ$LO: Compute-Efficient Meta-Generalization of Learned Optimizers

Dit paper introduceert een computatie-efficiënte meta-trainingmethode voor geleerde optimalisatoren die, gebaseerd op Maximal Update Parametrization ($\mu$P), aanzienlijk betere generalisatie toont naar bredere, diepere en langere trainingsopgaven dan standaard methoden binnen hetzelfde rekenbudget.

De kern

µLO: De Slimme Optimist die Alles Kan Leren (Zelfs Grotere Dingen)

Stel je voor dat je een groep jonge studenten (neuronale netwerken) wilt leren wiskunde oplossen. Normaal gesproken gebruik je een vaste, handgeschreven methode om ze te helpen, zoals een standaardrekenmachine of een bekende formule (zoals Adam of SGD). Deze methoden werken goed, maar ze zijn een beetje stijf. Ze weten niet hoe ze zich moeten aanpassen als de opgave plotseling veel moeilijker wordt.

Het Probleem: De "Grote" Opgave
In de wereld van kunstmatige intelligentie proberen wetenschappers nu een slimmere methode te bouwen: een "geleerde optimizer". Dit is eigenlijk een kleine AI die zelf leert hoe je andere AI's moet trainen. Het idee is prachtig: deze kleine AI leert van duizenden kleine wiskundepuzzels en wordt dan een meester-coach.

Maar er is een groot probleem: zodra deze meester-coach wordt ingezet op een veel grotere puzzel (bijvoorbeeld een netwerk dat veel breder of dieper is dan de oefeningen die hij zag), raakt hij in paniek. Hij probeert te hard, de cijfers exploderen en hij faalt. Het is alsof je een coach die alleen kleine kinderen heeft getraind, plotseling vraagt om een team van professionele voetballers te trainen. Hij weet niet hoe hij moet omgaan met de grootte van het team.

De Oplossing: De "µP" Recept
De auteurs van dit paper hebben een oplossing bedacht die ze µLO noemen. Ze hebben een nieuw "recept" (een manier van instellen) bedacht dat ze µP (Maximal Update Parametrization) noemen.

Om dit te begrijpen, gebruik een analogie met bouwen:

• De Oude Manier (SP): Stel je voor dat je een huis bouwt. Als je de breedte van het huis verdubbelt, verdubbel je simpelweg het aantal bakstenen. Maar de mortel (de optimizer) die je gebruikt, is niet aangepast. Als het huis te groot wordt, zakt de fundering in elkaar of stort het dak in. De "mortel" is niet sterk genoeg voor de nieuwe schaal.
• De Nieuwe Manier (µLO met µP): Met het nieuwe recept pas je de mortel aan voordat je begint. Je zorgt ervoor dat de kracht van de mortel perfect past bij de grootte van het huis, of het nu klein of gigantisch is. Je past ook de manier aan waarop je de bakstenen legt. Hierdoor blijft het huis stabiel, ongeacht of je 100 of 10.000 bakstenen gebruikt.

Wat hebben ze ontdekt? (De Magische Effecten)
De onderzoekers hebben getoond dat hun nieuwe coach (µLO), die is getraind op kleine, simpele netwerken, verrassend goed presteert op veel zwaardere taken:

1. Breder zijn: Hij kan netwerken trainen die 8 keer breder zijn dan wat hij ooit heeft gezien. De oude methoden gaven hier direct de geest.
2. Dieper zijn: Zelfs als de netwerken 5 keer dieper zijn (meer lagen), blijft de nieuwe coach stabiel. Dit is verrassend, want het recept was eigenlijk alleen bedoeld voor "breedte", maar het werkt ook voor "diepte".
3. Langer trainen: Hij kan netwerken trainen gedurende 25 keer zo lang als de oefenperiode. De oude methoden werden hier onstabiel en gaven het op.

Waarom is dit belangrijk?
Vroeger moest je om een optimizer te maken die op grote schaal werkt, duizenden TPU's (supercomputers) maandenlang laten rekenen op enorme hoeveelheden data. Dat is extreem duur en onbetaalbaar voor de meeste mensen.

Met deze nieuwe methode (µLO) kunnen ze een optimizer trainen op kleine, goedkope computers (slechts een paar uur op een gewone GPU), en diezelfde optimizer werkt daarna perfect op de grootste, duurste modellen. Het is alsof je een kind leert fietsen op een kleine fiets in de tuin, en het kind kan daarna direct een motorfiets besturen zonder te vallen.

Kort samengevat:
Deze paper laat zien dat je door slimme aanpassingen in de "bouwregels" van een AI-optimizer, kunt voorkomen dat hij in paniek raakt bij grote taken. Je kunt een slimme coach trainen op kleine taken, en hij zal verrassend goed presteren op enorme, complexe taken, zonder dat je daarvoor een fortuin aan rekenkracht hoeft uit te geven. Het is een stap in de richting van AI die echt "leert hoe het moet", in plaats van alleen maar te herhalen wat hij heeft gezien.

Technische samenvatting

Titel: µLO: Compute-efficiënte meta-generalisatie van geleerde optimaliseerders

Publicatie: ICLR 2026 (voorgesteld)
Auteurs: Benjamin Thérien, Charles-Étienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, Eugene Belilovsky.

---

1. Het Probleem

Geleerde optimaliseerders (Learned Optimizers of LOs) zijn neurale netwerken die zijn getraind om andere neurale netwerken (de "optimizees") te trainen. Hoewel ze het potentieel hebben om de trainingsduur van deep learning-modellen aanzienlijk te verkorten, kampen ze met een fundamenteel probleem: meta-generalisatie.

• Generalisatieprobleem: Bestaande LOs (zoals VeLO) presteren goed op taken die lijken op hun meta-trainingstaken, maar falen vaak wanneer ze worden toegepast op ongezette taken (out-of-distribution).
• Specifieke beperkingen:
  • Ze hebben moeite om modellen te optimaliseren die breder (meer hidden units) of dieper (meer lagen) zijn dan die tijdens de meta-training zijn gezien.
  • Ze generaliseren slecht naar langere trainingshorizons (meer iteraties dan de "unroll length" tijdens meta-training).
• Oorzaak: Traditionele parametrisatie (Standard Parametrization - SP) zorgt ervoor dat de activaties en updates in bredere netwerken instabiel worden (ze "blazen op" of divergeren), wat de convergentie verhindert.

2. Methodologie

De auteurs introduceren µLO, een aanpak die de Maximal Update Parametrization (µP) toepast op bestaande state-of-the-art geleerde optimaliseerders.

A. Toepassing van µP op Geleerde Optimaliseerders

µP is een parametrisatieregel die hyperparameters (zoals leersnelheden) toestaat om te worden overgedragen naar modellen van verschillende groottes zonder opnieuw te hoeven tunen. De auteurs hebben dit theoretisch afgeleid voor twee populaire LO-architecturen:

1. small_fc_lopt: Een feed-forward netwerk dat per-parameter updates berekent.
2. VeLO: Een architectuur die een LSTM-hypernetwerk gebruikt om de parameters van de optimizer te genereren.

De µP-regels voor LOs:
Om µP te bereiken, moeten drie aspecten van het optimizee-netwerk en de optimizer worden aangepast:

1. Initialisatie: Gewichten in verborgen lagen worden geïnitieerd met variatie $1/FAN_IN$, terwijl uitgaande lagen variatie 1 hebben.
2. Pre-activatie vermenigvuldiging: De pre-activaties van uitgaande lagen worden vermenigvuldigd met $1/FAN_IN$ tijdens de forward pass.
3. Update-schaal: De update van de optimizer ($w_t = w_{t-1} - \alpha \dots$) wordt herschaald. Voor verborgen lagen wordt de update vermenigvuldigd met $1/FAN_IN$om te voorkomen dat de updates te groot worden naarmate de breedte ($n$) toeneemt.

B. Meta-training Recept

De auteurs stellen een eenvoudige meta-trainingstrategie voor:

• Train de µLOs op een distributie van MLP-taken met verschillende breedtes (bijv. 128, 512, 1024).
• Dit in tegenstelling tot eerdere methoden die vaak op één specifieke breedte trainden of extreem veel compute vereisten (zoals VeLO-4000 met 4000 TPU-maanden).
• De meta-training gebeurt met een budget dat vergelijkbaar is met standaard parametrisatie (geen extra compute-kosten).

3. Belangrijkste Bijdragen

1. Theoretische Afleiding: De auteurs hebben de juiste µP-parametrisatie afgeleid voor small_fc_lopt en VeLO en bewezen dat deze voldoet aan de µP-eisen (stabiliteit van activaties en maximale updates in de oneindig-brede limiet).
2. Systematische Evaluatie: Ze hebben een uitgebreide testset opgezet om generalisatie te testen op variabelen zoals breedte, diepte, dataset en trainingsduur.
3. Empirische Resultaten: Ze tonen aan dat µLOs aanzienlijk beter generaliseren dan SP-LOs en zelfs handmatig ontworpen optimaliseerders (zoals AdamW) die specifiek voor elke taak zijn getuned.

4. Resultaten

De experimenten tonen overtuigende resultaten op diverse uitdagingen:

• Generalisatie naar bredere netwerken:

  • µLOs presteren soepel op netwerken tot 8x breder dan de meta-trainingstaken (bijv. van 1024 naar 8192 hidden units).
  • Standaard LOs (SP) divergeren vaak al bij 2x of 4x de breedte.
  • µLOs behalen een lagere eindloss dan handmatig getuned AdamW en µAdam op deze grote, ongezette taken.

• Generalisatie naar diepere netwerken (Onverwacht):

  • Hoewel µP theoretisch primair voor breedte is ontworpen, generaliseren µLOs verrassend goed naar netwerken 5x dieper dan tijdens meta-training (van 3 lagen naar 16 lagen).
  • SP-LOs divergeren vaak direct bij diepere netwerken, terwijl µLOs stabiel blijven. De auteurs hypothetiseren dat dit komt door de stabilisatie van pre-activaties.

• Generalisatie naar langere trainingshorizons:

  • µLOs kunnen trainen voor 25x meer stappen dan de maximale "unroll length" tijdens meta-training (1000 stappen) zonder te divergeren.
  • SP-LOs falen vaak na 1000-8000 stappen.

• Efficiëntie:

  • Deze verbeteringen worden bereikt met een zeer laag compute-budget (µLOM is getraind op slechts ~100 GPU-uren), in tegenstelling tot eerdere grote LO's die duizenden TPU-maanden nodig hadden.

5. Betekenis en Conclusie

Dit werk is een doorbraak in het veld van "Learning to Optimize" (L2O). Het bewijst dat:

1. Stabiliteit is cruciaal: Het toepassen van µP op geleerde optimaliseerders lost het fundamentele stabiliteitsprobleem op dat generalisatie naar grotere modellen belemmert.
2. Kostenefficiëntie: Het is niet nodig om enorme hoeveelheden data of compute te gebruiken om een generaliserende optimizer te bouwen; een slimme parametrisatie en een eenvoudige trainingsrecept volstaan.
3. Praktische toepasbaarheid: µLOs bieden een haalbare weg naar het creëren van universele optimaliseerders die kunnen worden ingezet op grote, ongezette taken (zoals grote taalmodellen of vision transformers) zonder dat er per taak hyperparameters hoeven te worden getuned.

De auteurs concluderen dat µLOs een veelbelovende richting vormen voor de ontwikkeling van echt universele, compute-efficiënte optimaliseerders voor de toekomst van deep learning.