技術要約：Hyperballの最適化

問題提起

Muonのような行列ベースのオプティマイザは、言語モデルの事前学習においてAdamWに対して大幅な高速化を実現している。しかし、経験的な観察によれば、モデルサイズやデータスケールが増大するにつれ、特に標準的な定数デカップル重み減衰（constant decoupled weight decay）を使用している場合に、これらの利得が減少することが示されている。大規模な設定（例：1.2Bパラメータのモデル）では、AdamWに対するMuonのスピードアップは、約30%から約10%へと縮小する。この性能低下が、より高い計算スケールにおいても高い効率性を維持できる最適化戦略の探索を動機付けている。

核心的な問題は、デカップルされた重み減衰のメカニズムにある。スケール不変なレイヤー（正規化を備えた現代のTransformerにおいて一般的）において、重み減衰は単なる正則化としてではなく、間接的な「角度学習率（angular learning rate）」、すなわち重み行列の方向が変化する速度の制御器として機能する。理論的分析は、重み減衰がハイパーパラメータ（$\eta$ および $\lambda$）によって決定される平衡状態へと重みのノルムを駆動することを示唆している。この平衡ノルムは、結果として実効的な角度ステップサイズを決定する。しかし、この間接的なキャリブレーションはスケールが変わるにつれて最適ではなくなる可能性があり、それが観察されたオプティマイザの性能低下につながっている。

手法：Hyperball

著者らは、重み減数のスケーリング限界に対処するために設計された、シンプルなオプティマイザ・ラッパーであるHyperballを提案する。Hyperballは、重み減衰による重みノルムの間接的かつソフトな制御を、重み行列とその更新に対するフロベニウスノルム（Frobenius norm）への明示的な制約に置き換える。

コアメカニズム

パラメータ行列 $W_t$ に対して更新 $u_t$ を生成するベースとなるオプティマイザ（例：AdamまたはMuon）が与えられたとき、Hyperballは以下の2つの条件を強制する：

1. 重み行列のフロベニウスノルム $\|W_t\|_F$ は、一定の半径 $R$（初期重みのノルムとして初期化される）に固定される。
2. 更新方向は正規化され、ステップサイズは明示的にスケーリングされる。

更新ルールは以下のように定義される：
$$W_{t+1} = R \cdot \text{Normalize}\left( W_t - \eta_t R \cdot \text{Normalize}(u_t) \right)$$
ここで、$\text{Normalize}(X) = X / \|X\|_F$ である。

幾何学的には、これは最適化の軌跡を半径 $R$ の超球（hypersphere）の表面上に制約する。各ステップは、正規化された更新方向に、$\eta_t R$ に比例する距離で移動し、直ちに球体へと投影される。これにより、重みの大きさと更新の方向を切り離し、オプティマイザが主に重みの方向を通じてナビゲートすることを可能にする。

実装の詳細

• 適用範囲: Hyperballは、prenorm TransformerにおけるAttentionおよびMLPの重み行列に適用される。
• 除外対象: 埋め込み（Embeddings）、正規化ゲイン（normalization gains）、およびスカラーパラメータは、それらのノルムが意味的な情報を持つため、標準的なオプティマイザ（例：Adam）を使用して更新される。
• 計算量: フロベニウス球への投影は、行列あたり $O(N^2)$ であり、SVDを必要とするスペクトルノルム投影（$O(N^3)$）よりも計算コストが低い。著者らは、LLMにおいて典型的な高安定ランク（high stable-rank）のレジームでは、フロベニウス制約がスペクトル制約と同様に振る舞うと指摘している。

主な貢献

1. Hyperballラッパー: 重みと更新のノルムを固定するバリアントへと、あらゆるベース・オプティマイザを変換する汎用的なラッパー。
2. 理論的洞察: 本論文は、スケール不変な損失関数において、デカップルされた重み減衰が間接的に角度ステップサイズを設定するという理論的枠組みを提供している。Hyperballは、重み減衰の平衡ダイナミクスに依存するのではなく、半径 $R$ と正規化された更新長を直接固定することで、このメカニズムを明示化する。
3. スケーラビリティの向上: 標準的な重み減衰ベースラインが性能を維持できなくなるような大規模なスケールにおいても、最適化のスピードアップを維持できることを示している。
4. 学習率の転送性: Hyperballは、重み減衰ベースラインと比較して、異なるモデルの幅や深さの間での学習率の転送性を大幅に改善する。

実験結果

著者らは、Qwen3スタイルのモデル（最大1.2Bパラメータ）およびより大きなスケール（最大8Bパラメータ）を用いて、様々なベンチマークでHyperballを評価した。

• トークン等価スピードアップ: 1.2Bパラメータのモデルにおいて、Muon Hyperball (MuonH) は、AdamWベースラインに対して20–30%のトークン等価スピードアップを達成した。対照的に、標準的な重み減衰を用いたMuon (MuonW) は、このスケールでは約10%のスピードアップしか達成できなかった。MuonHとMuonWの差は、訓練時間が長くなるにつれて拡大した。
• ベンチマーク性能:
  • Marin speedrunベンチマークにおいて、AdamHおよびMuonHは、約10%大きい重み減衰ベースラインと同等の性能を示した。
  • 公開されている modded-nanogpt Track 3 最適化ベンチマークにおいて、Hyperballバリアント（KL-SOAP-Hを含む）は、重み減衰のカウンターパートよりも速く低い検証損失に到達した。例えば、KL-SOAP-Hは3125ステップで検証損失3.2780に到達し、チューニングされたMuonWベースラインを上回った。
• ハイパーパラメータの転送:
  • 深さのスケーリング: 深さを4から512までスイープした際、Hyperballバリアントの最適な学習率のドリフトは約1.4倍であった。これに対し、AdamWおよびMuonWでは2–4倍であった。
  • 幅のスケーリング: 同様に、幅を128から2048まで変化させた際も、Hyperballは1.4倍のドリフトのみで安定した学習率ウィンドウを維持した。
• オーバーフィット設定: データスケーリング実験（130Mパラメータモデル、1B–128Bトークン）において、MuonHはMuonW（$L_\infty = 3.079$）と比較して、より低い漸近検証損失（$L_\infty = 3.065$）を達成した。

意義と主張

本論文は、Hyperballが、特に行列ベースのオプティマイザを使用する場合において、大規模言語モデルの学習に対する、デカップルされた重み減に代わる、より堅牢で効率的な選択肢を提供すると主張している。

• 明示的な制御: 重みノルムと更新ノルムを明示的に固定することで、Hyperballは、角度学習率を間接的に制御するために重み減衰に依存することを排除する。これにより、より予測可能な最適化ダイナミクスが実現される。
• スケーラビリティ: この手法は、モデルおよびデータスケールが増大しても、行列ベースのオプティマイザ（Muonなど）の効率性の利点を維持することに成功しており、現在の事前学習の実践における重要なボトルネックに対処している。
• 設計上の選択: 著者らは、Hyperballを、重み減衰のスケーリングへの依存（これは転送性に影響を与える）を明示的なパラメータ（半径 $R$）へと変える設計上の選択肢として位置づけており、それによって異なるモデルアーキテクチャ間での学習率の転送性を向上させている。
• 概念実証: 本研究は、重みノルムを明示的に制御するオプティマイザを設計することで、重み減衰を完全に排除できる可能性があるという概念実証であり、正規化ネットワークにおけるオプティマイザ設計の新しい方向性を示唆している。

著者らは、特定の制約の選択については謙虚な姿勢を保っており、フロベニウスノルムは計算効率が高く理論的にも裏付けられているものの、特定のモデルやオプティマイザには他の制約（スペクトル、行方向、列方向など）の方が適している可能性があると述べている。また、重み正規化された訓練の軌跡に関する、より厳密な理論の必要性についても強調している。