A little outer momentum helps DiLoCo, but only at large H

On a 14k-parameter character model, outer momentum of 0.1 to 0.3 trims 0.01 to 0.02 nats off long local runs and hurts short ones; 0.5 and up still hurts.

Research findings · 1 min read · Source
researchdistributed-trainingdilocoouter-momentum

An earlier step in research/lowcomm-training found that DiLoCo's outer momentum (0.5 and 0.9) made loss worse once the inner learning rate was tuned. That grid skipped everything between 0 and 0.5, so we filled it in: momentum 0.1, 0.2 and 0.3, four workers, 640 steps per worker, five seeds, held-out loss.

Hmomentum 00.10.20.3
162.18312.19302.20362.2264
642.21012.19622.19242.1898
1602.21772.20892.20702.2134

At H=64 every seed improves with momentum 0.1 (paired difference -0.014). At H=160 the gain is smaller and peaks at 0.2. At H=16 momentum hurts, steadily. Synchronous AdamW is 2.1585, so the gap at large H shrinks by a fifth to a third; it does not close.

So the earlier "momentum hurts" holds from 0.5 up, not below. Limits: a 14k parameter model, inner learning rate not re-tuned per momentum, no correction for nine comparisons. The log and raw numbers are in research/lowcomm-training/LOG.md (milestone 2g).

More in Research findings

All writing · All in this topic · All tools and research