Technische Samenvatting: µpscaling Kleine Modellen: Geprincipeerde Warm Starts en Hyperparameter Transfer

Probleemstelling

Moderne deep learning-workflows omvatmen doorgaans het trainen van modelfamilies op meerdere schalen. Kleine modellen worden gebruikt voor snelle prototyping en scaling law-analyse, terwijl grotere modellen vanaf nul worden getraind voor uiteindelijke implementatie om te voldoen aan diverse hardware- en latentiebeperkingen. Deze workflow is computationeel verspillend, omdat de kennis die tijdens het trainen van kleinere modellen is verworven, wordt weggegooid.

Hoewel "model upscaling" (het initialiseren van grotere modellen vanuit getrainde kleinere modellen) is voorgesteld om deze verspilling te beperken, kampen bestaande methoden (bijv. Net2Net) met twee kritieke beperkingen:

1. Hyperparametergevoeligheid: Upscaled modellen vereisen vaak hyperparameters (met name leersnelheden en ruisniveaus) die aanzienlijk verschillen van het basismodel of modellen die vanaf nul zijn getraind. Het afstemmen hiervan op de doelgrote schaal is evenredig duur.
2. Gebrek aan Theoretische Onderbouwing: Eerdere benaderingen vertrouwen op empirische heuristieken. Het blijft onduidelijk of het afstemmen van hyperparameters op kleinere modellen en het extrapoleren via scaling laws geldig is in de context van upscaling, aangezien de trainingsdynamiek van een upscaled model verschilt van zowel het basismodel als een vanaf nul getraind groot model.

Methodologie

De auteurs stellen een theoretisch onderbouwde framework voor µpscaling voor, die functie-behoudende gewichtstransformaties combineert met de Maximal Update Parametrization (µP) om efficiënte hyperparameter transfer mogelijk te maken.

1. Dynamische Equivalentie

De kern van de theoretische bijdrage is de uitbreiding van "statische equivalentie" (waarbij een verbreed model hetzelfde resultaat produceert als een basismodel bij initialisatie) naar dynamische equivalentie.

• Mechanisme: Door gewichten te dupliceren en specifieke breedte-afhankelijke rescalings toe te passen (gebaseerd op de ⊗ Kronecker-product met matrices van enen), kan een breder model worden geconstrueerd dat exact dezelfde functie parametriseert als het basismodel.
• Optimizer Coördinatie: Cruciaal is dat de auteurs bewijzen dat als optimizer-hyperparameters (leersnelheid, weight decay en hulpparameters zoals $\epsilon$) worden geschaald volgens specifieke regels afgeleid van de homogeniteit van de update-functie, het bredere model een identiek trainingspad volgt in de functieruimte als het basismodel.
• Generaliteit: Deze dynamische equivalentie geldt voor algemene architecturen (uitdrukbaar in het Ne⊗or⊤ program framework, inclusief MLPs, ResNets en Transformers) en algemene entrywise optimizers (SGD, Adam, AdamW).

2. Het Upscaling Algoritme

Om de extra capaciteit van het bredere model te benutten, introduceren de auteurs een "copy-and-perturb" procedure:

1. Verbreding (Widening): Construeer een dynamisch equivalent breder model door gewichten en optimizer-states te dupliceren en te herschalen.
2. Symmetrie-doorbreking (Symmetry Breaking): Injecteer een kleine hoeveelheid ruis ($\Delta$) in de parameters. Zonder ruis blijft het model beperkt tot de laagdimensionale subspace van het basismodel. De ruis stelt het model in staat om deze subspace te verlaten en de toegenomen capaciteit te exploiteren.
3. µP Scaling: De ruisgrootte en de optimizer-hyperparameters voor het upscaled model worden geschaald volgens µP-regels. Dit zorgt ervoor dat de trainingsdynamiek in de oneindige-breedte limiet "optimaal" is, waarbij het signaal van het basismodel behouden blijft terwijl de geïnjecteerde ruis betekenisvol bijdraagt zonder te verdwijnen of te exploderen.

3. Hyperparameter Transfer

Het framework maakt zero-shot hyperparameter transfer mogelijk. Omdat de oneindige-breedte limiet van de upscaled trainingsdynamiek rigoureus wordt gekarakteriseerd met behulp van gemodificeerde Tensor Programs, kunnen de optimale hyperparameters (leersnelheid en ruisniveau) die zijn afgestemd op een klein "proxy" upscaled systeem ($n_0 \to kn_0$) direct worden overgedragen naar het doel systeem op grote schaal ($n \to kn$).

Belangrijkste Bijdragen

1. Dynamische Equivalentie Theorie: Het artikel stelt vast dat, onder specifieke gewichts- en hyperparameter-rescalings, een upscaled model dynamisch equivalent is aan zijn basismodel gedurende het gehele trainingsproces, en niet alleen bij initialisatie. Dit breidt de bestaande theoretische fundamenten uit om het volledere trainingspad te omvatten.
2. Geprincipeerd Upscaling Algoritme: Een verenigde procedure voor upscaling die gewichtsduplicatie, symmetrie-doorbrekende ruisinjectie en µP-gebaseerde hyperparameter scaling integreert. Dit biedt een rigoureus alternatief voor heuristische methoden zoals Net2Net.
3. Hyperparameter Transfer Garanties: De auteurs bieden het eerste rigoureuze framework voor het transfereren van hyperparameters in de upscaling setting. Door de Tensor Program machinery uit te breiden om mid-training upscaling te accommoderen, bewijzen zij dat optimale hyperparameters kunnen worden afgestemd op kleine, goedkope proxy-systemen en kunnen worden overgedragen naar grootschalige doelwitten, wat de afstemkosten drastisch vermindert.
4. Oneindige-breedte Analyse: Een gemodificeerd Tensor Program framework wordt geïntroduceerd om de oneindige-breedte limiet van upscaled training te karakteriseren, waarbij wordt aangetoond dat pre-activaties convergeren naar i.i.d. blokken met niet-triviale covariantie-structuren, die verschillen van standaard niet-upscaled training.

Experimentele Resultaten

De methode werd geëvalueerd op drie architecturen (MLPs, ResNets, GPT-2) over verschillende datasets (Forest Cover Type, CIFAR-100, FineWeb) en optimizers (SGD, AdamW).

• Prestaties: In alle settings convergeerden upscaled modellen sneller en bereikten ze een lagere uiteindelijke trainingsverlies vergeleken met modellen die vanaf nul zijn getraind bij dezelfde breedte.
• Hyperparameter Transfer: Experimenten bevestigden dat optimale leersnelheden en ruisniveaus die zijn afgestemd op kleine systemen (bijv. breedte $n_0=100$) robuust werden overgedragen naar veel grotere systemen (bijv. breedte $n=500$ of $n=1024$).
• Efficiëntie: Door hyperparameter tuning uit te voeren op een kleine schaal in plaats van op de doelgrote schaal, bereikte de methode aanzienlijke computationele versnellingen in de tuningkosten:
  • 23,6× versnelling voor MLPs ($n/n_0 = 5$).
  • 16,0× versnelling voor ResNets ($n/n_0 = 4$).
  • 48,2× versnelling voor GPT-2 ($n/n_0 = 10$).
• Generalisatie Nuance: Hoewel het trainingsverlies verbeterde, vertoonde de upscaled ResNet op CIFAR-100 een slechtere generalisatie (validatienauwkeurigheid) dan het vanaf nul getrainde model. De auteurs merken op dat dit consistent is met hun theorie, die de trainingsdynamiek beheerst en niet de generalisatie, en dat dit een bekend kenmerk is van overgeparameteriseerde CNN's op beperkte data.

Betekenis en Claims

Het artikel claimt het eerste rigoureuze theoretische framework te bieden voor hyperparameter transfer in model upscaling. De betekenis ligt in:

• Brug tussen Theorie en Praktijk: Het verplaatsen van upscaling van empirische heuristieken naar een geprincipeerde aanpak die geworteld is in oneindige-breedte limiettheorie en Tensor Programs.
• Kostenreductie: Het mogelijk maken van het afstemmen van hyperparameters voor massieve modellen op kleine proxy-systemen, waardoor de exploratie van grootschalige architecturen computationeel haalbaarder wordt.
• Universaliteit: Het framework is toepasbaar op standaard architecturen en optimizers, in tegen tegenstelling tot eerdere methoden die vaak architectuur-specifiek waren.

De auteurs blijven bescheiden over generalisatie en erkennen dat hun theorie zich richt op trainingsdynamiek en dat generalisatie onder upscaling een open vraag blijft. Ze merken ook op dat de methode momenteel alleen van toepassing is op breedte-wijze upscaling, waarbij diepte-wijze upscaling wordt overgelaten aan toekomstig werk.