技術的サマリー：閉形式による最終層最適化

問題定義

深層ニューラルネットワークは通常、すべての層を均一に扱う確率的勾配降下法（SGD）またはその変種を用いて最適化されます。しかし、二乗誤差関数の下では、バックボーンによって生成された特徴量が与えられた場合、ニューラルネットワークの線形最終層の最適重みは閉形式解（具体的にはリッジ回帰解）として得られます。

標準的な SGD は、最終層を勾配ステップで更新されるパラメータの集合として扱い、この解析的な最適性を無視しています。これによりミスマッチが生じます。バックボーンのパラメータ（$\theta$）は反復的な最適化を必要とする一方で、最終層の重み（$W$）は理論的には直接計算可能だからです。さらに、ミニバッチを用いた確率的設定において、各バッチに対して閉形式解を単純に計算すると、その特定のバッチへの過剰適合が深刻化し、不安定さを引き起こします。これを活用しようとする既存のアプローチ（因果推論などにおけるもの）は、しばしば閉形式の逆行列を通じた逆伝播を必要とし、計算コストが高くなるか、あるいは最終層をバックボーンから分離することで最適化の不安定さを招いています。

手法

著者らは、最終層をバックボーンパラメータの暗黙的かつ決定論的な関数として扱い、最終層が各ステップで最適であることを保証しつつ、バックボーンのみを最適化する訓練手順を提案します。

1. 理論的基盤（決定論的ケース）

固定されたバックボーンパラメータ $\theta$ に対して、最適最終層 $W^\star(\theta)$ は正則化された二乗誤差を最小化します：
$$W^\star(\theta) = \arg\min_W \sum \|y_i - W\phi_\theta(x_i)\|^2 + \beta\|W\|_F^2$$
著者らはエンベロープ定理を活用し、合成損失 $L^\star(\theta) = L(W^\star(\theta), \theta)$ の $\theta$ に関する勾配が、$W$ を $W^\star(\theta)$ に固定した元の損失の勾配と同一であることを示します：
$$\nabla_\theta L^\star(\theta) = \nabla_\theta L(W^\star(\theta), \theta)$$
この重要な結果は、$W^\star(\theta)$ を計算するために必要な行列逆演算を通じた逆伝播を行わなくてもよいことを意味します。代わりに、アルゴリズムは以下の 2 つを交互に行います：

1. $W_{t+1} = W^\star(\theta_t)$ の計算（閉形式更新）。
2. $\theta_{t+1} = \theta_t - \alpha \nabla_\theta L(W_{t+1}, \theta_t)$ の更新（バックボーンを通じた標準的な逆伝播）。

2. 実用的適応（確率的ミニバッチ）

確率的設定では、単一のミニバッチ上で $W^\star$ を計算すると過剰適合を引き起こします。これを解決するため、著者らは現在の閉形式解と以前の推定値 $W_t$ との距離を罰する近傍正則化子を導入します。目的関数は以下のようになります：
$$L_{B_t, W_t}^{\text{prox}}(W, \theta) = \sum_{(x,y) \in B_t} \|y - W\phi_\theta(x)\|^2 + \lambda \|W - W_t\|_F^2$$
この近傍損失の閉形式解は以下の通りです：
$$W^\star_{B_t, W_t}(\theta) = (Y_t \phi_\theta(X_t)^\top + \lambda W_t)(\phi_\theta(X_t)\phi_\theta(X_t)^\top + \lambda I)^{-1}$$
この定式化は、最終層上で近似カルマンフィルタを実装するものであり、前のバッチからの情報を蓄積しつつ現在のバッチに適応します。アルゴリズムは SGD によってバックボーンを更新し、その後、この近傍閉形式解を用いて最終層を再適合させます。

3. 分類への拡張

この手法は主に回帰のために設計されていますが、ワンホットラベルを二乗誤差のターゲットとして扱うことで分類に拡張されます。出力は $\arg\max$ 操作によって解釈されます。著者らは、このヒューリスティックが小規模から中規模のクラス数では機能するものの、確率の正規化を保証するものではないと指摘しています。

主要な貢献

1. 最適化フレームワーク：逆行列演算を通じた逆伝播を必要とせず、エンベロープ定理を活用して、標準的な SGD パイプラインに閉形式による最終層更新を統合する実用的なアルゴリズム。
2. 確率的正則化：最終層を以前の推定値に向けて正則化することでミニバッチ設定における過剰適合を防ぐ近傍定式化。これは近似カルマンフィルタリングとして解釈されます。
3. 理論的解析：
   • NTK 領域：関数空間における損失の非凸性にもかかわらず、ニューラルタンジェントカーネル（NTK）無限幅極限において大域的最適解への収束が証明されました。
   • ワンステップ解析：勾配降下ステップに対する閉形式ステップの優位性を定量化する理論的 bound が提示され、その利得は特徴量行列の条件数（特異値の分散）に関連することが示されました。
4. 実証的検証：多様なタスクにおいて、標準的な SGD および単純な閉形式リッジ回帰に対する改善が実証されました。

実験結果

本手法（"$\ell_2$ c.f. proximal ($\lambda$)"）は、標準的な二乗誤差 SGD/Adam および単純な閉形式リッジベースライン（"$\ell_2$ c.f. ridge ($\beta$)"）に対して評価されました。

• 量子化学（QM9）：近傍法は標準的な SGD を一貫して上回りました。特に、単純なリッジベースラインが収束に失敗する小バッチサイズにおいて顕著でした。また、バックボーンオプティマイザ（SGD 対 Adam）の選択に対する頑健性を示しました。
• フーリエニューラルオペレーター（FNO）：粘性バーガース方程式のオペレータ学習において、近傍法はバッチサイズにかかわらず有効であり続けましたが、単純なリッジベースラインは小バッチで失敗しました。この手法は高次元の空間自由度を効果的に処理しました。
• 深層特徴量道具変数（DFIV）：因果推論において、この手法は訓練後に全データセット上で最終層を再推定する高コストな必要性を排除し、計算効率を維持しながら再推定戦略と同等の性能を達成しました。
• 分類（CIFAR-100、ImageNet）：CIFAR-100 において、小規模から中規模のバッチサイズではクロスエントロピーと同等かそれ以上の精度を達成しました。しかし、クラス数が多数の ImageNet においてはクロスエントロピーを下回り、高カテゴリー分類における二乗誤差の限界を浮き彫りにしました。

意義と主張

本論文は、最終層の閉形式解を活用することが、最適化の加速と安定性の向上、特に回帰タスクや小バッチ領域において、実用的かつ効果的な戦略であると主張しています。

• 理論的意義：この研究は、関数空間における修正された損失の臨界点に関する最初の解析を提供し、問題に内在する非凸性を克服して NTK 領域における大域的最適解への収束を証明しました。
• 実用的意義：近傍定式化は、ミニバッチ閉形式更新に内在する過剰適合の問題を解決し、手法を頑健にして標準的な訓練パイプラインとの互換性を可能にします。
• 限界：著者らは明示的に、この手法が現在二乗誤差に限定されていると述べています。少数のクラスを持つ分類では良好に機能しますが、大規模分類（ImageNet）ではクロスエントロピーを下回るため、クロスエントロピー損失への適応が今後の研究に必要な方向性であることを示唆しています。また、この手法は最終層の次元にスケーリングする行列逆演算による計算オーバーヘッドを伴い、小バッチ領域では非常に大きな最終層を持つ場合、標準的な SGD よりも遅くなります。ただし、このコストはバックボーンのサイズが増大するにつれて減少します。