Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

技術要約：通信効率の高いパイプライン並列化のための学習可能な部分空間圧縮

問題提起
地理的に分散し、ヘテロジニアスなハードウェア上で大規模言語モデル（LLM）をトレーニングする場合、特にパイプライン並列化において、重大な通信ボトルネックに直面します。パイプライン並列化は、メモリ制限を克服するためにモデルパラメータをデバイス間で分割しますが、フォワードパスおよびバックワードパスの両方において、ステージ間での中間活性化（activations）の交換を必要とします。低帯域幅のネットワークでは、このステージ間の通信が支配的なボトルネックとなります。

Subspace Networks (SSNs) のような既存のアプローチは、直交行列によって定義される固定かつ共有された低ランクの部分空間へ活性化を投影することで、これを軽減しようと試みています。しかし、この手法は、重みを共有部分空間内に維持するための制約付き最適化や、修正されたオプティマイザの使用を含む、侵入的なアーキテクチャ変更を必要とします。経験的に、これらの固定されたグローバルな部分空間は、圧縮比が高まるにつれて大幅な性能低下を招きます。これは、訓練中の活性化の進化する幾何学的構造に適応できないためです。さらに、投影行列を制約なしにナイーブに学習させると、それらがスティフェル多様体（Stiefel manifold：直交行列の集合）から逸脱してしまい、効果的な圧縮に必要な等長性を破壊し、深刻な性能低下を引き起こします。

手法：多様体認識型投影学習 (Manifold Aware Projection Learning: MAPL)
著者らは、ステージ間の圧縮を、固定されたアーキテクチャ上の制約としてではなく、学習可能な幾何学的投影として扱う手法である Manifold Aware Projection Learning (MAPL) を提案しています。MAPLの核となる構成要素は以下の通りです：

1. ステージごとの学習可能な直交射影器: 各パイプラインステージの境界に、単一のグローバルな部分空間を強制するSSNとは異なり、MAP_Pはそれぞれ独自の学習可能な直交射影器 $A_p \in \mathbb{R}^{d \times r}$ を備えています。これらの射影器は スティフェル多様体 ($St(d, r)$) に制約されており、$A_p^\top A_p = I_r$ を保証します。これにより、各ステージは自身のタスクに最適な圧縮部分空間を発見し、継続的に適応することが可能になります。
2. 多様体制約付き最適化 (SPEL): 標準的な勾配更新中に射影器が多様体から逸脱することを防ぐため、著者らは スティフェル多様体上のスペクトル最急降下法 (Spectral Steepest Descent on the Stiefel Manifold: SPEL) を採用しています。このアルゴリズムは、毎ステップで接射影とリトラクション操作（Polar Express法による）を用いて射影器を更新し、直交制約が厳密に維持されることを保証します。
3. 因子分解されたアンカー埋め込み: 残差ストリームにおける高ランクでトークン依存のオフセットを、射影器の容量を使い果たすことなく扱うために、MAPLは因子分解されたアンカー埋め込みを導入しています。オフセットは、学習可能な低ランク行列 $E_{small}^p$ と凍結された直交行列 $P_p$ に分解されます。受信側のステージでこれらのオフセットを再構成するために、整数値のトークンIDのみが送信され、通信オーバーヘッドは無視できるほど小さくなります。
4. ベクトル量子化 (VQ): さらなる圧縮のために、本手法は低ランク投影された表現に対する残差ベクトル量子化 (RVQ) をサポートしています。辞書通信を償却するためのストリーミング・コードブック同期プロトコルを使用することで、最小限の帯域幅コストで高い圧縮比（最大16倍）を実現しています。

主な貢献

• ステージごとの部分空間学習: 共有されたグローバルな部分空間を強制することは、トークンが一致する設定において学習を劣化させることを本論文は示しています。代わりに、MAPLは各ステージが個別の、スティフェル制約を持つ圧縮部分空間を学習することを可能にし、これにより異なる深さにおける活性化の固有の低ランク構造により良く適合します。
• 多様体制約付き最適化: 著者らは、「多様体からの逸脱（manifold escape）」がナイーブな射影器学習の主要な失敗モードであることを特定しました。彼らは直交性を維持するためにSPELを導入し、制約なしの学習は固定ベースラインよりも厳密に劣ることを示しました。
• 効率的な再構成: 学習可能な射影器と因子分解されたアンカー埋め込みを組み合わせることで、MAPLは無視できるほどの通信オーバーヘッドでフルランクの活性化再構成を実現し、SSNが必要とする侵入的なアーキテクチャ分解を排除します。
• パレート・フロンティアの改善: 本手法は、検証クロスエントロピー対通信コストのパレート・フロンティアを追跡し、先行手法と比較して、性能低下を最小限に抑えつつ高い圧縮を実現しています。

実験結果
著者らは、DCLMコーパスを用いて、パイプライン構成 $P=4$ および $P=8$ ステージのデコーダーのみのLLaMAモデル（150Mから1Bパラメータまで）でMAPLを評価しました。

• 性能 vs 圧縮率: MAPLは一貫してSSNベースラインを上回っています。4倍圧縮における150Mモデルにおいて、MAPLは検証損失を非圧縮ベースラインの 0.84% 以内に抑えましたが、SSNベースラインは 8～12% 低下しました。
• スケーラビリティ: 1Bパラメータ規模での8倍圧縮において、MAPLは損失の低下をわずか 1.38～2.02% に抑えたのに対し、SSNベースラインは最大 26% の低下を被りました。
• ダウンストリームタスク: ゼロショット評価（HellaSwag, PIQA, ARC）において、MAPLは非圧縮ベースラインに極めて近く（500M/1Bスケールで平均精度差は約1ポイント以内）、一方でSSNベースラインは最大9ポイントの精度低下を経験しました。
• 極限圧縮: ベクトル量子化（MAPL+VQ）を組み合わせた場合、本手法は16倍の圧縮を達成しましたが、（例：1Bパラメータにおいて約2.3%の損失増加）という緩やかな追加の劣化となりました。ただし、ダウンストリームタスクの精度については、非量子化のMAPLと比較してより顕著な低下が見られました。

意義と主張
本論文は、ステージ間のトランスフォーマー活性化は、固定されたグローバルな部分空間によって課されるのではなく、タスクの目的関数と共同で学習可能な、適応的で低次元の構造を自然に持つことを主張しています。圧縮をスティフェル多様体上の学習可能な幾何学的投影として扱うことで、MAPLは、訓練の質とダウンストリームの性能を維持しながら、積極的な通信削減（4倍から16倍）を可能にします。本研究は、多様体制約付きの表現学習が、分散型基盤モデルの訓練における通信コストを削減するための実用的なメカニズムであることを浮き彫りにしており、既存の部分空間ベースのアプローチと比較して、性能と圧縮の優れたトレードオフを提供しています。

限界
著者らは、現在の実験が最大1Bパラメータまでのモデルに限定されていることを認めており、より大規模なスケールおよび現実世界のヘテロジニアスなネットワーク条件下でのさらなる評価が必要であるとしています。また、MAPLは性能の低下を大幅に軽減するものの、極端な圧縮率、特にベクトル量子化が適用された場合には依然として性能が低下することも指摘しています。