技術的サマリー：非同期 DiLoCo 向けのコサインゲート付き Adam 減衰（CGAD）

1. 問題定義

非同期の Decoupled DiLoCo（DDiLoCo）システムにおいて、ワーカーはグローバルパラメータサーバーと疑似勾配を交換する前に、複数の局所最適化ステップを実行する。このアーキテクチャにおける重要な課題は**古さ（staleness）**である。ネットワーク遅延、ストレーガー（遅延）ワーカー、または非同期予算により、外側ラウンド $t$ で適用される疑似勾配は、$\tau$ ラウンド前に計算されたものである可能性がある。

DiLoCo の標準的な外側オプティマイザは、ネステロフ・モメンタム付き SGD である。この手法は、勾配の年齢（$\tau$）に条件付けすることなく、すべての入力疑似勾配に一様な更新ステップを適用する。$\tau$ が増加するにつれて、モメンタムバッファは損失地形の古くなったビューを蓄積する。その結果、オプティマイザは現在の損失を減少させる方向がもはや存在しない方向へ過剰なステップを踏むようになり、特にモデル規模が大きくなるにつれて不安定性や発散を招く。非同期モメンタム-SGD に関する既存の理論的解析では、収束率が最大遅延 $\tau_{max}$ の二乗に比例する因子（$\tau_{max}^2$）によって劣化することが示唆されている。

2. 手法：コサインゲート付き Adam 減衰（CGAD）

著者は、疑似勾配の年齢に明示的に条件付けした更新を行う、ネステロフ外側オプティマイザの代替として機能する**コサインゲート付き Adam 減衰（CGAD）**を提案する。

中核メカニズム

CGAD は、標準的な Adam 更新を修正し、フラグメント $f$ に対する各入力疑似勾配 $g_f$ を、Adam の第一モーメントおよび第二モーメントバッファに入る前に因子 $\sigma(\tau)$ によってスケーリングする：
$$\sigma(\tau) = \gamma(\tau) e^{-\alpha \tau}$$

ここで：

• 指数減衰（$e^{-\alpha \tau}$）： 古くなった勾配の情報減衰をモデル化する。$\tau$ が大きくなるにつれて、勾配のモメンタムバッファへの寄与は指数関数的に減少する。
• コサインゲート（$\gamma(\tau)$）： 以下のように定義される滑らかなカットオフ関数である：
  $$\gamma(\tau) = \frac{1}{2} \left(1 + \cos\left(\frac{\pi \tau}{\tau_{cut}}\right)\right) \cdot \mathbb{I}[\tau \le \tau_{cut}]$$
  このゲートは、新鮮な勾配に対して 1 となり、選択されたカットオフ $\tau_{cut}$ において 0 まで滑らかに減衰し、$\tau > \tau_{cut}$ に対しては厳密に寄与をゼロにする。

主要な設計選択

1. バッファ前ゲート： スケーリング $\sigma(\tau)$ は、勾配が第一モーメント（$m$）および第二モーメント（$v$）バッファを更新する前に適用される。これにより、古くなった勾配が第二モーメント（$v$）を膨張させることが防がれ、さもなければ後続の新鮮な更新が減衰してしまうのを回避する。
2. フラグメント別変種（PA-CGAD）： 部分同期スケジューラ（例：Streaming DiLoCo）の場合、個々のフラグメントの年齢（$a_f$）を追跡し、$\max(\tau, a_f)$ に基づいて更新をゲートする。これにより、ネットワーク遅延とスケジューリング遅延の両方に対する統一された古さモデルが可能となる。
3. ハイパーパラメータ： この手法は、$\alpha$（減衰率）と $\tau_{cut}$（カットオフ）の 2 つのハイパーパラメータを導入する。デフォルト値（$\alpha=0.2, \tau_{cut}=32$）はモデル規模を超えて転移可能である。

3. 理論的貢献

本論文は、滑らかな非凸目的関数に対する理想的なゲート付き適応更新の非漸近収束境界を提供する。

• 古さバイアスの独立性： 収束境界の古さバイアス項は、実現された最大遅延 $\tau_{max}$ ではなく、ゲートハイパーパラメータ $\alpha$（具体的には $O(1/e^\alpha)$）に依存する。
• 標準的解析との対比： 非同期モメンタム-SGD の標準的な解析は、$\tau_{max}^2$ に比例する複雑さ項をもたらす。これを $\alpha$ に依存する項に置き換えることで、ゲートメカニズムは理論的に、$\tau_{cut} \ge 1/\alpha$ が満たされる限り、収束ペナルティを最悪ケースの遅延分布から切り離す。
• 範囲： 著者は、これは経験的統計から導出されるのではなく、有界な正規化ステップを仮定しているため、経験的 Adam に対する完全な収束定理というよりは、ゲートメカニズムの動機付けであると指摘している。

4. 実験結果

実験は、制御された遅延（最大 $\tau=16$）を用いて、25 M、1 B、7 B パラメータの 3 つの規模における Llama 風言語モデルの事前学習で行われた。

安定性と収束

• 1 B 規模： $\tau=16$ において、標準的なネステロフ手法は発散し（損失 $\approx 319$）、CGAD は 7.44 の損失を達成し、$\tau=0$ の性能（7.48）と同等であった。ネステロフ手法は $\tau \ge 8$ で収束性を完全に失った。
• 25 M 規模： CGAD および他の古さ認識ベースライン（Adam-Decay、SDM、Delayed Nesterov）はすべて安定していたが、ネステロフは $\tau \ge 4$ で発散した。この規模では、古さ認識手法間の差は限定的であった。

規模保険とロバスト性

最も顕著な発見は、メモリ制約条件下（bf16 重み、8 ビット AdamW）での7 B 規模において現れる：

• シード間分散： 最も近いベースラインであるAdam-Decay（コサインカットオフなしの CGAD）は、モデルが 25 M から 7 B にスケールするにつれて、シード間の標準偏差（$\sigma$）が 27 倍増加した（$\tau=8$ で 0.07 から 1.89 へ）。対照的に、CGAD の $\sigma$ はほぼ一定（約 0.38）であった。
• 単一ショットリスク： 「平均 + $\sigma$"を単一ショット展開リスクの代理指標として使用すると、7 B / $\tau=8$ における Adam-Decay のリスク（13.19）は、確率レベルの損失（10.83）を 2.4 ナット上回り、単一の生産実行における失敗の確率が高いことを示唆している。一方、CGAD のリスク（11.28）はこの閾値を十分に下回っていた。
• 失敗のメカニズム： 著者は、7 B における Adam-Decay の失敗を、完全にゼロにされない古くなった勾配から第二モーメントバッファ（$v$）に蓄積された二乗丸めノイズに起因すると帰属している。CGAD のコサインカットオフは、$\tau_{cut}$ を超える寄与を厳密にゼロにすることでこれを防止する。

調整感度

• ネステロフ vs CGAD： ネステロフを再調整（モメンタムを完全に除去）しても、高い遅延において調整済み CGAD 構成の性能には及ばなかった。調整済みネステロフは遅延領域ごとに異なるハイパーパラメータを必要としたのに対し、CGAD は全範囲にわたってロバスト性を維持した。
• ハイパーパラメータ感度： 最適な $\alpha$ は規模によってわずかにシフトする（25 M で 0.2、$\tau=8$ における 1 B で 0.4）が、デフォルトの $\alpha=0.2$ は全遅延スペクトルにわたってロバストである。

5. 意義と主張

本論文は、CGADが、特に大規模な制御遅延下でのネステロフ・モメンタムの脆さに対処する、非同期 DiLoCo 学習のための最もロバストなスケーリング選択であると主張する。

• 規模保険： 主要な貢献はコサインカットオフであり、これは「規模保険」として機能する。指数減衰のみ（Adam-Decay）は小規模モデルやフル精度設定では十分であるが、ハードカットオフは、圧縮されたオプティマイザ状態（8 ビット/bf16）で動作する最先端規模モデル（7 B 以上）にとって決定的である。これは、壊滅的なシード間分散と単一ショット失敗につながる古くなった分散の蓄積を防ぐ。
• ドロップイン互換性： CGAD は既存の Adam 実装に対して 1 行の変更のみを要し、規模を超えて転移可能なデフォルト値を持つ 2 つのハイパーパラメータを追加するだけである。
• 実用的影響： 最先端規模の非同期学習において、この手法は、成功して出荷される単一の学習実行と、失敗する（確率レベルの損失を上回る）実行との違いをもたらす。

著者は理論的範囲については謙虚であり、収束境界は理想的な更新に適用され、経験的優位性は特に DiLoCo ファミリーおよび Llama 風事前学習の文脈内で示されていることを認めている。また、実験で使用された内部ステップ数（$H=8$）は、生産環境の DiLoCo 設定（$H=100-500$）よりも低いことを指摘しているが、感度調査により、より高い $H$ でもランキングは維持されることを示唆している。