Rethinking Bregman Divergences in Kronecker-Factored Optimizers

技術要約：Kronecker分解を用いたBregmanダイバージェンスの再考

問題提起
Shampoo型のオプティマイザは、計算複雑性を削減するために、勾配共分散行列 $C = E[\text{vec}(G)\text{vec}(G)^\top]$ をクロネッカー積構造 $L \otimes R$ を用いて近似する。近年の研究（Lin et al., 2026）は、これらを $C$ と $L \otimes R$ の間のBregman行列ダイバージェンスの最小化問題として定式化することで、様々なShampooのバリアントを統一した。しかし、一般的な高密度な共分散行列に対しては、厳密なクロネッカー分解は不可能であり、近似誤差は避けられないという決定的なギャップが残されている。異なるBregmanダイバージェンスがこの誤差を共分散スペクトラム全体にどのように分配し、結果として得られるプリコンディショナがスペクトラムのすべての部分を等しく信頼しているのかどうかは不明である。さらに、勾配共分散は真のヘシアン（Hessian）とは大きく異なることが多いため、すべてのスペクトル成分がプリコンディショニングに適しているわけではない。ここで対処される中心的な問いは、「厳密なクロネッカー適合が不可能な場合、構造化されたオプティマイザは、共分散スペクトラムのどの部分を近似し、信頼すべきか？」である。

手法および分析
著者らは、Bregman近似のスペクトル幾何学と、共分散固有空間の経験的な信頼性という2つの主要な分析的観点を通じてこの問題にアプローチしている。

1. Bregman近似のスペクトル幾何学：
   本論文は、Frobenius、von Neumann、およびLogDetダイバージェンスが近似誤差をどのように異なって罰するかを特徴付けている。ダイバージェンスを周辺スペクトル項とクロススペクトル結合項 $g_F(\lambda_i, \omega_j)$ に分解することで、著者らは以下を示す：

   • Frobeniusダイバージェンス： 結合項 $-\lambda_i \omega_j$ は双線形な大きさのペナルティを課し、上位（大きな固有値）のスペクトラムを強く強調する。
   • von Neumannダイバージェンス： 項 $-\lambda_i \log \omega_j$ は上位スペクトラムへの選好を保持するが、それを対数的に減衰させる。
   • LogDetダイバージェンス： 項 $\lambda_i / \omega_j$ は相対的な比率の歪みを生じさせ、近似を小さな固有値（テイルスペクトラム）に対して非常に敏感にする。

   定理2は、これらの最小化問題の停留点が、生の残差 $(S^* - C)$ ではなく、ダイバージェンスで重み付けされた残差 $\nabla^2 F(S^*)[S^* - C]$ によって決定されることを確立している。その結果、得られるクロネッカー因子は、選択されたダイバージェンスの特定のスペクトル選好（例：LogDet因子は小さな固有値に敏感であり、Frobenius因子は支配的な部分空間に焦点を当てる）に沿ったものとなる。

2. 共分散プリコンディショニングのスペクトル信頼性：
   著者らは、SST-2で訓練されたTransformerモデルを用いて、勾配共分散とヘシアン固有空間の間の整列（アライメント）を経験的に調査している。結果（表3）は、上位の共分散固有空間は、下位の固有値成分よりもヘシアンと実質的に良く整列していることを示している。テイルスペクトラムは、プリコンディショニングにおいてノイズが多く信頼性が低いことが観察されており、これは先行研究（Kunstner et al., 2019; Thomas et al., 2020）とも一致している。

3. 提案される解決策：部分空間認識型クロネッカー・オプティマイザ（BregTop）：
   ダイバージェンスによって誘導されるスペクトル選好と、実際のスペクトラムの信頼性との間のミスマッチに着想を得て、著者らは BregTop を提案している。このオプティマイザは、両方の観察結果の強みを組み合わせている：

   • 上位に敏感なダイバージェンス（例：von NeumannまたはFrobenius）を使用してクロネッカー因子を選択する。
   • 固有値ベースのプリコンディショニングを、上位の部分空間（結合クロネッカースペクトラムの最大 $\lceil \rho mn \rceil$ 個のエントリ）内でのみ適用する。
   • 相補的な下位部分空間に対しては、誤解を招く可能性のある曲率情報でプリコンディショニングを試みるのではなく、適応的な等方性加速定数 ($\chi_{\rho, q}$) を適用する。

   理論的には、これはプリコンディショナを、曲率でスケールされた上位成分と、下位空間のための一様な加速成分へと分解する（定理4）。

主な貢献

• スペクトル幾何学の特性付け： 本論文は、Frobenius、von Neumann、およびLogDetダイバージェンスがどのように固有値依存のペナルティを誘発するかについての形式的な特性付けを提供し、停留クロネッカー因子が生の誤差ではなく、ダイバージェンスで重み付けされた残差によって支配されることを示している。
• 経験的なスペクトル信頼性： 全体の共分散スペクトラムはプリコンディショニングには信頼できない一方で、上位スペクトラムはヘシアンと強い整列を維持していることを示し、部分空間認識型のアプローチを正当化している。
• 部分空間認識型オプティマイザの設計： 著者らは、信頼できる上位スペクトラムにおける固有値ベースのスケーリングと、ノイズの多い下位スペクトラムにおける等方性加速を適応的に切り替えるアルゴリズムであるBregTopを導入している。

実験結果
著者らは、4層のTransformerをSST-2データセット全体で訓練し、BregTopを標準的なShampoo、KL-Shampoo、VN-Shampoo、およびF-Shampooと比較評価した。使用された指標は、EMA平滑化された訓練損失が0.05を下回るまでに必要なステップ数である。

• BregTop-VN-v2 は、標準的なShampoo（1119ステップ）およびベースラインのVN-Shampoo（1127ステップ）を上回り、目標損失に 1003ステップ で到達した。
• BregTopの各バリアントは、対応するダイバージェンスのベースライン（例：BregTop-VNはVN-Shampooを、BregTop-FはF-Shampooを）を一貫して改善しており、上位部分空間をプリコンディショニングのために隔離することの有用性を検証している。

意義および主張
本論文は、構造化されたオプティマイザの設計は、クロネッカー分解における避けられない近似誤差と、共分散スペクトラムの信頼性の変動を考慮すべきであると主張している。異なるBregmanダイバージェンスが本質的に異なるスペクトル領域を優先することを認識し、かつ上位スペクトラムのみがヘシアンと信頼性高く整列していることを認識することで、提案された部分空間認識型のアプローチは、より堅牢なプリコンディショニング戦略を提供する。著者らは、この設計がSong et al. (2025) で観察された学習ダイナミクス（下位部分空間で学習が進むこと）と一致していると述べている。したがって、そこで等方性加速を適用することは、上位部分空間のプリコンディショニングによって高曲率方向へのノイズの多い更新を抑制しつつ、必要な勾配フローを維持することにつながる。結果は、このハイブリッドなアプローチが、完全な二次の情報を必要とせずにステップ効率を向上させることを示唆している。