🤖 machine learning

DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers

Dit artikel introduceert DASH, een aanzienlijk snellere implementatie van de Distributed Shampoo-optimizer die gebruikmaakt van 3D-tensorstacking voor verbeterde GPU-benutting en nieuwe inverse-wortel-solvers (Newton-DB en Chebyshev-benaderingen) om een versnelling van tot wel 5,6x te bereiken terwijl de convergentiekwaliteit gelijk blijft of wordt verbeterd.

Oorspronkelijke auteurs: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

Gepubliceerd 2026-06-26

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een reusachtige robot probeert te leren een nieuwe taal te spreken. Om dit te doen, heb je een "optimizer" nodig—een slimme coach die de hersenen van de robot (de parameters) aanpast na elke les om fouten te minimaliseren.

Jarenlang was de meest populaire coach Adam, een methode die snel is maar een beetje "lui". Het kijkt alleen naar hoeveel elke individuele neuron moet veranderen, zonder te kijken naar hoe neuronen samenwerken in teams.

Maak kennis met Shampoo, een veel slimmere coach. In plaats van naar neuronen één voor één te kijken, kijkt Shampoo naar hoe ze in groepen met elkaar interageren. Dit leidt tot beter leren en modellen die later gemakkelijker te comprimeren zijn. Echter, er is een addertje onder het gras: Shampoo is ontzettend traag. Het is alsof je een genie als coach hebt die zoveel tijd besteedt aan het berekenen van de perfecte zet, dat de robot nauwelijks oefening krijgt.

Het paper introduceert DASH (Distributed Accelerated SHampoo), een nieuw systeem dat deze genie-coach op de snelheid van een sprinter laat draaien zonder zijn intelligentie te verliezen. Dit is hoe ze het deden, met behulp van eenvoudige analogieën:

1. De "Stapeling"-truc (Batched Block Preconditioning)

Het Probleem:
Stel je voor dat je een enorme bibliotheek aan boeken (de data) hebt die georganiseerd moeten worden. De oude manier (Distributed Shampoo) was om één boek te pakken, het te organiseren, terug te leggen, het volgende boek te pakken, het te organiseren, enzovoort. Zelfs als je 1.000 werkers (GPU's) had, stonden ze meestal te wachten tot de vorige persoon klaar was. Dit is inefficiënt.

De DASH-oplossing:
DASH verandert de workflow. In plaats van boeken één voor één te behandelen, stapelt het ze in nette, uniforme 3D-torens. Nu kunnen de werkers een hele toren grijpen en alle boeken in die toren tegelijkertijd organiseren.

De Analogie: Het is het verschil tussen een kassier die artikelen één voor één scant versus een lopende band die een hele doos met boodschappen direct in een machine scant.
Het Resultaat: Deze "stapeling" zorgt ervoor dat de krachtige grafische chips (GPU's) van de computer op volle capaciteit werken, waardoor de optimizer-stappen tot wel 5,6 keer sneller worden.

2. De "Short-cut" Wiskunde (Efficient Inverse-Root Solvers)

Het Probleem:
Om zijn werk te doen, moet Shampoo bij elke stap een zeer moeilijke wiskundige puzzel oplossen: het vinden van de "inverse square root" van een gigantische matrix. De oude manier om dit op te lossen was als het zoeken naar een speld in een hooiberg door elk stukje hooi één voor één te controleren (een methode genaamd Eigen-Value Decomposition). Het is accuraat, maar pijnlijk traag.

De DASH-oplossing:
De auteurs hebben twee nieuwe "shortcuts" geïntroduceerd om deze puzzel op te lossen:

Newton-DB: Een slimme iteratieve methode die met elke gok dichter bij het antwoord komt, zoals een GPS die je route herberekent terwijl je rijdt.
Chebyshev Polynomials: Een wiskundige benadering die het antwoord zeer snel raadt.
De Analogie: In plaats van elke weg in een doolhof af te lopen om de uitgang te vinden (de oude manier), zijn deze nieuwe methoden als het hebben van een kaart die de algemene richting aangeeft, waardoor je in een sprint naar de uitgang kunt rennen.

3. De "Liniaal"-kalibratie (Matrix Scaling)

Het Probleem:
Wanneer je deze shortcuts gebruikt, kan de wiskunde instabiel worden als de getallen te groot of te klein zijn. De oude methode gebruikte een "liniaal" (Frobenius norm) die te lang was, waardoor de getallen werden uitgerekt en de shortcuts veel meer stappen nodig hadden om te convergeren. Het was alsof je een kleine mier probeerde te meten met een meetlint van 30 meter; de precisie gaat verloren.

De DASH-oplossing:
De auteurs realiseerden zich dat ze een betere liniaal nodig hadden. Ze ontwikkelden een techniek genaamd Multi-Power-Iteration.

De Analogie: In plaats van de lengte van de mier te raden met een gigantisch meetlint, gebruiken ze een gespecialiseerde, hogeprecisie schuifmaat die perfect bij de mier past. Dit zorgt ervoor dat de wiskunde snel convergeert en niet crasht door numerieke fouten.

4. De Resultaten

Het paper testte DASH op een groot taalmodel (Llama met 953 miljoen parameters).

Snelheid: DASH voltooide het "denkgedeelte" van de trainingsstap 5,6 keer sneller dan de vorige beste versie.
Kwaliteit: Ondanks dat het veel sneller was, waren de modellen die met DASH getraind werden net zo goed, of zelfs iets beter, in het leren. Sterker nog, de nieuwe "Newton-DB" shortcut produceerde modellen met de laagste foutmarges (perplexity) van alle geteste methoden.

Samenvatting

Beschouw Shampoo als een Ferrari-motor die vastzat in de file omdat de bestuurder de panoramische route koos. DASH is diezelfde Ferrari-motor, maar nu heeft het:

Een bredere snelweg (Stacking blocks) zodat het sneller kan rijden.
Een GPS-shortcut (Newton-DB) om files te vermijden.
Betere navigatietools (Multi-Power-Iteration) om te zorgen dat het niet verdwaalt.

Het resultaat is dat de "slimme" optimizer (Shampoo) niet langer te traag is om te gebruiken, wat het een praktische keuze maakt voor het trainen van de volgende generatie AI-modellen.

Technische Samenvatting: DASH – Snellere Shampoo via Batched Block Preconditioning en Efficiënte Inverse-Root Solvers

Probleemstelling

Shampoo is een vooraanstaande benaderde tweede-orde optimizer die bekend staat om het produceren van modellen met lagere activatie-outliers en verbeterde compressibiliteit. De brede adoptie ervan wordt echter gehinderd door aanzienlijke computationele overhead, primair door de dure operatie van het berekenen van inverse matrixwortels ( $A^{-1/4}$ ) voor preconditioneringsmatrices. Hoewel Distributed Shampoo probeert dit te verzachten door preconditioners op te splitsen in $B \times B$ blokken om de complexiteit te verminderen van $O(n^3)$ naar $O(Bn^2)$ , lijden bestaande implementaties onder kritieke efficiëntiekloven:

Sequentiële Verwerking: Standaardimplementaties verwerken blokken sequentieel, waardoor ze er niet volledig in slagen de moderne GPU-parallelliteit te benutten.
Inefficiënte Solvers: De standaardafhankelijkheid van Eigenwaarde-Decompositie (EVD) is moeilijk te parallelliseren op GPU's. Alternatieve iteratieve methoden zoals Coupled-Newton (CN) worden vaak standaard uitgeschakeld vanwege zorgen over numerieke stabiliteit en convergentiesnelheid.
Suboptimale Schaling: Huidige methoden schalen matrices vaak met behulp van de Frobenius-norm, wat computationeel goedkoop is maar de eigenwaarden naar nul kan duwen, waardoor er aanzienlijk meer iteraties nodig zijn om een streefprecisie te bereiken.

Methodologie

Het artikel stelt DASH (Distributed Accelerated SHampoo) voor, een hoogwaardige implementatie die is ontworpen om de kloof tussen theoretische efficiëntie en hardwarecapaciteiten te overbruggen. De methodologie rust op twee primaire pijlers:

1. Batched Block Preconditioning (Systeemniveau Optimalisatie)

In plaats van preconditioneringsblokken sequentieel te verwerken, stapelt DASH onafhankelijke blokken op in 3D-tensoren.

3D Tensor Stacking: Gradiëntblokken $G_{ij}$ worden gestapeld in een matrix met de vorm $(N, B, B)$ , waarbij $N$ het totaal aantal blokken is.
Parallelle Executie: Dit maakt het mogelijk om de inverse root-procedures (EVD, CN, of nieuwe iteratieve methoden) toe te passen op de gehele batch van blokken gelijktijdig met behulp van batched matrixvermenigvuldigingsprimitieven (bmm).
GPU Benutting: Deze architecturale wijziging maximaliseert de throughput van GPU tensor cores, vermindert geheugenfragmentatie en maakt efficiënt gebruik van half-precision (FP16) rekenkunde voor bepaalde solvers.

2. Efficiënte Inverse-Root Solvers (Algoritmische Optimalisatie)

Het artikel introduceert en analyseert geavanceerde lineair-algebraïsche benaderingen om EVD en CN te vervangen of aan te vullen:

Newton-Denman-Beavers (NDB) Iteratie: Een nieuwe benadering voor Shampoo die zowel vierkants als inverse wortels berekent. De auteurs optimaliseren de eerste iteratie om redundante matrixvermenigvuldigingen te elimineren, waardoor de overhead wordt teruggebracht tot 3 matmuls per stap (in de praktijk effectief minder).
Chebyshev Polynomial Approximaties: Het gebruik van Clenshaw's algoritme om matrixmachten te benaderen, wat een alternatieve iteratieve solver biedt.
Multi-Power-Iteration voor Schaling: Om convergentieproblemen aan te pakken, vervangt het artikel de standaard Frobenius-norm schaling door een Multi-Power-Iteration methode.
- Mechanisme: Het schat de spectrale radius ( $\lambda_{max}$ ) met behulp van een pool van startvectoren (tot 32) die parallel via Power-Iteration worden verwerkt.
- Voordeel: Dit levert een nauwkeurigere schaleringsfactor op dan de Frobenius-norm, waardoor de eigenwaarden in een regime blijven waar iteratieve solvers snel convergeren (dichter bij 1), wat het aantal vereiste iteraties vermindert en de numerieke stabiliteit verbetert.

Belangrijkste Bijdragen

DASH Implementatie: Een GPU-efficiënt framework dat sequentiële loops vervangt door parallelle 3D-tensoroperaties, waarmee tot 5,6× snellere optimizerstappen worden bereikt vergeleken met goed geoptimaliseerde Distributed Shampoo.
Nieuwe Solvers: De introductie van de Newton-DB (NDB) iteratie en Chebyshev polynomial benaderingen als levensvatbare, snellere alternatieven voor EVD en CN. Het artikel demonstreert dat NDB een lagere validatie-perplexity oplevert dan zowel Coupled-Newton als standaard EVD.
Convergentie Analyse: Een gedetailleerde gedragsanalyse die aantoont dat standaard Frobenius-norm schaling suboptimal is, wat vaak leidt tot tragere convergentie. Het artikel biedt intuïtieve verklaringen voor de verschillende convergentiegedragingen van CN versus NDB.
Multi-Power-Iteration: Een efficiënte, half-precision implementatie van Power-Iteration die spectrale radii robuust schat, wat optimale matrixschaling mogelijk maakt die voldoet aan de convergentiecriteria voor NDB en CN met minder stappen.

Experimentele Resultaten

De auteurs evalueerden DASH op een Llama-model met 953M parameters (embedding grootte 2048, sequentielengte 1024) getraind op de C4-dataset.

Snelheid:
- DASH vermindert de runtime van de optimizer step met 4× tot 5,6× vergeleken met Distributed Shampoo, afhankelijk van de inverse root methode en de gebruikte precisie.
- De meest significante versnelling (5,6×) werd waargenomen bij het vergelijken van Distributed Shampoo's EVD (FP32) tegen DASH's CN (FP16) met GPU-level stacking.
- Zelfs met frequente preconditioner updates ( $f=1$ ), behoudt DASH aanzienlijke snelheidvoordelen.
Nauwkeurigheid (Validatie Perplexity):
- DASH komt in bijna alle stabiele configuraties overeen met Distributed Shampoo.
- NDB levert consistent een lagere validatie-perplexity op dan CN en presteert vaak beter dan EVD.
- Het gebruik van Power-Iteration scaling met NDB resulteerde in de laagste validatie-perplexity (11,68) vergeleken met Frobenius scaling (11,76), wat de theoretische analyse valideert.
Precisie:
- De implementatie ondersteunt FP16 voor CN en EVD zonder de perplexity te verslechteren, wat de runtime verder verkort.
- NDB in FP16/BF16 bleek in hun experimenten numeriek instabiel te zijn, hoewel de auteurs aangeven dat dit een gebied is voor toekomstig onderzoek.

Betekenis en Claims

Het artikel claimt dat DASH de computationele barrière voor het gebruik van tweede-orde optimizers zoals Shampoo aanzienlijk verlaagt, waardoor hoogwaardige preconditioning praktisch bruikbaar wordt voor grootschalige deep learning.

Brug tussen Theorie en Hardware: Door algoritmisch ontwerp (iteratieve solvers, matrixschaling) af te stemmen op hardwarecapaciteiten (tensor cores, gebatchte operaties), stelt DASH onderzoekers in staat om eigenschappen van optimizers te onderzoeken zonder prohibitieve looptijden.
Numerieke Stabiliteit vs. Snelheid: Het werk daagt de aanname uit dat iteratieve methoden inherent minder stabiel of trager zijn dan EVD. Door een betere schaling (Multi-Power-Iteration) en geoptimaliseerde iteraties (NDB) te introduceren, demonstreren de auteurs dat iteratieve methoden een superieure convergentie en lagere perplexity kunnen bereiken terwijl ze aanzienlijk sneller zijn.
Toekomstige Richtingen: De auteurs suggereren dat de superieure prestaties van Newton-DB deuren openen voor dynamische solverselectie op basis van de block condition numbers en het stabiliseren van NDB in lagere-precisie formaten met behulp van stochastische of foutcorrectiemethoden.

Het artikel concludeert dat de primaire bottleneck van Shampoo niet de wiskundige complexiteit was, maar gefragmenteerde executie en numerieke instabiliteit, zaken die DASH effectief aanpakt.