💬 NLP

Can Model Merging Improve Aggregation in DiLoCo?

Dit artikel overbrugt de kloof tussen modelmerging en gedistribueerd leren door aan te tonen dat het aanpassen van de Iso-C merging techniek met Nesterov momentum (IsoLoCo) bestaande DiLoCo-methoden aanzienlijk overtreft in low-communication gedistribueerde training, met name naarmate het aantal lokale workers toeneemt.

Oorspronkelijke auteurs: Stefan Horoi, Benjamin Thérien, Guy Wolf, Eugene Belilovsky

Gepubliceerd 2026-07-07

📖 4 min leestijd☕ Koffiepauze-leesvoer

CC BY 4.0

Oorspronkelijke auteurs: Stefan Horoi, Benjamin Thérien, Guy Wolf, Eugene Belilovsky

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 probeert een gigantisch, superintelligent AI-brein te trainen. Hiervoor heb je een enorme hoeveelheid computerkracht nodig. Meestal gebruiken bedrijven een "data-parallel" aanpak: ze zetten honderden computers zij aan zij op die allemaal precies dezelfde taak op exact hetzelfde moment uitvoeren, terwijl ze constant updates naar elkaar roepen. Dit is snel, maar het vereist dure, hogesnelheidskabels die elke computer met elkaar verbinden. Als één computer traag is of de kabel zwak is, moet de hele rij wachten.

Het Probleem: De "Eiland"-aanpak
Om geld te besparen en zwakkere verbindingen te kunnen gebruiken, ontwikkelden onderzoekers een methode genaamd DiLoCo. Stel je voor dat je in plaats van een enkele lijn veel aparte "eilanden" van computers hebt.

Elk eiland traint een tijdje (laten we zeggen 30 stappen) zijn eigen versie van het AI-brein onafhankelijk van elkaar.
Ze praten tijdens die 30 stappen niet met elkaar.
Na 30 stappen sturen ze een samenvatting van wat ze hebben geleerd (een "pseudo-gradiënt") naar een centraal knooppunt.
Het knooppunt middelt deze samenvattingen en werkt het hoofdbrein bij.

De Fout:
De paper vond een gebrek in deze "Eiland"-aanpak. Naarmate je meer eilanden (werkers) toevoegt of ze langer laat trainen voordat ze praten, raakt de AI in de war. Het is alsof een groep mensen samen probeert een tekening te maken zonder met elkaar te praten: als ze te lang op hun eigen eiland werken, beginnen ze verschillende dingen te tekenen, en wanneer ze proberen hun werk te combineren, wordt het resultaat een rommeltje. In technische termen: de "updates" van de verschillende eilanden beginnen met elkaar te botsen, wat de prestaties van de AI verslechtert.

De Inspiratie: Het Samenvoegen van Expertmodellen
Onafhankelijk daarvan was een andere groep onderzoekers bezig met Model Merging. Stel je voor dat je een expertkok hebt die leerde koken in de Italiaanse keuken en een andere die leerde koken in de Japanse keuken. Je wilt één kok die beide kan.

De Oude Manier: Gewoon de recepten van beide experts middelen. Dit mislukt vaak omdat de ingrediënten botsen (bijv. sojasaus mengen met slagroom).
De Nieuwe Manier (Task Arithmetic): In plaats van het hele recept te middelen, kijk je naar het verschil tussen het definitieve recept van de expert en het oorspronkelijke basisrecept. Je combineert dan zorgvuldig alleen die verschillen.

Het grote "Aha!"-moment van de paper was de realisatie dat DiLoCo eigenlijk hetzelfde doet als Model Merging, maar dan in een lus.

Het "Basismodel" is het AI-brein van de vorige keer dat iedereen met elkaar praatte.
De "Eilanden" zijn de experts die op hun eigen eiland hebben getraind.
De "Updates" die ze terugsturen zijn de "verschillen" (of taakvectoren).

De paper betoogt dat de reden waarom DiLoCo rommelig wordt met veel eilanden, dezelfde is als de reden waarom het samenvoegen van recepten rommelig wordt: Interferentie. De updates van verschillende eilanden vechten tegen elkaar.

De Oplossing: IsoLoCo
De auteurs besloten een techniek uit de wereld van Model Merging te lenen om DiLoCo op te lossen. Ze testten verschillende samenvoegingstechnieken en ontdekten dat er één genaamd Iso-C (Isotrope samenvoeging) het beste werkte.

Zie Iso-C als een "harmoniseerder" voor de updates. Wanneer de eilanden hun updates terugsturen:

De harmoniseerder kijkt naar de "vorm" van de updates.
Als één eiland te hard roept in een specifieke richting (een specifieke wiskundige richting), zet de harmoniseerder dit volume zachter om het gemiddelde te evenaren.
Dit voorkomt dat een enkel eiland domineert of botst met de anderen, wat zorgt voor een soepelere, meer gebalanceerde gecombineerde update.

Ze voegden deze "harmoniseerder" toe aan een "momentum"-functie (zoals een hardloper die zijn snelheid behoudt, zelfs als de weg hobbelig wordt) om een nieuwe methode te creëren die ze IsoLoCo noemen.

De Resultaten
De paper testte dit op verschillende groottes van AI-modellen en met verschillende aantallen eilanden (van 1 tot 128).

De Winnaar: IsoLoCo versloeg consequent de oorspronkelijke DiLoCo-methode.
Het Gat: Hoe meer eilanden je toevoegde, hoe groter het voordeel van IsoLoCo was. Bij 128 eilanden werd de oorspronkelijke methode quite rommelig, maar IsoLoCo bleef schoon en efficiënt.
De Snelheid: Ze maakten ook een "snelle modus" voor IsoLoCo met behulp van een wiskundige shortcut (Newton-Schulz iteraties), waardoor het proces veel sneller verliep zonder kwaliteitsverlies.

Samenvattend
De paper laat zien dat door gedistribueerde AI-training te behandelen als een "samenvoegen van experts"-probleem, we geavanceerde wiskundige trucs kunnen gebruiken om te voorkomen dat de computers met elkaar vechten. Dit stelt ons in staat om enorme AI-modellen te trainen over vele zwak verbonden computers zonder aan prestaties in te boeten, wat grootschalige AI-training goedkoper en schaalbaarder maakt.

Technische Samenvatting: Model Merging voor Aggregatie in DiLoCo

Probleemstelling

Distributed Low-Communication (DiLoCo) training is een methode geworden om grote taalmodellen te trainen over slecht verbonden rekenclusters ("eilanden") door globale synchronisatie te minimaliseren. In DiLoCo voeren workers meerdere lokale optimalisatiestappen (inner steps) uit voordat ze hun parameterupdates (pseudo-gradiënten) aggregeren via een globale optimizer. Hoewel dit effectief is bij het verminderen van communicatiekosten, lijdt DiLoCo onder prestatievermindering naarmate het aantal workers ( $R$ ) of het aantal lokale stappen ( $H$ ) toeneemt. Deze degradatie treedt op omdat lokale replica's aanzienlijk divergeren, waardoor de gemiddelde pseudo-gradiënten niet meer uitgelijnd zijn met de updates die worden geproduceerd door nauw gesynchroniseerde data-parallel training.

Gelijktijdig zijn Model Merging-technieken ontwikkeld om onafhankelijk gefinetunede modellen te combineren tot één enkel model. Een centrale uitdaging in model merging is "task interference", waarbij parameterupdates van verschillende taken conflicteren, wat leidt tot prestatiedalingen, vooral naarmate het aantal te mergen modellen toeneemt.

De auteurs identificeren een structurele analogie tussen deze twee velden: de pseudo-gradiënt aggregatiestap in DiLoCo (met een SGD outer optimizer) is wiskundig equivalent aan Task Arithmetic model merging. In deze analogie:

De gedeelde globale model aan het begin van een outer step komt overeen met de pre-trained initialisatie.
De lokale worker replica's komen overeen met de gefinetunede "expert" modellen.
De lokale updates (pseudo-gradiënten) komen overeen met "task vectors".
Het aantal workers komt overeen met het aantal taken.

De kern van het probleem dat wordt aangepakt, is of geavanceerde model merging technieken, specifiek ontworpen om task interference te verminderen, kunnen worden ingezet om de aggregatiestap in DiLoCo te verbeteren, waardoor de prestaties worden hersteld naarmate het aantal workers of lokale stappen toeneemt.

Methodologie

Theoretische Brug

Het artikel stelt vast dat de outer update van DiLoCo met SGD een directe toepassing is van task arithmetic. Door de aggregatie van pseudo-gradiënten te beschouwen als een model merging probleem, stellen de auteurs voor om de standaard gemiddelde operatie te vervangen door geavanceerde merging algoritmen die ontworpen zijn om negatieve interferentie te verminderen.

Experimentele Evaluatie van Merging Methoden

De auteurs hebben systematisch verschillende state-of-the-art merging methoden getest als directe vervanging voor de aggregatiestap van DiLoCo. De geëvalueerde methoden waren:

Sparsity-based: TIES (TrIm, Elect Sign) en DARE (Drop And REscale).
Orthogonalization-based: TSV (Task Singular Vector), Iso-C (Isotropic merging in Common subspace), en Iso-CTS.

Bevindingen over Merging Methoden:

Sparsity methoden (TIES, DARE): Presteerden vergelijkbaar met standaard DiLoCo SGD bij zeer lage sparsity levels. De prestaties verslechterden naarmate de sparsity toenam, wat suggereert dat ze geen duidelijk voordeel boden ten opzichte van eenvoudige middeling in deze context.
TSV: Presteerde slecht (hogere validatie loss) en was computationeel te duur, omdat het meerdere Singular Value Decompositions (SVDs) per stap vereiste.
Iso-CTS: Presteerde beter dan de baseline maar was computationeel duur en bood geen significant voordeel ten opzichte van Iso-C in de DiLoCo setting (waar lokale taken kort en vergelijkbaar zijn).
Iso-C: Bleek de meest veelbelovende kandidaat. Het verbetert de prestaties door het singular value spectrum van de geaggregeerde pseudo-gradiënt te vervangen door het gemiddelde, waardoor hoogenergetische richtingen worden afgevlakt om een gebalanceerde common subspace te verkrijgen. Het vereist slechts één SVD per parametermatrix, wat een gunstige kost-prestatie verhouding biedt.

Voorgesteld Algoritme: IsoLoCo

Gebaseerd op het succes van Iso-C, stellen de auteurs IsoLoCo (Isotropic Distributed Low-Communication) voor. IsoLoCo past Iso-C aan voor gedistribueerde training door:

Isotropic Correction: Voor 2D parametermatrices (gewichtslagen) berekent het de SVD van de gemiddelde pseudo-gradiënt, vervangt de singular values door hun gemiddelde, en reconstrueert de update. Dit reguleert het spectrum terwijl de subspace structuur behouden blijft.
Nesterov Momentum: In tegen tegenstelling tot standaard Iso-C, incorporeert IsoLoCo Nesterov momentum in de outer optimizer. Dit is gemotiveerd door eerder werk dat het voordeel van momentum in FedOpt/DiLoCo settings aantoont. De gecorrigeerde pseudo-gradiënt wordt in een momentum buffer gevoed om de uiteindelijke update te genereren.

Implementatie Optimalisatie

Om de computationele kosten van SVD aan te pakken, hebben de auteurs een snelle implementatie van IsoLoCo ontwikkeld met behulp van Newton–Schulz iteraties. Dit vervangt de standaard SVD door iteratieve orthogonalisatie en gebruikt een RMS proxy voor het singular-value gemiddelde. Deze optimalisatie levert tot twee ordes van grootte snellere orthogonalisatie op met een verwaarloosbare degradatie in prestaties.

Belangrijkste Resultaten

Prestaties over Worker Counts

Experimenten werden uitgevoerd op Llama-stijl modellen (178M, 512M en 1B parameters) met de DCLM dataset.

Opschalen met Workers: IsoLoCo presteert aanzienlijk beter dan standaard DiLoCo naarmate het aantal workers toeneemt.
- Bij de 178M schaal met $R=128$ , verminderde IsoLoCo de validatie loss kloof ten opzichte van de data-parallel baseline met 5,8% vergeleken met DiLoCo.
- Bij $R=128$ vertoonde IsoLoCo een degradatie van 9,7% ten opzichte van data-parallel training, terwijl DiLoCo een degradatie van 16,5% liet zien.
- IsoLoCo presteerde zelfs beter dan de data-parallel baseline bij lage worker counts ( $R \le 4$ ).
Model Grootte: Het voordeel van IsoLoCo ten opzichte van DiLoCo blijft bestaan over verschillende model groottes (178M tot 1B), hoewel de relatieve verbetering iets kleiner is voor grotere modellen (1,2% bij 1B vs. ~2,4% bij 178M), waarschijnlijk omdat grotere modellen van nature beter schalen met het aantal workers.

Impact van Inner Steps

IsoLoCo werd geëvalueerd met toenemende synchronisatie-intervallen ( $H = 30$ tot $240$). Het presteerde consistent beter dan DiLoCo over alle intervallen, wat duidt op robuustheid tegen langere periodes van lokale divergentie.

Vergelijking met Momentum-gebaseerde DiLoCo

Opvallend genoeg zorgde de combinatie van DiLoCo SGD met Iso-C aggregatie (zonder momentum) al voor betere prestaties dan de standaard DiLoCo (die Nesterov momentum gebruikt). Het toevoegen van Nesterov momentum aan Iso-C (wat IsoLoCo creëert) vergrootte het prestatieverschil verder, wat aangeeft dat de merging-gebaseerde aggregatie en momentum mechanismen complementair zijn.

Compatibiliteit met Muon

Het artikel merkt op dat IsoLoCo kan worden samengesteld met de Muon inner optimizer (MuLoCo). In deze configuratie behaalde IsoLoCo de laagste validatie loss onder alle methoden voor $R \ge 32$ , terwijl IsoLoCo met een AdamW inner optimizer MuLoCo evenaarde of overtrof bij hoge worker counts met aanzienlijk lagere orthogonalisatiekosten.

Betekenis en Claims

Het artikel claimt dat het bekijken van low-communication gedistribueerde training door de lens van model merging een krachtig kader biedt voor het verbeteren van de schaalbaarheid. De primaire bijdragen zijn:

Theoretische Unificatie: Het vaststellen van een directe analogie tussen de pseudo-gradiënt aggregatie van DiLoCo en task-arithmetic model merging.
Empirische Validatie: Het aantonen dat model merging technieken, specifiek Iso-C, als superieure aggregatieregels voor gedistribueerde training kunnen dienen, waarbij ze eenvoudige middeling en zelfs momentum-gebaseerde DiLoCo overtreffen.
IsoLoCo: Het voorstellen van een praktisch, momentum-enhanced algoritme dat de prestatiekloof tussen low-communication training en data-parallel gouden standaarden aanzienlijk verkleint, met name in regimes met veel workers of lange synchronisatie-intervallen.
Efficiëntie: Het leveren van een Newton-Schulz-gebaseerde implementatie die het SVD-zware merging proces computationeel haalbaar maakt voor iteratieve gedistribueerde training.

De auteurs concluderen dat merging-geïnspireerde aggregatie een effectieve strategie is voor low-communication gedistribueerde training, wat een pad biedt om training te schalen over geografisch verspreide of slecht verbonden clusters zonder kwaliteitsverlies van het model.