HTMuon: Improving Muon via Heavy-Tailed Spectral Correction

本論文は、Muon 学習アルゴリズムの重みスペクトルが過度に正規化される問題を Heavy-Tailed Self-Regularization 理論に基づいて解決し、LLM の事前学習や画像分類において最先端の性能を達成する「HTMuon」という新しい最適化手法を提案するものである。

要約

この論文は、人工知能（AI）を学ぶための「新しい学習の教科書」のようなもの、HTMuonという新しい技術を紹介しています。

AI が本を読むように大量のデータから学習する際、その学習の進め方を決めるのが「オプティマイザ（最適化アルゴリズム）」という仕組みです。最近、Muon（ミュオン）という新しい学習方法が注目されていましたが、この論文は「Muon は素晴らしいけど、もっと良くできる！」と提案しています。

わかりやすくするために、いくつかの比喩を使って説明しましょう。

1. 従来の方法（Muon）の課題：「全員に同じ声で歌わせる合唱団」

AI の学習は、合唱団が歌を練習しているようなものです。

• Muon の仕組み: 指揮者が「全員、同じ大きさの声で歌って！」と指示します。これは、歌の「重要な部分（大きな声が必要なパート）」も、「ノイズ（雑音）」も、すべて同じ強さで扱おうとするものです。
• 問題点: 実際には、合唱団の中には「本当に上手なパート（重要な情報）」と「ただの雑音（ノイズ）」が混ざっています。Muon は「全員同じ声で」というルールを厳格に守りすぎるため、ノイズが混じったパートにも力を入れてしまい、重要なパートが埋もれてしまうことがあります。また、この「全員同じ」というルールが、AI の脳（重み行列）を「平らで退屈な状態」にしてしまい、本当の天才的な学習（複雑なパターンを捉えること）ができなくなってしまうのです。

2. 発見されたヒント：「天才的な脳は『偏り』を持っている」

研究者たちは、よく訓練された天才的な AI は、その脳の中に**「偏り（ヘビーテール）」**を持っていることに気づきました。

• 比喩: 天才的な音楽家は、すべての音を均等に鳴らすのではなく、「ここは激しく、ここは静かに」という極端な差をつけて表現します。この「偏り」があるほど、AI は複雑な世界を理解し、上手に学習できることがわかっていました。これを「ヘビーテール自己正則化（HT-SR）理論」と呼びます。
• Muon の欠点: 前述の「全員同じ声」というルールは、この「偏り」を消してしまい、AI を平凡な存在にしてしまうのです。

3. 解決策：HTMuon（「賢い指揮者」）

そこで登場するのが、この論文が提案するHTMuonです。

• どんな仕組み？: HTMuon は、Muon の「全員同じ声」というルールを少しだけ緩めます。具体的には、**「重要なパート（大きな声）はそのまま、ノイズのパート（小さな声）は少しだけ小さくする」**という調整を加えます。
• 魔法の調整: 数式では「パワースケール（p）」という値を使って調整しますが、イメージとしては「ノイズを少しだけ抑え、重要な情報を強調する」ようなフィルターをかけることです。
• 結果: これにより、AI の脳は「平らな状態」ではなく、「重要な部分に集中した、偏りのある（ヘビーテールな）状態」になります。その結果、AI はより深く、より正確に学習できるようになります。

4. 実際の効果：「より賢く、より速く」

この新しい方法（HTMuon）を試したところ、以下のような素晴らしい結果が出ました。

• LLM（大規模言語モデル）: 文章生成 AI を学習させたところ、従来の Muon よりも**「困惑度（PPL）」が大幅に下がりました**。これは、AI がより自然で正確な文章を生成できるようになったことを意味します。
• 画像認識: 写真の分類タスクでも、より高い精度を達成しました。
• 互換性: 既存の Muon の改良版（NorMuon など）の上に乗せるだけで使える「プラグイン」のような役割も果たし、さらに性能を向上させました。

5. 理論的な裏付け：「なぜこれが動くのか？」

研究者たちは、これが単なる運の良さではなく、数学的に正しいことも証明しました。

• 数学的な裏付け: HTMuon は、ある特定の数学的なルール（Schatten-q ノルム制約）の下で「最も急な坂を降りる（最急降下）」方法と等しいことがわかりました。つまり、AI が最も効率的に学習できる道筋を、数学的に保証しているのです。

まとめ

この論文は、**「AI を教えるとき、全員に同じことを強要するのではなく、重要な情報に集中し、ノイズを適切に抑える『偏り』を持たせることが、天才的な学習の鍵だ」**と教えています。

HTMuon は、その「偏り」を自動的に作り出す賢い学習方法です。これにより、今後登場する AI は、より少ないデータで、より深く、より賢く学習できるようになるかもしれません。

---

一言で言うと：
「AI の学習を『全員一様に』やるのではなく、『重要な部分に集中して』やるように調整した新しい方法で、AI がもっと賢くなったよ！」というお話です。

技術概要

論文「HTMuon: Improving Muon via Heavy-Tailed Spectral Correction」の技術的サマリー

本論文は、大規模言語モデル（LLM）の学習において近年注目されているオプティマイザ「Muon」の改良版として、HTMuon（Heavy-Tailed Muon）を提案するものです。Muon の更新則が持つ「重たい尾部（Heavy-tailed）」な特性の欠如を指摘し、Heavy-Tailed Self-Regularization (HT-SR) 理論に基づいて、より重たい尾部を持つスペクトルを誘発する新しい手法を開発しました。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細をまとめます。

---

1. 問題定義 (Problem)

近年、LLM の学習において、パラメータ間の幾何学的な依存関係を捉える行列ベースのオプティマイザ「Muon」が優れた性能を示しています。Muon は、モメンタム行列の特異値分解（SVD）を行い、すべての特異値を 1 に設定して直交化（Orthogonalization）する更新則を採用しています。

しかし、著者らは以下の 2 つの根本的な問題点を指摘しました。

1. ノイズ支配方向への過剰な強調:
   Muon の更新則は、モメンタム行列のすべての特異値を 1 に強制します。これにより、小さな特異値を持つ方向（通常はノイズが支配的な方向）に対しても、大きな特異値を持つ方向（信号が支配的な方向）と同様の重みを割り当ててしまいます。これは学習の初期段階では有効かもしれませんが、学習の後半段階ではノイズへの感度を高め、モデルの性能を制限する可能性があります。
2. 重たい尾部スペクトルの抑制:
   深層学習の理論（HT-SR 理論）によれば、高品質に学習されたニューラルネットワークの重み行列は、その固有値スペクトル（ESD: Empirical Spectral Density）が「重たい尾部（Heavy-tailed）」を持つ傾向があります。尾部が重いほど（べき乗則の指数 $\alpha$ が小さいほど）、モデルの汎化性能が高いことが知られています。
   一方、Muon の直交化更新則は、更新行列のスペクトルを「軽い尾部（Light-tailed）」にし、結果として学習された重み行列のスペクトルも軽くなってしまいます。これは HT-SR 理論の観点から、モデルが到達可能な最終的な性能の上限を制限していると考えられます。

2. 手法 (Methodology)

著者は、Muon のパラメータ間の依存関係を捉える能力を維持しつつ、更新行列と重み行列のスペクトルを「より重たい尾部」にするオプティマイザ HTMuon を提案しました。

核心的なアイデア

HTMuon は、Muon の更新則における特異値の処理を以下のように変更します。

• Muon: モメンタム行列 $M_t$ の SVD を行い、特異値行列 $\Sigma_t$ のすべての要素を 1 に置き換えて更新を行います（$O_t = U_t V_t^\top$）。
• HTMuon: 特異値 $\sigma_i$ を $p$ 乗（$0 < p < 1$）して重み付けを行います。
  $$O_t = U_t \Sigma_t^p V_t^\top$$
  ここで、$p$ はハイパーパラメータです。

理論的根拠と効果

• $p$ の役割:
  • $p=1$ の場合、SGDM（確率的勾配降下法＋モメンタム）に帰着し、パラメータ間の依存関係を捉えられません。
  • $p=0$ の場合、Muon に帰着し、軽い尾部のスペクトルになります。
  • $p \in (0, 1)$: 行列ベースの依存関係捕捉能力を維持しつつ、小さな特異値（ノイズ方向）を相対的に減衰させ、大きな特異値（信号方向）を強調します。これにより、更新行列および最終的な重み行列のスペクトルがより重たい尾部（Heavy-tailed）を持つようになります。
• 理論的解釈:
  著者は、HTMuon が Schatten-$q$ ノルム制約下における最急降下法（Steepest Descent）に対応することを理論的に示しました（Muon は Schatten-$\infty$ ノルム制約下での最急降下法に対応します）。また、滑らかな非凸環境における収束解析も行い、Muon や SGDM と同等のサンプル複雑性の上限を持つことを証明しています。

実装上の工夫

SVD の計算コストを削減するため、2 つの高速化実装も提案しています。

1. 間隔更新: 全てのステップで HTMuon を適用するのではなく、一定ステップごとに適用し、それ以外は Muon を使用する。
2. 数値反復法による近似: 完全な SVD の代わりに、Newton-Schulz 反復法を用いて分数べき乗（Fractional Power）を近似する「HTMuon NS」を実装し、計算オーバーヘッドを大幅に削減しました。

3. 主要な貢献 (Key Contributions)

1. Muon の限界の理論的・実証的解明:
   Muon の直交化更新則が、ノイズ支配方向への偏りを生み出し、HT-SR 理論に基づく重たい尾部スペクトルの形成を阻害していることを示しました。
2. HTMuon の提案と実証:
   重たい尾部スペクトルを誘発する新しいオプティマイザ HTMuon を提案し、LLM の事前学習（LLaMA シリーズ）および画像分類（ResNet, ViT）において、Adam、AdamW、Muon、およびその派生手法（NorMuon, AdaMuon, COSMOS など）を凌駕する性能を達成しました。
   • LLaMA-60M (C4 データセット): Muon に対して Perplexity が 0.92 改善。
   • LLaMA-135M (C4 データセット): Muon に対して Perplexity が 0.98 改善。
   • 画像分類: CIFAR-100/10 および ImageNet-1K において、Muon や SGDM よりも高い精度を達成。
3. 理論的保証:
   HTMuon が Schatten-$q$ ノルム制約下での最急降下法に対応すること、および非凸滑らかな環境における収束性を理論的に証明しました。
4. 実用性の確保:
   計算コストの削減に向けた高速化実装（HTMuon NS）を開発し、大規模モデル（LLaMA-1B）の学習においても Muon よりも優れた性能を維持しつつ、実行時間を現実的な範囲に抑えることを示しました。

4. 実験結果 (Results)

• LLM 事前学習:
  C4 データセットでの LLaMA モデル（60M, 135M, 350M, 1B）および OpenWebText での GPT-2 small の学習において、HTMuon はすべてのベースライン（Adam, AdamW, Muon, Sophia, GaLore など）に対して一貫して低い Perplexity を記録しました。特に、Muon 単体と比較して、Perplexity の改善幅は最大で 0.98 に達しました。
• 画像分類:
  CIFAR-100/10 および ImageNet-1K における ResNet や ViT の学習でも、HTMuon は Muon や SGDM よりも高い精度を達成しました。
• スペクトル分析:
  学習されたモデルの重み行列の固有値スペクトルを分析した結果、HTMuon を使用した場合、Muon に比べてべき乗則の指数 $\alpha$ が小さく（より重たい尾部）、かつスペクトルノルムやフロベニウスノルムが小さくなる傾向が確認されました。これは HT-SR 理論の予測と一致し、汎化性能の向上と相関しています。
• ダウンストリームタスク:
  LLaMA-1B モデルの事前学習後、7 つの常識推論タスク（Commonsense Reasoning）でゼロショット評価を行ったところ、HTMuon は平均スコアで Muon よりも 1.05 ポイント上回る結果となりました。

5. 意義と結論 (Significance)

本論文の HTMuon は、オプティマイザ設計において「幾何学的依存関係の捕捉」と「重たい尾部スペクトルの誘発」という、一見相反する可能性のある 2 つの目標を両立させた画期的なアプローチです。

• 理論と実践の融合: HT-SR 理論という統計力学・ランダム行列理論に基づく知見を、実用的なオプティマイザの設計に直接反映させた点に大きな意義があります。
• 既存手法の汎用性: HTMuon は単独で強力であるだけでなく、既存の Muon 派生手法（NorMuon など）の上にプラグインとして機能し、さらに性能を向上させることも示されました。
• 大規模学習への適用: 計算コストの課題を克服する高速化手法を提案しており、LLaMA-1B 規模以上の大規模モデルの学習においても有効であることを実証しました。

結論として、HTMuon は LLM の学習効率と最終的なモデル品質を向上させるための、理論的裏付けのある強力な新しいオプティマイザとして位置づけられます。