技術要約：重み空間の対称性を利用した曲率近似

問題定義

損失関数の曲率（ヘッセ行列）を正確に推定することは、二次の最適化、ベイズ深層学習、継続学習、モデル圧縮を含む、多くの機械学習のサブフィールドにおいて不可欠である。しかし、現代の深層ネットワークに対して、完全なヘッセ行列を構築、保存、および逆行列計算を行うことは、計算量およびメモリの観点から極めて困難である。 (E)KFAC、Shampoo、Soapといった既存の近似手法は、特定のブロック対角構造やクロネッカー分解構造を仮定することでスケーラビリティを実現してきたが、これら特定の構造が「なぜ」機能するのか、あるいはモデルのアーキテクチャに基づいてどのように体系的に導出できるのかを説明する、原理的な理論的枠組みがこの分野には欠けている。さらに、従来の作業は、ニューラルネットワークに固有の（例：隠れユニットの置換不変性など）既知の重み空間の対称性が課す構造的制約を、概して見落としてきた。

手法

著者らは、重み空間の対称性によって生成される「群の軌道（group orbit）」上で解析的に平均化することにより、ヘッセ行列を近似する新しいフレームワークを提案している。核心となる洞察は、多くの損失関数が特定のパラメータ変換（群作用）に対して不変であるという事実に基づいている。これは、損失の勾配がこれらの変換の下で共変（equivariant）であることを意味する。

1. 軌道平均と勾配の共変性

現在のパラメータ集合 $w$ と勾配 $g$ が与えられたとき、著者らは、対称群 $G$ からの対称変換を適用することによって到達可能なすべての $(w^{(i)}, g^{(i)})$ のペアを検討する。損失は不変であるため、軌道上の任意の点における勾配は、群作用によって元の勾配に関連付けられる。これらの勾配を軌道全体で平均化することにより、本手法は複数の勾配計算を必要とせずに曲率情報を集約する。

2. 構造化されたヘッセ行列近似

本手法は、パラメータの差と勾配の差の軌道平均に基づくセカント条件（secant condition）を解くことにより、構造化されたヘッセ行列近似 $\hat{H}$ を導出する。

• 交換子代数（Commutant Algebra）: 得られるヘッセ行列近似は、対称群の交換子代数に属する。これは、近似が、単位行列と全要素1の行列のクロネッカー積からなる少数の基底要素（疎なバイナリテンソル）の線形重ね合わせとして表現できることを意味する。
• 因数分解: 全体行列を保存する代わりに、本手法はこれらの基底要素の重みとなるスカラーまたは行列の「因子（factors）」のみを保存する。これにより、扱いやすい推定、保存、および逆行列計算が可能になる。
• 簡略化: 著者らは、パラメータの差の軌道平均が単位行列に比例すると仮定する実用的な簡略化を提案している。これにより、曲率推定値 $\hat{H}_g = S_g^{1/2}$ が得られる（ここで $S_g$ は勾配の外積の軌道平均である）。

3. 「Symo」オプティマイザ

このフレームワークは、Symo（Symmetry-based optimizer）と呼ばれる新しいオプティマイザとして具体化される。更新則は $w_{t+1} = w_t - \eta (\hat{H}_g)^{-1} g_t$ である。著者らは、ネットワークの対称性の記号的仕様を受け取り、必要な因子を計算するための効率的なPyTorchルーチンを生成するジャストインタイム・コンパイラを提供している。

主な貢献

1. 新しいヘッセ行列近似: 対称群の軌道上で平均化することにより、単一の勾配計算から構造化されたヘッセ行列推定値を導出する手法。得られる近似は証明可能なほど高度に構造化されており、交換子代数内に存在する。
2. ShampooとMuonの理論的統一: 本フレームワーク内の特定の対称群の選択が、Shampoo および Muon オプティマイザの更新則を正確に復元することを証明した。
   • 対称群を特定の符号付き置換部分群（例：入力/出力層には恒等写像を適用し、隠れ層には符号付き置換を適用する）に制限することで、導出される曲率行列は、Shampoo/Muonと数学的に等価な、クロネッカー因子を持つブロック対角行列となる。
   • これにより、これらのオプティマイザの有効性に対する原理的な説明が提供される。すなわち、これらは暗黙的に重み空間の対称性を利用しており、近似精度と計算コストの間の「スイートスポット」を突いているのである。
3. 複雑性の制御: フレームワークは、ヘッセ行列近似の複雑さが、選択された対称群のサイズによって直接制御されることを示している。より大きな群は、より低次元の交換子代数（より少ない因子、より軽い計算）をもたらすが、より大きな軌道（潜在的に低い局所精度）を生む。より小さな群は、より多くの項をもたらすが、より高い局所的な忠実度を実現する。
4. 実証的検証:
   • ヘッセ行列の品質: 小規模なMLPを用いた実験により、提案された近似（特に簡略化された $\hat{H}_g$）が、負の勾配方向における真のヘッセ行列に対して高いコサイン類似度を達成し、軌道平均されたパラメータの推定を上回り、Shampooの性能に匹敵することを示した。
   • オプティマイザの等価性: MNISTオートエンコーダおよびnanoGPT（小規模な言語モデル）において、SymoオプティマイザはMuonとほぼ同一の訓練曲線を生み出し、特定の群制約下での理論的等価性を検証した。両者はAdamを大幅に上回る性能を示した。

結果

• 精度とコストのトレードオフ: 本論文では、異なる群の選択を経験的に調査している。「より大きな群」（すべての層の置換を許可する）は、依然としてAdamよりも優れた性能を示すが、内部層に置換を制限しブロック対角近似を使用する「より小さな群」の性能には及ばない。これは、アーキテクチャ固有の不変性に一致するように群を制限することが、最適な性能を得るために重要であることを示唆している。
• ランク欠損: 著者らは、完全な軌道平均行列 $S_w$ がランク欠損を起こしやすく、それが逆行列計算を不安定にすることを指摘している。提案された簡略化（$S_w \approx I$）はこの問題を回避し、完全な正定値解よりも優れた曲率推定を経験的に提供する。
• スケーラビリティ: 本手法はTransformerアーキテクチャ（nanoGPT）へのスケールにも成功しており、対称性を考慮したアプローチが単純なMLPを超えて適用可能であることを示している。

意義と主張

本論文は、既存の二次オプティマイザを見るための統一的な理論的レンズを提供することを主張している。ShampooやMuonを対称性に基づくヘッセ行列近似の特定事例として位置づけることで、本研究はこれらの手法を、ヒューリスティックなエンジニアリングから原理的な数学的導出へと昇華させている。

著者らは、本フレームワークが以下の機能を提供すると述べている：

• 計算予算に応じて適切な対称群を選択することにより、カスタマイズされた曲率近似を構築するための体系的なレシピを提供する。
• モデルの対称性から構造化された曲率への原理的な経路を提供し、異なるモデルクラスのための新しいオプティマイザの創出を促す可能性がある。
• 最適化以外の用途（不確実性推定、継続学習、モデル圧縮など）への応用を示唆しているが、これらは結果として示されたものではなく、将来の展望として提示されている。

本論文は、学習済みの重みにおけるテイラー展開の精度が厳密に限定されていないことや、プリコンディショナと真のヘッセ行列の間の固有ベクトル整列に関する形式的な保証が未解決の問題であることなど、限界についても謙虚に認めている。また、フレーム本フレームワークは因子の導出を自動化するものの、任意の新しいアーキテクチャに対して最適な群を選択することは、依然としてユーザー依存のタスクである。