技術要約: LoRDO – 非同期通信を伴う分散低ランク最適化

1. 問題提起

分散データ並列（DDP）による基盤モデルの分散学習は、インターコネクトの帯域幅と、オプティマイザの状態（例：Adamのモーメンタ）を保存するために必要なメモリオーバーヘッドによって、ますます制約を受けています。非同期通信（ローカル更新）を採用する戦略は、同期頻度を減らすことができますが、フルランクのオプティマイザ状態に伴うメモリおよび通信コストというボトルネックに依然として直面しています。

低ランク適応型オプティマイザ（GaLoreやLDAdamなど）は、勾配を低ランク部分空間に投影することで、これらの制約を緩和する道を提供します。しかし、これらの手法を非同期同期レジーム（infrequent-synchronization regime）へと一般化するには、重大な最適化上の課題が生じます：

1. ローカル投影（Local Projections）: 個々のワーカー上で、確率的勾配を用いてローカルに低ランク投影を計算することは、ワーカーあたりの実効バッチサイズが小さいために、大きなノイズを導入します。これは、サブスパン（部分空間）の低品質化と性能低下を招きます。
2. グローバル投影（Global Projections）: DDPのように、集約された擬似勾配（pseudo-gradients）から導出されたグローバルな投影行列を使用する方法は、より安定した基底を提供します。しかし、ローカル更新レジームにおいては、このアプローチは最適化の軌跡を固定された低ランク部分空間内に永久に制限してしまいます。フルランクの信号がなければ、投影基底が「刷新」されたり、初期の部分空間の外側にある方向を探索したりすることができないため、最適化が停滞します。

核心となる問いは、「通信効率の高い学習において、メモリおよび通信のオーバーヘッドを維持しつつ、いかにして部分空間の停滞を回避する高性能な低ランクオプティマイザを設計できるか？」という点です。

2. 手法: LoRDO

著者らは、低ランク最適化と非同期同期を統合するフレームワークである LoRDO (Low-Rank Distributed Optimization) を提案しています。この手法は、3つの主要なメカニズムを通じて、投影の安定性と部分空間の探索との間の緊張関係に対処します。

A. グローバル投影行列

各ワーカーが自身のローカルな確率的勾配に基づいて独自の投影行列 $Q_m$ を決定する（これは高い分散に苦しむ）代わりに、LoRDOはグローバル投影行列 $Q_t$ を採用します。この行列は、$K$ 回のローカルステップにおけるモデルパラメータの総変化量を示す集約された擬似勾配 ($\Delta_t$) を用いて、同期境界時に計算されます。

• 利点: 擬似勾配は実効バッチサイズが $M \times B$（$M$ はワーカー数、$B$ はバッチサイズ）となり、ローカル勾配と比較して分散が抑えられた、より安定した信号を提供します。また、擬似勾配に組み込まれた曲率情報を取り込むことで、より高品質な基底推定を可能にします。

B. フルランク準双曲モーメンタム (Full-Rank Quasi-Hyperbolic Momentum: QHM)

グローバル投影に固有の「停滞する部分空間」の問題を解決するために、LoRDOはローカル更新にフルランク準双曲モーメンタム項を注入します。

• メカニズム: 更新ルールは、低ランクモーメンタムと、第2モーメンタムでスケールされたフルランク勾配信号を組み合わせます。
  $$\text{Update} \propto (1-\omega)\frac{\hat{G}_t}{\mu(\sqrt{\hat{v}_t} + \epsilon)} + \omega Q_t \left[ \frac{u_t}{\sqrt{v_t} + \epsilon} \right]$$
• 効果: この注入により、集約された擬似勾配が固定されたランク-$r$の部分空間に閉じ込められるのを防ぎます。これにより、異なるワーカーからのランク-$r$の部分空間を時間的に集約することが可能になり、最終的な完全な部分空間探索を実現します。これにより、通信オーバーヘッドを増やすことなく、全解空間の探索能力を回復させます。

C. エラーフィードバックとモーメンタム回転

LoRDOは、低ランク圧縮中に失われた情報を補うために、ローカルでエラーフィードバックを取り入れています。さらに、新しい投影行列の計算に続いて、オプティマイザの状態（モーメンタム）を**回転（rotate）**させ、それらが同じ部分空間上に蓄積され続けるようにします。これはLDAdamから適応された手法です。

3. 主な貢献

1. 投影失敗モードの分析: 著者らは、ワーカーの小さなバッチサイズによる高分散がローカル投影に悪影響を与えること、またフルランクの信号がないグローバル投影が部分空間の停滞を誘発することを実証しました。
2. フル部分空間探索の回復: LoRDOは、集約された擬似勾配から派生したグローバル投影を維持しつつ、ローカル更新にフルランク勾配信号（QHM経由）を注入する低ランクオプティマイザを導入しました。これにより、通信やメモリのオーバーヘッドを増やすことなく、停滞を防ぎます。
3. DDPへの近似と効率性: 本フレームワークは、同期型低ランクDDPに対してニア・パリティ（パープレキシティの差が1%未満）を達成しつつ、オプティマイザのメモリと通信を最大 $8\times$ ～ $12\times$ 削減します。低ランクを必要とする厳しいメモリ制約下では、LoRDOはDDPよりも優れたパープレキシティを示します。

4. 実験結果

実験は、SmolLM2データ混合を用いて、16M、125M、720Mのパラメータを持つデコーダーのみのトランスフォーマーを用いて行われました。

• 停滞 vs 探索: 16Mモデルにおいて、QHMのないグローバル投影は学習に失敗（パープレキシティ 46.8）し、DDP (36.6) と比較して性能が低下しました。QHMを追加することで性能が回復し、LoRDO-Globalがローカル変種を上回ることが確認されました。
• 投影の品質: LoRDO-Globalは、特に低いランク（$r=8, 16$）において、集約された擬似勾配から得られる高品質な投影基底により、優れた安定性と性能を示しました。これは、ローカルのバッチサイズの変動に対しても、ローカル投影手法より感度が低いことを示しています。
• 720Mへのスケーリング: 720Mスケールにおいて、LoRDO-Global ($r=256$) は、低ランクDDP ($r=256$) に対してニア・パリティ（パープレキシティの差が1%未満：10.41 vs 10.34）を達成しました。また、ダウンストリームタスクの精度（ARC, HellaSwag, MMLU, PIQA）においてもDDPと同等の性能を示しました。
• 効率性の向上:
  • 通信: LoRDOは、低ランクDDPと比較して約 10倍、フルランクDDPと比較して約 25倍 の通信オーバーヘッドを削減します。
  • メモリ: フルランクの対応物と比較して、オプティマイザ状態のメモリオーバーヘッドを 8倍 削減します。
  • ウォークラックタイム（実時間）: 低帯域幅環境（10 Gbpsと想定）において、LoRDOは同期頻度の減少により、DDPを大幅に上回るパフォーマンスを発揮します。

5. 意義と主張

本論文は、LoRDOが、通信効率の高い低ランク最適化における「投影の安定性」と「部分空間の探索」という根本的な対立を解決することを主張しています。集約された擬似勾配による投影の安定性と、フルランク準双曲モーメンタムによる探索を利用することで、LoRDOは以下を可能にします：

• スケーラブルな学習: 通信およびメモリコストを大幅に削減しながら、大規模（最大720Mパラメータ）での競争力のある性能を実現。
• リソース効率: フルランクのオプティマイザが困難な、限られたリソース（低メモリ/低帯域幅）を持つハードウェアでの事前学習の実現可能性。
• 分散学習: モデルの学習を単一のデータセンターを超えて拡張することにより、効率的な分散環境における学習への道を開く。

著者らは、収束の形式的な証明が欠けていることや、現在は行列ベースのオプティマイザ（例：Muon）ではなく適応型オプティマイザ（Adam）に焦点を当てているといった限界についても言及しており、これらは今後の課題として残されています。しかし、小規模モデルで観察された傾向は、より大きなスケールでも保持されると考えています。