🤖 machine learning

MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

Dit artikel introduceert MKOR, een momentum-geactiveerde tweede-orde optimizer die rank-1 updates en Sherman-Morrison-gebaseerde inversie gebruikt om de computationele en communicatiecomplexiteit respectievelijk van cubisch naar kwadratisch en lineair te verminderen, waardoor de training en convergentie van grote taalmodellen aanzienlijk wordt versneld in vergelijking met de huidige state-of-the-art eerste- en tweede-orde methoden.

Oorspronkelijke auteurs: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

Gepubliceerd 2024-01-30✓ Author reviewed ⓘ

📖 7 min leestijd🧠 Diepgaand

Oorspronkelijke auteurs: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

In het uitgestrekte landschap van moderne kunstmatige intelligentie zijn de krachtigste systemen gebouwd op diepe neurale netwerken, complexe wiskundige structuren die leren patronen te herkennen, talen te vertalen en tekst te genereren. Om deze systemen te onderwijzen, moeten onderzoekers een massaal optimalisatieprobleem oplossen: het vinden van de perfecte set interne getallen, of gewichten, die het netwerk in staat stellen nauwkeurige voorspellingen te doen. Decennialang was de standaardaanpak om deze gewichten een duwtje te geven in de richting die de fout het snelst vermindert, een methode die vergelijkbaar is met een wandelaar die zich een weg naar beneden een mistige berg banen door de helling onder zijn voeten te voelen. Deze eerste-orde methode is betrouwbaar en breed gebruikt, maar kan traag zijn en vereist vaak duizenden stappen om de bodem van de vallei te bereiken.

Een ambitieuzere strategie, bekend als tweede-orde optimalisatie, probeert de wandelaar een kaart van het gehele terrein te geven, inclusief de kromming van de grond, wat veel snellere afdalingen mogelijk maakt. Door te begrijpen hoe de helling verandert, kunnen deze methoden theoretisch de optimale oplossing bereiken in veel minder stappen. Deze aanpak is echter historisch gezien te rekenintensief geweest om praktisch bruikbaar te zijn voor de enorme modellen die vandaag de dag worden gebruikt. Het berekenen van de benodigde kromminginformatie voor een groot netwerk vereist immense rekenkracht en geheugen, waardoor de methode vaak trager is dan de eenvoudigere aanpak die zij beoogt te vervangen. De uitdaging was om de snelheid van de kaartlezende wandelaar te vangen zonder de zware last van het dragen van de kaart.

Een team van onderzoekers van de University of Toronto, Rutgers University en het Texas Advanced Computing Center heeft een nieuwe tool ontwikkeld genaamd MKOR die dit dilemma oplost. Ze hebben een momentum-geactiveerde optimizer gecreëerd die een specifieke wiskundige afkorting gebruikt om de vorm van het terrein te benaderen zonder elk detail te berekenen. In plaats van te proberen de volledige, complexe kromming van het verlieslandschap te berekenen, gebruikt MKOR een techniek genaamd rank-1 updates. Deze aanpak vereenvoudigt de berekening door zich te concentreren op de meest significante richting van verandering, waardoor de computationele kosten effectief worden teruggebracht van een cubische relatie naar een kwadratische relatie. In praktische termen betekent dit dat de methode veel beter schaalt naarmate modellen groter worden, waardoor de knelpunten worden vermeden die voorheen voorkwamen dat tweede-orde methoden werden gebruikt in grote taalmodellen.

De onderzoekers testten deze nieuwe optimizer op een groot taalmodel bekend als BERT-Large-Uncased, waarbij ze experimenten uitvoerden op 64 high-performance grafische processoren. Ze vergeleken MKOR met de beste bestaande eerste-orde methoden en de meest geavanceerde beschikbare tweede-orde implementaties. De resultaten waren opmerkelijk. In een taak die verband houdt met vraagbeantwoording, had de standaard eerste-orde methode 1.536 stappen nodig om een specifiek nauwkeurigheidsniveau te bereiken. De nieuwe MKOR-methode bereikte dezelfde of een betere nauwkeurigheid in slechts 600 stappen. Bij het meten van de totale tijd die nodig is om de training te voltooien, voltooide MKOR de taak in iets meer dan drie uur, terwijl de standaardmethode bijna acht uur in beslag nam. Dit vertegenwoordigt een versnelling van 2,57 keer, wat de trainingstijd effectief met meer dan de helft verkort.

Naast pure snelheid adresseerde de studie een kritiek stabiliteitsprobleem dat tweede-orde methoden vaak teistert. Omdat deze methoden het leerproces aanpassen op basis van complexe wiskundige factoren, zijn ze gevoelig voor "exploderende gradiënten", een fenomeen waarbij de aanpassingen zo groot worden dat het trainingsproces crasht of divergeert. Eerdere oplossingen vereisten vaak het aanzienlijk vertragen van de leersnelheid om de veiligheid te waarborgen, wat de snelheidvoordelen tenietdeed. MKOR bevat een ingebouwde stabilisator die de omvang van deze aanpassingen monitort. Als de waarden te groot worden, schaalt het systeem ze automatisch opnieuw, waardoor wordt voorkomen dat de training ontploft terwijl het agressieve tempo van de tweede-orde aanpak behouden blijft. Dit stelt de optimizer in staat om snel te bewegen zonder het risico op instabiliteit.

De onderzoekers merkten ook op dat het voordeel van tweede-orde methoden het meest uitgesproken is in de vroege stadia van de training, wanneer het model de brede contouren van de data leert. Naarmate het model dichter bij de optimale oplossing komt, nemen de voordelen van de complexe krommingberekeningen af. Om dit aan te pakken, introduceerden zij een hybride versie van de optimizer, genaamd MKOR-H. Deze versie begint met de krachtige tweede-orde methode om de initiële leerfase te versnellen, maar schakelt automatisch over naar een eenvoudigere, eerste-orde methode zodra de trainingsverlies niet langer snel afneemt. Deze hybride aanpak stelde het systeem in staat om in nog minder stappen te convergeren, waarbij dezelfde uiteindelijke nauwkeurigheid als de standaardmethode werd bereikt in minder dan de helft van de tijd.

De studie onderzocht ook waarom deze specifieke wiskundige afkorting zo goed werkt. De onderzoekers analyseerden de datastructuren binnen het neurale netwerk en ontdekten dat de covariantie-matrices, die beschrijven hoe verschillende delen van het netwerk met elkaar interageren, vaak een lage rang (low-rank) hebben. Dit betekent dat, ondanks de hoge dimensionaliteit van de data, de informatie geconcentreerd is in een paar belangrijke richtingen. Door deze matrices te benaderen met behulp van rank-1 updates, legt MKOR de essentiële informatie vast die nodig is voor optimalisatie zonder de overhead van volledige matrixberekeningen. Het team verifieerde door uitgebreide tests dat deze benadering zeer weinig fout introduceert, waarmee werd bevestigd dat de vereenvoudiging niet ten koste gaat van de nauwkeurigheid.

Naast snelheid en stabiliteit vermindert de nieuwe methode ook aanzienlijk de hoeveelheid data die tussen verschillende processoren moet worden getransfereerd tijdens de training. In gedistribueerde computing, waar duizenden processoren samenwerken, kan communicatie een grote bottleneck worden. MKOR vermindert deze communicatiecomplexiteit van een kwadratische relatie naar een lineaire relatie, wat betekent dat wanneer er meer processoren worden toegevoegd, het systeem efficiënt schaalt zonder te worden vertraagd door vertragingen in de dataoverdracht. Deze efficiëntie stelt de methode in staat om de enorme batchgroottes aan te kunnen die vereist zijn door moderne grote taalmodellen, wat een beperking was voor eerdere tweede-orde benaderingen.

De implicaties van dit werk reiken verder dan alleen het trainen van één specif kind specifiek model. Door tweede-orde optimalisatie praktisch te maken voor grootschalige systemen, biedt MKOR een pad naar een efficiëntere training van kunstmatige intelligentie. De onderzoekers merkten op dat hun methode kan helpen om aanzienlijke hoeveelheden energie en tijd te besparen voor machine learning-praktici. Hoewel de technologie zelf een neutraal instrument is, zou de toepassing ervan in velden variërend van medische diagnostiek tot klimaatmodellering wetenschappelijke ontdekkingen kunnen versnellen. De auteurs erkenden echter ook dat efficiënte optimalisatiealgoritmen kunnen worden toegepast op elk model, inclusief die met negatieve maatschappelijke gevolgen, en benadrukten dat het verantwoorde gebruik van dergelijke technologie afhangt van de keuzes van de mensen die het hanteren.

Uiteindelijk presenteert het artikel een concrete oplossing voor een langdurig probleem in deep learning. Door momentum, rank-1 benaderingen en een hybride schakelmechanisme te combineren, hebben de onderzoekers een optimizer gecreëerd die sneller, stabieler en schaalbaarder is dan zijn voorgangers. De experimenten tonen aan dat het mogelijk is om de kracht van tweede-orde methoden te benutten zonder de verbijsterende computationele kosten die hen tot nu toe aan de zijlijn hielden. Terwijl het veld van kunstmatige intelligentie doorgaat met het pushen naar grotere en complexere modellen, bieden tools zoals MKOR een manier om deze systemen efficiënter te trainen, waardoor wat ooit een theoretisch voordeel was, een praktische realiteit wordt.

Technische Samenvatting: MKOR – Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

1. Probleemstelling

Second-order optimalisatiemethoden, zoals Natural Gradient Descent (NGD) en Kronecker-Factored Approximate Curvature (KFAC), bieden superieure convergentiesnelheden vergeleken met first-order methoden (bijv. SGD, Adam, LAMB) door gebruik te maken van kromminginformatie (Hessian of Fisher Information Matrix-benaderingen). De adoptie ervan bij het trainen van grootschalige Deep Neural Networks (DNN's), met name Transformers en Large Language Models (LLM's), wordt echter gehinderd door drie primaire knelpunten:

Computationele Complexiteit: Traditionele KFAC-methoden vereisen het inverteren van covariantie matrices van omvang $d \times d$ (waarbij $d$ de laagdimensie is), wat resulteert in een $O(d^3)$ complexiteit. Vergelijkbaar daarmee schalen Sherman-Morrison-Woodbury (SMW) gebaseerde methoden (SNGD) vaak met de batchgrootte $b$ ( $O(b^3)$ ). In Transformers, waar batchgroottes schalen met de sequentielengte, worden deze complexiteiten prohibitief.
Communicatie-overhead: Gedistribueerde training vereist het synchroniseren van second-order informatie (covariantie matrices en hun inversen) over workers. KFAC vereist het synchroniseren van $O(d^2)$ data, terwijl SNGD $O(bd + b^2)$ vereist, wat significante communicatieknelpunten creëert die schaalbaarheid beperken.
Update Frequentie en Stabiliteit: Vanwege hoge inversiekosten updaten state-of-the-art implementaties (zoals KAISA) second-order informatie onregelmatig (bijv. elke 100–1000 iteraties) met behulp van "verouderde" (stale) factoren. Dit vermindert de effectiviteit van de kromminginformatie. Bovendien zijn second-order methoden gevoelig voor numerieke instabiliteit (exploderende gradiënten, singuliere matrices) en vereisen ze vaak dempingsfactoren of precisie-aritmetica met hoge precisie, wat de computationele kosten verder verhoogt.

2. Methodologie: Het MKOR-algoritme

De auteurs stellen MKOR voor (Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates), die de bovenstaande problemen aanpakt via vier kernmechanismen:

Rank-1 Benadering & SM-gebaseerde Inversie:
In plaats van volledige covariantie matrices te berekenen en te inverteren, benadert MKOR de covariantie van activaties en gradiënten met behulp van rank-1 matrices afgeleid van het gemiddelde van samples ( $E[aa^T] \approx \bar{a}\bar{a}^T$ ). Door de Sherman-Morrison identiteit toe te passen op deze rank-1 updates, werkt MKOR de inverse van de Kronecker-factoren direct bij. Dit vermindert de computationele complexiteit van factorinversie van $O(d^3)$ naar $O(d^2)$ .
- Update Regel: De inverse factoren worden bijgewerkt met behulp van momentum ( $\gamma$ ) en de Sherman-Morrison formule, wat frequente updates mogelijk maakt zonder de kosten van volledige matrixinversie.
Lineaire Communicatiecomplexiteit:
In gedistribueerde omgevingen synchroniseert MKOR alleen de twee rank-1 benaderingsvectoren (omvang $2d$ ) tussen workers, in plaats van volledige covariantie matrices ( $4d^2$ ). Dit vermindert de communicatiecomplexiteit van kwadratisch $O(d^2)$ naar lineair $O(d)$ . Daarnaast gebruikt MKOR half-precisie (FP16) voor deze vectoren, wat de communicatiekosten verder halveert, een prestatie die moeilijk is voor andere second-order methoden vanwege numerieke gevoeligheid.
Stabilisatiemechanismen:
Om numerieke instabiliteit en exploderende gradiënten, die gebruikelijk zijn bij second-order methoden, te voorkomen:
- Norm-gebaseerde Stabilisator: Houdt de oneindige norm (infinity norm) van de inverse factoren in de gaten. Als deze een drempel overschrijdt, worden de factoren gemengd met de identiteitsmatrix ( $\hat{L} = \zeta L + (1-\zeta)I$ ) om divergentie te voorkomen.
- Gradiënt Rescaling: Gepreconditioneerde gradiënten worden geschaald om overeen te komen met de norm van de originele gradiënten, zodat learning rate schedulers effectief blijven.
- Theoretische Stabiliteit: De auteurs bewijzen (Lemma 3.1) dat de factoren positief-definiet blijven zonder de noodzaak van kunstmatige dempingsfactoren ( $\mu I$ ), en dat de enkele scalaire deling in de update regel gegarandeerd niet nul is.
Hybride MKOR (MKOR-H):
Omdat de voordelen van second-order methoden vaak afnemen naarmate de training convergeert (waarbij het loss landscape de identiteit nadert), introduceert MKOR-H een schakelmechanisme. Het monitort de snelheid van de reductie in de loss en schakelt automatisch terug naar een first-order optimizer als second-order updates geen versnelling meer bieden, waardoor de afweging tussen kosten en convergentiesnelheid wordt geoptimaliseerd.

3. Belangrijkste Bijdragen

Complexiteitsreductie: MKOR bereikt $O(d^2)$ computationele complexiteit en $O(d)$ communicatiecomplexiteit voor second-order updates, waardoor de knelpunten van $O(d^3)$ en $O(d^2)$ in KFAC en $O(b^3)$ in SNGD worden verlicht.
High-Frequency Updates: Door de inversiekosten te verlagen, maakt MKOR second-order updates bij elke iteratie (of zeer frequent) mogelijk, waardoor de noodzaak voor "verouderde" factoren vervalt en de kromminginformatie volledig wordt benut.
Robuustheid: De introductie van norm-gebaseerde stabilisatie en gradiënt rescaling stelt MKOR in staat om stabiel te opereren zonder dempingsfactoren of double-precision aritmetica, wat efficiënt gebruik van half-precision hardware mogelijk maakt.
Hybride Strategie: De MKOR-H variant schakelt dynamisch tussen second-order en first-order methoden om de efficiëntie gedurende de gehele trainingscyclus te behouden.

4. Experimentele Resultaten

De auteurs evalueerden MKOR op BERT-Large-Uncased met behulp van 64 A100 GPU's, waarbij het werd vergeleken met de state-of-the-art first-order optimizer LAMB en second-order implementaties KAISA (KFAC) en HyLo (SNGD), evenals Eva.

Versnelling: MKOR behaalde een 2.57× versnelling ten opzichte van LAMB en een 1.85× versnelling ten opzichte van KAISA op de pre-training taak. Specifiek verminderde het de trainingstijd van BERT-Large-Uncased van 8 uur naar 3 uur.
Convergentie: Op de SQuAD v1.1 question-answering taak convergeerde MKOR-H in 600 stappen (tegenover 1.536 voor LAMB) terwijl het vergelijkbare of betere F1-scores behaalde (90.64 vs 90.44).
GLUE Benchmark: MKOR behaalde nieuwe state-of-the-art metrieken op de GLUE dataset met 1.500 stappen. Opvallend genoeg, terwijl KAISA er niet in slaagde om zelfs na 1.563 stappen te convergeren naar de gemiddelde baseline nauwkeurigheid, convergeerde MKOR succesvol.
Schaalbaarheid: In tegen tegenstelling tot HyLo, dat divergeerde of niet in het GPU-geheugen paste door de hoge geheugenoverhead ( $O(bd + b^2)$ ), trainde MKOR succesvol met grote batchgroottes (tot 64k effectieve batchgrootte via accumulatie) zonder geheugenproblemen.
Geldigheid van Benadering: Experimenten op ResNet-50 en BERT toonden aan dat rank-1 benaderingen van covariantie matrices een verwaarloosbare fout veroorzaken, wat de kern van de benaderingsveronderstelling valideert.

5. Betekenis en Claims

Het artikel stelt dat MKOR er succesvol in slaagt de kloof te overbruggen tussen de theoretische voordelen van second-order optimalisatie en de praktische beperkingen van het trainen van moderne, grootschalige DNN's. Door de complexiteit van second-order updates te reduceren tot een niveau dat vergelijkbaar is met first-order methoden (wat betreft de schaling met modelgrootte) en tegelijkertijd de convergentievoordelen te behouden, maakt MKOR het volgende mogelijk:

Snellere Training: Significante reductie in de werkelijke trainingstijd (wall-clock time) voor LLM's.
Betere Convergentie: Het vermogen om te convergeren naar hogere nauwkeurigheidsmetrieken waar andere second-order methoden (zoals KAISA) falen of stagneren.
Schaalbaarheid: Het vermogen om te schalen naar grote batchgroottes en modeldimensies zonder de communicatie- en geheugenknelpunten die voorheen second-order methoden beperkten.

De auteurs benadrukken dat hun werk second-order optimalisatie levensvatbaar maakt voor het "large model" tijdperk, waarbij specifiek de schaalbaarheidsvraagstukken worden aangepakt die worden geïntroduceerd door de sequentielengte-afhankelijke batchgroottes in Transformer-architecturen. Zij concluderen dat MKOR een nieuwe state-of-the-art vertegenwoordigt voor efficiënte en effectieve training van diepe neurale netwerken.