Technische Samenvatting: GEFEN: OPTIMALISEERDE STOCHASTISCHE OPTIMIZER

Probleemstelling

Moderne deep learning leunt zwaar op de AdamW-optimizer vanwege de robuuste prestaties over diverse architecturen heen. Echter, AdamW houdt bewegende gemiddelden van het eerste en tweede moment bij voor elke parameter, typisch in FP32. Dit resulteert in een geheugenvoetafdruk voor de optimizer-toestand van ongeveer twee keer de omvang van de modelparameters (ongeveer 8 bytes per parameter). Voor grootschalige modellen (bijv. 13B parameters) kan deze toestand alleen al ~97 GiB aan geheugen consumeren, wat de omvang van trainbare modellen, maximale batchgroottes en trainingsdoorvoer op geheugenbeperkte hardware ernstig beperkt.

Bestaande geheugenefficiënte alternatieven kampen met aanzienlijke trade-offs:

• Gekwantiseerde optimizers (bijv. Adam8bit, Adam4bit): Verminderen het geheugen, maar vertrouwen vaak op handmatig ontworte, suboptimale codebooks, vaste blokgroottes (wat hyperparameters introduceert) en zijn mogelijk niet universeel toepasbaar (bijv. het vermijden van embedding-kwantisatie).
• Gestructureerde optimizers (bijv. Adam-mini): Delen statistieken van het tweede moment op basis van handmatige regels die Hessian block-diagonale structuren nabootsen. Deze vereisen echter architectuurspecifieische metadata (bijv. het aantal attention heads) en missen een sterke theoretische rechtvaardiging voor waarom het groeperen dat uitgelijnd is met de Hessian voordelig is.

Het doel is om een optimizer te ontwerpen die de geheugenvoetafdruk aanzienlijk vermindert (verder dan Adam-mini) terwijl de AdamW-prestaties behouden blijven, zonder dat er handmatige configuratie, extra hyperparameters of het herafstemmen van de leersnelheid nodig is.

Methodologie

De voorgestelde methode, Gefen, bereikt geheugenreductie door twee primaire mechanismen: automatische deling van het tweede moment en geleerde kwantisatie van het eerste moment.

1. Theoretische Motivatie

De aanpak is gebaseerd op een theoretisch resultaat (Theorem 4.1) dat aantoont dat voor parameters met grote gemengde Hessian-entries, de ratio van hun gekwadrateerde gradiënten naar één wordt beperkt. Dit impliceert dat parameters met een hoge Hessian-affiniteit van nature vergelijkbare magnitudes hebben in hun gekwadrateerde gradiënten, wat hen ideale kandidaten maakt voor het delen van tweede-moment schattingen. Hoewel het direct berekenen van de Hessian op schaal onuitvoerbaar is, observeren de auteurs dat de structuur van de gekwadrateerde gradiënten in de eerste trainingsstap correleert met deze Hessian-structuur.

2. Automatische Blokpartitionering

In plaats van te vertrouwen op handmatige regels of Hessian-berekeningen, extrapoleert Gefen de optimale blokstructuur uit de gekwadrateerde gradiënten van de eerste optimalisatiestap ($g_0$).

• Algoritme: Voor een platgeslagen parametertensor test het algoritme alle geschikte delers $p$ van de tensoromvang $n$. Het hervormt de gekwadrateerde gradiënten in blokken van grootte $p$ en berekent de heterogeniteit binnen het blok $E(p)$ (de variantie van de gekwadrateerde gradiënten binnen de blokken).
• Selectiestrategie: Het algoritme selecteert de periode $p$ die de eerste duidelijke daling in $E(p)$ oplevert naarmate de blokgrootte toeneemt, waarbij de voorkeur wordt gegeven aan grotere blokken voor geheugenefficiëntie terwijl de gradiënt-gelijkenis gewaarborgd blijft. Als er geen betekenisvolle daling wordt gevonden of als de optimale $p < 8$ is, valt het terug op $p=1$ (geen deling).
• Gebruik: Dit enkele partitioneringsschema wordt gebruikt voor zowel het delen van het tweede moment als het kwantiseren van het eerste moment.

3. Exacte Dynamic Programming Kwantisatie

Om het geheugen van het eerste moment (momentum) te verminderen, gebruikt Gefen een leerbaar, exact kwantisatiecodebook in plaats van vaste of heuristische schema's.

• Codebook Leren: Met behulp van de initiële gradiënten construeert de methode een histogram van genormaliseerde gradiëntwaarden. Vervolgens lost het het optimale scalaire kwantisatieprobleem op met behulp van Dynamic Programming (DP) (geïnspireerd door Wang & Song, 2011) om het exacte optimale codebook voor de gegeven distributie en bitbreedte te vinden.
• Stabiliteit: In tegen tegenstelling tot Lloyd-Max, dat kan lijden onder "contractie" (het krimpen van het kwantisatiebereik over de tijd) en gevoeligheid voor initialisatie, dwingt Gefen de codebook-uiteinden vast op $-1$ en $1$. Dit zorgt ervoor dat het kwantisatiebereik gedurende de hele training behouden blijft.
• Efficiëntie: Het codebook wordt één keer aan het begin van de training geleerd (met behulp van CUDA voor histogramvorming en CPU/Numba voor DP) en wordt gedurende de gehele trainingsrun hergebruikt.

4. Algoritme Flow

Gefen werkt als een drop-in vervanging voor AdamW:

1. Initialisatie: Partitioneer parameters op basis van $g_0$ en leer het kwantisatiecodebook.
2. Update Stap:
   • Bereken gradiënten.
   • Update het eerste moment (momentum) door de dekwantiseerde codebookwaarden toe te passen, momentum toe te passen, opnieuw te kwantiseren en de blokwijze absolute maximum bij te werken.
   • Update het tweede moment door het gemiddelde van de gekwadrateerde gradiënten binnen elk blok te berekenen en een exponentieel bewegend gemiddelde toe te passen.
   • Pas bias-correctie toe en update parameters met de gedeelde tweede-moment en de gekwantiseerde eerste-moment.

Belangrijkste Bijdragen

1. Geheugenefficiëntie: Gefen vermindert de optimizer-geheugenvoetafdruk met ongeveer 8× vergeleken met AdamW (een reductie van ~2.0× parametergrootte naar ~0.25×), wat overeenkomt met een reductie van ~6.5 GiB per miljard parameters.
2. Automatische Adaptatie: De methode vereist geen architectuurspecifieke metadata, handmatige groeperingsregels of extra hyperparameters boven de standaard AdamW-defaults. Het extrapoleert automatisch blokstructuren uit de data.
3. Theoretische Onderbouwing: Biedt een theoretische link tussen Hessian-entries en de ratio van gekwadrateerde gradiënten, wat het delen van tweede-moment statistieken rechtvaardigt voor Hessian-uitgelijnde parameters.
4. Optimale Kwantisatie: Introduceert een exacte, histogram-gebaseerde dynamic programming aanpak voor het kwantiseren van momentum, waardoor de suboptimaliteit en instabiliteit van iteratieve methoden zoals Lloyd-Max wordt vermeden.
5. Prestatie-gelijkwaardigheid: Toont aan dat deze reducties niet ten koste gaan van de modelkwaliteit, waarbij de AdamW-prestaties worden geëvenaard over diverse taken.

Experimentele Resultaten

De auteurs evalueerden Gefen tegenover AdamW, Adam-mini, Adam8bit, Adam4bit, Adafactor en SM3 over CNNs, ResNets, GPT-2 en Llama 3 modellen.

• Pieke Geheugenvoetafdruk: Gefen behaalde consistent de laagste pieke geheugengebruik onder de vergeleken methoden. Voor een model met 1.5B parameters verminderde Gefen de optimizer-toestand van 11.2 GiB (AdamW) naar 1.5 GiB.
• Prestatie (Kwaliteit): Over alle instellingen (OpenWebText, C4, ImageNet, MNIST) behaalde Gefen validatieverliezen en accuraatheid die vergelijkbaar zijn met AdamW. Opvallend genoeg presteerde het beter dan Adam4bit en Adam-mini in verschillende metrieken terwijl het dezelfde leersnelheid-schema behield.
• Doorvoer: De verminderde geheugenvoetafdruk maakte grotere microbatches mogelijk in gedistribueerde training:
  • FSDP (Llama 3-1.5B): Gefen maakte een microbatch-grootte van 2 mogelijk (vs. 1 voor anderen), wat resulteerde in een toename van de doorvoer met 56% (tokens/sec).
  • DDP (GPT-2-1B): Gefen maakte training mogelijk met een microbatch-grootte van 2, terwijl AdamW zelfs een microbatch-grootte van 1 niet kon huisvesten. Dit resulteerde in een doorvoerverbetering van 21% ten opzichte van Adam-mini.
• Stabiliteit: Ablatie-studies bevestigden dat het geleerde codebook stabiel blijft tijdens de training, wat de beslissing valideert om het alleen bij initialisatie te leren.

Betekenis en Claims

Het artikel positioneert Gefen als een praktische, "drop-in" vervanging voor AdamW die het geheugenprobleem bij grootschalige training aanpakt zonder prestatieverlies of de noodzaak voor handmatige afstemming.

• Praktische Impact: Door de optimizer-toestand te verminderen, stelt Gefen gebruikers in staat om grotere modellen te trainen of grotere batchgroottes te gebruiken op vaste hardware, wat de trainingsdoorvoer en de efficiëntie van monsterverwerking direct verbetert.
• Generaliteit: In tegen tegenstelling tot eerdere methoden die vertrouwen op handmatige heuristieken of specifieke tensornamen, is Gefen architectuur-agnostisch en past het zich automatisch aan de data aan.
• Complementariteit: De auteurs merken op dat de optimizer-toestandcompressie van Gefen orthogonaal is aan gradiënt-communicatiecompressie (bijv. 1-bit Adam) en gedistribueerde trainingsstrategieën (FSDP/DDP), wat betekent dat het gecombineerd kan worden met deze technieken voor verdere winst.

De auteurs bieden een volledige Python-implementatie met gefuseerde CUDA-kernels, waarbij de nadere nadruk wordt gelegd op reproduceerbaarheid en eenvoudige integratie in bestaande trainingspipelines.