StoSignSGD：構造的確率性による「不偏」な修正でSignSGDを大規模言語モデル学習に対応させる

概要: SignSGD のような符号（サイン）に基づく最適化アルゴリズムは、分散学習や大規模な基盤モデルの学習において驚異的な性能を示し、大きな注目を集めています。実証的には優位であるにもかかわらず、SignSGD は ReLU、max-pool、mixture-of-experts によって現代の機械学習では至るところに存在する非滑らかな目的関数に対して発散することで知られています。この根本的な制約を克服するために、私たちは
\textbf{StoSignSGD} を提案します。これは、不偏な更新ステップを維持しつつ、符号演算子に構造的な確率性（stochasticity）を注入するアルゴリズムです。（オンライン）凸最適化の領域において、理論解析により StoSignSGD は SignSGD の非収束問題を厳密に解決し、下界に一致する鋭い収束率を達成することを示します。より難しい非凸・非滑らかな最適化に対しては、既存の定義を包含する一般化された停留測度（stationary measures）を導入し、StoSignSGD が次元要因によって、これまでで最良として知られている計算量の境界を改善することを証明します。実験的には、StoSignSGD は多様な大規模言語モデル（LLM）学習のレジームにわたって頑健な安定性と優れた効率を示します。特に、（AdamW が壊滅的に失敗する）低精度 FP8 の事前学習において、StoSignSGD は非常に安定しており、既存のベースラインに対して 1.44\times から 2.14\times の顕著な高速化を実現します。さらに、数学的推論タスクで 7B LLM をファインチューニングすると、StoSignSGD は AdamW と SignSGD の両方に対して大幅な性能向上をもたらします。最後に、その成功のメカニズムを解剖するために、あらゆる一般的な最適化手法をその不偏な符号ベースの対応物へと変換できる符号変換フレームワークを開発します。このフレームワークを用いて、StoSignSGD の中核となる構成要素を分解し、提案するアルゴリズム設計上の選択を実験的に検証するための包括的なアブレーションスタディを提示します。