確率的符号およびスペクトル降下法のための非ユークリッド勾配ノイズスケールを用いた適応的バッチサイズ

TL;DR（結論）

従来の適応的バッチサイズ制御はSGDのユークリッド幾何学を前提としていたが、本研究ではsignSGDやspecSGD（Muon）といった非ユークリッド幾何学を用いる最適化手法に対応した新しい勾配ノイズスケール（GNS）を導出した。 提案手法では、各最適化手法の双対ノルム（L1ノルムや核ノルム）を用いて勾配のノイズを測定し、分散データ並列学習におけるローカルなミニバッチ勾配を活用することで、計算負荷を抑えつつリアルタイムで正確なGNSを推定する仕組みを構築した。 1億6000万パラメータのLlamaモデルを用いた実験において、提案する非ユークリッドGNSに基づく適応的バッチサイズ制御は、従来の固定バッチサイズと比較して検証損失を維持したまま、SignumやMuonの学習ステップ数を最大66%削減することに成功した。

なぜこの問題か

近代的な機械学習システムにおいて、ハードウェアの利用効率を最大化することは極めて重要な課題であり、一般的にはバッチサイズを大きくすることでGPUのスループットを向上させている。しかし、バッチサイズを固定したまま大規模化すると、サンプルあたりの学習進捗が低下し、最終的なモデルの品質やサンプル効率が悪化するというトレードオフが存在する。この問題を解決するために、多くの現場では学習の進行に合わせてバッチサイズを段階的、あるいは線形に増加させるヒューリスティックな手法が採用されているが、これらは調整が困難で壊れやすいという欠点がある。より原理的なアプローチとして、勾配の推定誤差を定量化する勾配ノイズスケール（GNS）を用いてバッチサイズを動的に調整する戦略がある。 しかし、既存のGNS理論は確率的勾配降下法（SGD）が前提とするユークリッド幾何学に基づいて導出されており、現代の深層学習で多用される非ユークリッド的な最適化手法との間に根本的な不一致が生じている。…

核心：何を提案したのか

本研究の核心は、signSGD（Signum）やspecSGD（Muon）といった最適化手法の幾何学的構造に自然に適合する、非ユークリッド的な勾配ノイズスケールを導出したことにある。具体的には、各最適化手法が採用するノルムに対応した双対ノルムを用いて、勾配の推定誤差を測定する新しい指標を定義した。signSGDに対してはL1ノルム（マンハッタン距離）に基づくGNSを、行列の重みを扱うspecSGDに対しては核ノルム（Schatten-1ノルム）に基づくGNSをそれぞれ導出している。これにより、最適化手法の探索方向のバイアスを理論的に制御し、サンプル効率を最大化するクリティカルバッチサイズ（CBS）を特定することが可能となった。…

続きはログイン/プランで閲覧できます。