Technische Samenvatting: Cosine-Gated Adam-Decay (CGAD) voor Gekoppeld DiLoCo

1. Probleemstelling

In asynchrone Gekoppelde DiLoCo-systemen (DDiLoCo) voeren workers meerdere lokale optimalisatiestappen uit voordat ze pseudo-gradiënten uitwisselen met een globale parameter-server. Een kritieke uitdaging in deze architectuur is veroudering (staleness): de pseudo-gradiënten die in de buitenste ronde $t$ worden toegepast, kunnen $\tau$ ronden eerder zijn berekend als gevolg van netwerklatentie, trage workers (stragglers) of asynchrone budgetten.

De standaard buitenste optimizer voor DiLoCo is SGD met Nesterov-momentum. Deze methode past een uniforme update-stap toe op elke inkomende pseudo-gradiënt zonder te conditioneren op de ouderdom ($\tau$). Naarmate $\tau$ toeneemt, accumuleert de momentum-buffer een verouderd beeld van de verlieslandschap. Bijgevolg onderneemt de optimizer stappen die richtingen overschrijden die niet langer het huidige verlies doen dalen, wat leidt tot instabiliteit of divergentie, vooral naarmate de modelgrootte toeneemt. Bestaande theoretische analyses van asynchrone momentum-SGD suggereren dat convergentiesnelheden verslechteren met een factor evenredig aan $\tau_{max}^2$ (het kwadraat van de maximale vertraging).

2. Methodologie: Cosine-Gated Adam-Decay (CGAD)

De auteurs stellen Cosine-Gated Adam-Decay (CGAD) voor, een drop-in vervanging voor de Nesterov-buitenste optimizer die updates expliciet conditioneert op de ouderdom van de pseudo-gradiënt.

Kernmechanisme

CGAD wijzigt de standaard Adam-update door elke inkomende pseudo-gradiënt $g_f$ voor fragment $f$ te schalen met een factor $\sigma(\tau)$ voordat deze de eerste- en tweede-momentbuffers van Adam binnenkomt:
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

Waarbij:

• Exponentiële Decay ($e^{-\alpha \tau}$): De informatieveroudering van verouderde gradiënten modelleert. Naarmate $\tau$ groeit, wordt de bijdrage van de gradiënt aan de momentumbuffers exponentieel verminderd.
• Cosine Gate ($\gamma(\tau)$): Een gladde afkapfunctie gedefinieerd als:
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  Deze gate is gelijk aan 1 voor frisse gradiënten, daalt glad af naar 0 bij een gekozen afkapwaarde $\tau_{cut}$, en zet bijdragen strikt op nul voor $\tau > \tau_{cut}$.

Belangrijkste Ontwerpkies

1. Pre-Buffer Gating: De schaling $\sigma(\tau)$ wordt toegepast op de gradiënt voordat deze de eerste moment ($m$) en tweede moment ($v$) buffers updatet. Dit voorkomt dat verouderde gradiënten de tweede moment ($v$) opblazen, wat anders latere frisse updates zou verzwakken.
2. Per-Fragment Variant (PA-CGAD): Voor partial-sync planners (bijv. Streaming DiLoCo) houdt de methode de ouderdom van individuele fragmenten ($a_f$) bij en gate-updates op basis van $\max(\tau, a_f)$, waardoor een verenigd verouderingsmodel mogelijk is voor zowel netwerkvertraging als planningsvertragingen.
3. Hyperparameters: De methode introduceert twee hyperparameters, $\alpha$ (decay rate) en $\tau_{cut}$ (cutoff), met standaardwaarden ($\alpha=0.2, \tau_{cut}=32$) die over modelgroottes heen overdraagbaar zijn.

3. Theoretische Bijdragen

Het paper biedt een niet-asymptotische convergentiebound voor een geïdealiseerde gated-adaptieve update op gladde niet-convexe doelfuncties.

• Onafhankelijkheid van Verouderingsbias: De term voor verouderingsbias in de convergentiebound hangt af van de gating-hyperparameter $\alpha$ (specifiek $O(1/e^\alpha)$) in plaats van de gerealiseerde maximale vertraging $\tau_{max}$.
• Contrast met Standaard Analyse: Standaard analyses van asynchrone momentum-SGD leveren een complexiteitsterm op die schaalt met $\tau_{max}^2$. Door deze te vervangen door een term die afhankelijk is van $\alpha$, koppelt het gating-mechanisme theoretisch de convergentiestraf los van de worst-case vertragingverdeling, mits $\tau_{cut} \ge 1/\alpha$.
• Reikwijdte: De auteurs merken op dat dit een motivatie is voor het gating-mechanisme in plaats van een volledige convergentiestelling voor empirische Adam, omdat het uitgaat van gebonden genormaliseerde stappen in plaats van deze af te leiden uit empirische statistieken.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op Llama-stijl taalmodel-pretraining op drie schalen: 25 M, 1 B en 7 B parameters, met gecontroleerde vertragingen tot $\tau=16$.

Stabiliteit en Convergentie

• 1 B Schaal: Bij $\tau=16$ divergeerde het standaard Nesterov-recept (verlies $\approx 319$), terwijl CGAD een verlies van 7,44 bereikte, vergelijkbaar met de prestaties bij $\tau=0$ (7,48). De Nesterov-methode verloor bij $\tau \ge 8$ de convergentie volledig.
• 25 M Schaal: CGAD en andere verouderingsbewuste baselines (Adam-Decay, SDM, Delayed Nesterov) bleven allemaal stabiel, terwijl Nesterov divergeerde bij $\tau \ge 4$. Op deze schaal waren de verschillen tussen verouderingsbewuste methoden marginaal.

Schaalverzekering en Robuustheid

De meest significante bevinding doet zich voor bij de 7 B schaal onder geheugengeconstrueerde omstandigheden (bf16-weights, 8-bit AdamW):

• Seed-tot-Seed Variantie: De dichtstbijzijnde baseline, Adam-Decay (CGAD zonder de cosine-afkap), vertoonde een 27-voudige toename in de seed-tot-standaardafwijking ($\sigma$) naarmate het model groeide van 25 M naar 7 B (van 0,07 naar 1,89 bij $\tau=8$). Daarentegen bleef de $\sigma$ van CGAD vlak (ongeveer 0,38).
• Single-Shot Risico: Het gebruik van "mean + $\sigma$" als proxy voor het risico van single-shot implementatie, betekende dat het risico van Adam-Decay bij 7 B / $\tau=8$ (13,19) de kans-niveau verlies (10,83) met 2,4 nats overschreed, wat impliceert dat er een hoge waarschijnlijkheid is van falen in één productie-run. Het risico van CGAD (11,28) bleef ver onder deze drempel.
• Mechanisme van Falen: De auteurs schrijven het falen van Adam-Decay bij 7 B toe aan de accumulatie van afgeronde kwadratische ruis in de tweede-momentbuffer ($v$) van verouderde gradiënten die nooit volledig op nul worden gezet. De cosine-afkap in CGAD voorkomt dit door bijdragen boven $\tau_{cut}$ strikt op nul te zetten.

Afstemmingsgevoeligheid

• Nesterov vs. CGAD: Zelfs toen Nesterov opnieuw werd afgestemd (momentum volledig verwijderd), kon het de prestaties van een afgestemde CGAD-configuratie bij hoge vertragingen niet evenaren. Afgestemd Nesterov vereiste verschillende hyperparameters voor verschillende vertragingregimes, terwijl CGAD robuustheid behield over het volledige bereik.
• Hyperparameter Gevoeligheid: De optimale $\alpha$ verschuift licht met de schaal (0,2 voor 25 M, 0,4 voor 1 B bij $\tau=8$), maar de standaardwaarde $\alpha=0.2$ blijft robuust over het volledige vertragingsspectrum.

5. Betekenis en Claims

Het paper beweert dat CGAD de meest robuuste schaalkeuze is voor asynchrone DiLoCo-training, en specifiek de broosheid van Nesterov-momentum onder grote gecontroleerde vertragingen aanpakt.

• Schaalverzekering: De primaire bijdrage is de cosine-afkap, die fungeert als "schaalverzekering". Waar alleen exponentiële decay (Adam-Decay) volstaat voor kleinere modellen of settings met volledige precisie, is de harde afkap kritiek voor frontier-scale modellen (7 B+) die werken met gecomprimeerde optimizer-staten (8-bit/bf16). Het voorkomt de accumulatie van verouderde variantie die leidt tot catastrofale seed-tot-seed variantie en single-shot falen.
• Drop-in Compatibiliteit: CGAD vereist slechts een één-regelige wijziging in bestaande Adam-implementaties en voegt twee hyperparameters toe met standaardwaarden die over schalen heen overdraagbaar zijn.
• Praktische Impact: Voor asynchrone training op frontier-schaal biedt de methode het verschil tussen een enkele trainingsrun die succesvol wordt uitgebracht en een die faalt (boven het kans-niveau verlies uitkomt).

De auteurs blijven bescheiden wat betreft de theoretische reikwijdte, erkend dat de convergentiebound van toepassing is op een geïdealiseerde update en dat de empirische superioriteit specifiek wordt aangetoond binnen de DiLoCo-familie en Llama-stijl pretraining contexten. Zij merken op dat het aantal inner-steps ($H=8$) gebruikt in experimenten lager is dan productie DiLoCo-instellingen ($H=100-500$), hoewel gevoeligheidsstudies suggereren dat de rangschikking behouden blijft bij hogere $H$.