技術要約：フィッシャー幾何学的な鋭さと、平坦な極小値に対するSGDの暗黙的バイアス

1. 問題提起

ディープラーニングにおける支配的な直感は、確率的勾配降下法（SGD）が「平坦な極小値（flat minima）」を暗黙的に優先し、それらの平坦な極小値は「鋭い（sharp）」ものよりも汎化性能が高いというものである。しかし、この物語の理論的基礎は、Dinhら [2017] によって批判的に覆された。彼らは、損失ヘッシアンのトレースや最大固有値といった標準的なユークリッド的な平坦さの尺度がいかに再パラメータ化に対して不変ではないかを実証した。

具体的には、ニューラルネットワークにおいて、重みをスケーリング（例：$W_1 \to \alpha W_1, W_2 \to \alpha^{-1} W_2$）することで、ネットワークの入出力関数を正確に維持したまま、ユークリッド的な鋭さを任意に変えることが可能である。汎化はパラメータ化ではなく関数の特性であるため、不変ではない鋭さの尺度は、汎化を説明する上で意味をなし得ない。不変な尺度に向けた部分的な歩み（例：Jang et al. [2022]; Kristiadi et al. [2023]）はあるものの、「不変な鋭さ」、「SGDの暗黙的バイアス」、および「汎化境界」を統合する枠組みは依然として欠如していた。

2. 手法

2.1 リーマン的鋭さ ($S_R$)

著者らは、パラメータ空間におけるユークリッド計量を、モデルの予測分布に固有のリーマン構造を誘導するフィッシャー情報行列（FIM） $G(\theta)$ に置き換える。彼らは、極小値 $\theta^*$ におけるリーマン的鋭さを次のように定義する：
$$S_R(\theta^*) = \text{tr}\left( G(\theta^*)^{-1} \nabla^2 L(\theta^*) \right)$$
ここで $\nabla^2 L$ は訓練損失のヘッシアンである。ユークリッド的な鋭さとは異なり、$S_R$ は統計的多様体の自然な幾何学に対する曲率を測定する。

2.2 理論的枠組み：確率微分方程式（SDE）としてのSGD

著者らは、ミニバッチSGDを連続時間SDEとして定式化する。彼らは、勾配ノイズの共分散がFIMに比例するという仮定3を導入する：
$$\text{Cov}[\nabla \tilde{L}(\theta^*)] = \frac{\eta}{B} G(\theta^*)$$
ここで $\eta$ は学習率、$B$ はバッチサイズである。この仮定の下で、彼らはSGD過程の定常分布 $\pi^*$ を導出する。

2.3 汎化境界

著者らは、幾何学的バイアスをPAC-Bayesフレームワークを用いてテスト性能へと結びつける。彼らは訓練された極小値を中心とする事後分布を構築し、$S_R$ によって明示的に制御される汎化境界を導出する。

2.4 実験的推定

フルFIMの計算は困難であるため、著者らは対角近似による経験的FIM推定器（モデル自身の分布から予測をサンプリングするもの）と、ヘッシアンのためのHutchinsonの確率的トレース推定器を利用する。彼らは、真のFIMは不変性を保証するものの、実験で使用される対角近似による経験的推定器は、その不変性を近似的にしか継承していないことを認めている。

3. 主な貢献

1. 再パラメータ化不変な鋭さの定義： 本論文は $S_R$ を定義し、それが滑らかで関数保存的な再パラメータ化に対して不変であることを証明した（補題1）。これにより、Dinh et al. [2017] による批判に直接答えている。
2. 定常分布の解析： 著者らはミニバッチSGDの定常分布を導出し（定理1）、その確率質量がリーマン的に平坦な極小値に指数関数的に集中することを証明した（系1）。これにより、リーマン的な平坦な極小値への厳密な暗黙的バイアスが確立された。
3. PAC-Bayes 汎化境界： 汎化ギャップが $S_R$ によって明示的に制御される汎化境界（系2）が確立された。これにより、リーマン的な平坦さとテスト性能との間の理論的な繋がりが提供された。
4. 実験的検証： MNISTおよびCIFAR-10を用いた実験により、$S_R$ が、ユークリッド的な鋭さ ($S_E$) とは異なり、異なるオプティマイザ（SGD, Adam）、バッチサイズ、学習率を横断して汎化性能を追跡することを確認した。$S_R$ の $\eta/B$ に対するスケーリングは、理論的予測と一致している。

4. 結果

4.1 不変性と整合性

• 不変性： 真のFIMは厳密に不変であるが、経験的な対角近似推定器は任意のスケーリングに対して完全には不変ではない。しかし、それはユークリッド的な鋭さに比べれば、重みのスケーリングに対して大幅に安定している。
• ノイズ構造： 実験的テスト（セクション5）によれば、勾配ノイズの共分散とFIMは、様々な構成において方向的に整合しており（余弦類似度 $\approx 1.0$）、両者は一致している。ただし、スカラー倍の比例定数は、特に小バッチ・高ノイズの設定においては、必ずしも $\eta/B$ とは一致しない。著者らは、このスカラーの主張を修正し、スカラー因子 $\alpha(\theta)$ を許容するように仮定を洗練させた。これにより、主要な理論的結論（平坦な極小値へのバイアス）は、スカラーの主張が不完全であっても堅牢であることが証明されている。

4.2 スケーリング則と汎化

• バッチサイズ ($B$)： MNISTにおいて、 $B$ を大きくすると $S_R$ が大幅に増加し、テスト精度が低下する。これは $S_R \propto B/\eta$ という予測と一致している。CIFAR-10では、小バッチ領域での最適化の困難さにより関係性はより複雑になるが、高い $B$ が高い $S_R$ につながるという傾向は保持されている。
• 学習率 ($\eta$)： $\eta$ を大きくすると、一般に $S_R$ は減少し、汎化性能が向上する（訓練が不安定になる点までは）。
• オプティマイザ：
  • SGD： 小バッチ設定において、一貫して低い $S_R$ と高い精度を持つ極小値を見出す。
  • Adam： MNISTにおいて、低い $S_R$ と高い精度を達成する。CIFAR-10では、AdamはSGDよりも低い $S_R$ を持つ極小値を見出すが、精度はわずかに低くなる。著者らは、これはAdamの適応的プリコンディショニングが実効的な計量を変化させているためであると示唆している。すなわち、プリコンディショニングされたオプティマイザの暗黙的バイアスは、それ自身の計量に関する平坦さに向かっている。

4.3 学習ダイナミクス

学習の軌跡は、小バッチSGDが訓練の進行とともに、段階的に鋭い盆地（basin）を脱出し、より平坦な領域（減少する $S_R$）へと落ち着いていく様子を示している。大バッチSGDは、高い $S_R$ 値で安定する。これは、鋭い極小値を脱出するためのノイズが不足していることと整合している。

5. 意義と主張

本論文は、なぜ平坦な極小値が汎化するのかについて、初の厳密で再パラメータ化不変な説明を提供することを主張している。統計的多様体のリーマン幾何学（FIM経由）に鋭さの根拠を置くことで、著者らは「平坦さは単なるパラメータ化のアーティファクトである」という根本的な異議を解決している。

本研究の意義は以下の通りである：

• 理論的解決： 「平坦さと汎化」という直感と「再パラメータ化の批判」との間の溝を、モデルの関数に固有の鋭さの尺度を定義することで埋めた。
• 暗黙的バイアスのメカニズム： SGDのノイズ構造（FIMに比例する）が、自然にリーマン的に平坦な極小値へとオプティマイザを導くことを数学的に導出した。
• 実用的な指標： $S_R$ は、特に異なるハイパーパラメータやオプティマイザを用いて訓練されたモデルを比較する場合において、ユークリッド的な鋭さよりも信頼できる予測因子であることを示した。

著者らが認める限界事項：

• 経験的な対角FIM推定器は、完全な再パラメータ化不変性を持っていない（K-FACのような構造化された推定器であれば完全な不変性が得られる）。
• ノイズ共分散のスカラーの主張（$\alpha = \eta/B$）は近似であり、方向性は成立するものの、データセットやアーキテクチャに応じて大きさは変化する。
• $S_R$ と汎化の関係は、オプティマイザが収束している場合に最も情報量が多い（非常に高い学習率などの不安定な領域では、低い $S_R$ は良好な汎化を保証しない）。