Fast Generalization after Interpolation via Critically Damped Momentum Optimization

技術要約：補間後の高速な汎化を実現する臨界減衰モーメンタム最適化

問題提起
高次元かつ低サンプル数（ゲノミクス、医学、金融などで一般的）のレジームにおいて、機械学習モデルは訓練性能においては完璧に近い補間（interpolation）を達成する一方で、未知のデータに対する汎化に失敗することがよくあります。この乖離は、訓練損失は同一であるが、その複雑さとテスト性能が大きく異なる多くのパラメータ構成が、損失関数内に存在するために発生します。ウェイトディケイ（重み減衰）などの正則化手法は、最適化をより単純な解へとバイアスさせるために伝統的に用いられてきましたが、「グロッキング（grokking）」という現象は、時間的な断絶を明らかにしています。すなわち、モデルはまずデータを記憶（memorize）し、その後かなり後になって初めて汎化へと移行するのです。この遅延した遷移は計算効率が悪く、汎化のために補間閾値に達するよりも数桁多くのエポックを必要とします。核心となる問題は、受動的なダイナミクスに頼るのではなく、モデルが補間閾値に達した直後に、いかにして効率的に低複雑度の補間解を選択するかという点にあります。

手法
著者らは、適合フェーズと汎化フェーズを明示的に分離するように設計された二相式最適化戦略であるGROKtimizerを提案しています。

1. 理論的枠組み:

   • 補間後のダイナミクス: 本論文では、補間後のレジームを局所的な二次力学系としてモデル化しています。補間解付近の損失関数が、ヘッセ行列（$H$）と$L_2$正則化項（$\lambda$）を含む二次形式で近似できるという仮定の下で、平坦な方向（$H$の零空間）に沿ったダイナミクスは、減衰振動子と同様の挙動を示します。
   • 臨界減衰（Critical Damping）: 著者らは、**臨界減衰モーメンタム（CDM）**が、この平坦な盆地内においてパラメータノルムを最小化するための最適な軌跡を提供することを証明しています。モーメンタム・パラメータ $\beta$ を $\beta = 1 - 2\sqrt{\lambda\eta}$ （ここで $\eta$ は学習率）に調整することで、オプティマイザは臨界減衰を実現します。
   • 収束の加速: 理論上、CDMはこのレジームにおいて古典的な勾配降下法（GD）に対して二次的なスピードアップを提供します。GDの収束タイムスケールが $\lambda^{-1}$ に比例するのに対し、CDMは $\lambda^{-1/2}$ に比例するタイムスケールで収束します。これにより、最小ノルム解に到達するために必要なエポック数の二次的な削減が証明可能です。
   • 学習率の推定: ヘッセ行列のフル行列を計算することなく実用的にCDMを実装するために、著者らはAitken補正を用いたパワーイテレーション法を用いて、ヘッセ行列の最大固有値（$\lambda_{max}$）を推定する方法を提案しています。これにより、安定性の境界における最適な学習率 $\eta_{max}$ をオンラインで算出することが可能になります。

2. アルゴリズムの実装 (GROKtimizer):

   • フェーズ1（補間）: 最適化は、補間閾値への迅速な収束を優先するため、ウェイトディケイなしで実行されます。このフェーズでは、訓練データを素早く適合させるために標準的なモーメンタムが利用されます。
   • フェーズ2（汎化）: 補間に達した時点で、オプティマイザはウェイトディケイを伴う臨界減衰モーメンタムへと切り替わります。モーメンタム・パラメータは、理論から導出された臨界減衰条件に合わせて調整され、ウェイトディケイが適用されることで、パラメータを補間多様体内の最小ノルム解へと駆動します。

主な貢献

1. 補間後レジームの特性付け: 著者らは、遅延した汎化フェーズを、損失関数の平坦な方向における減衰振動子の問題として形式的に定義し、グロッキング現象をヘッセ行列の零空間の幾何学と結びつけました。
2. 理論的なスピードアップ: 臨界減衰モーメンタムが、標準的な勾配降下法よりも二次的に速く低ノルム補間解への収束を加速することを証明し、局所二次モデルの下での一次近似手法における最適性を確立しました。
3. GROKtimizer: 非正則化された迅速な適合から、臨界減衰によるノルム最小化へと動的に切り替わる、二相式の最適化スケジュールを導入しました。
4. 実証的検証: GROKtimizerは、合成グロッキング・ベンチマークおよび実世界の高次元データセットにおいて検証されており、標準的なオプティマイザ（Adam, AdamW, SGD, Muon）に対して一貫した改善を示しています。

実験結果
本論文では、以下のベンチマークでGROKtimizerを評価しています。

• 合成ベンチマーク: Modular Addition、Binary Addition、Sparse Parityなどのタスクにおいて、GROKtimizerはベースラインを大幅に上回りました。例えば、Gaussianタスクでは、最強のベースラインと比較して検証損失を数桁減少させました。
• 実世界の高次元データ: 生物医学的分類タスク（Leukemia, TCGA）および分子特性予測（QM9）において、GROKtimizerは最高の精度と最低の検証損失を達成しました。これらの利得は、特徴量がサンプル数を大幅に上回るデータ不足の状況において最も顕著でした。
• サンプル効率: 低データMNIST実験（クラスあたり10〜50サンプル）において、GROKtimizerは最大の優位性を示し、明示的なノルム最小化が、解空間が低決定（underdetermined）である場合にサンプル効率を向上させることを裏付けました。
• 言語モデルの事前学習: WikiText2およびBabyLMモデルを用いた実験では、スケーリングの対称性や適応的な更新により、パラメータノルムと検証損失の直接的な相関は（トランスフォーマーにおいては）それほど明確ではないものの、二相式のパターンは依然として成立していました。補間後にノルム最小化フェーズに切り替えることで、検証損失のさらなる減少が得られました。

意義と主張
本論文は、事後的な補間訓練を受動的な待機期間としてではなく、明示的な最適化問題として扱うことで、GROKtimizerが「グロッキング」問題に対する原理的な解決策を提供すると主張しています。損失関数の幾何学を利用することで、この手法は低複雑度の解の選択を加速します。

著者らは、自身の知見を**平坦な極小値仮説（flat-minima hypothesis）**と整合させ、ヘッセ行列の零空間の次元が、平坦さと複雑さの堅牢な尺度として機能すると論じています。彼らは、この零空間内でのノルムを最小化することは、より平坦で汎化性能の高い極小値へとナビゲートすることと同等であると示唆しています。

限界
著者らは以下の限界を認めています。

• このアプローチは主にデータ制約のあるレジーム向けに設計されており、データの豊富さが性能を支配する大規模なビジョンや言語の事前学習においては、重要性が低い可能性があります。
• 本手法は、モデルが補間に達するのに十分な容量を持っていることを前提としており、過小パラメータ化されたモデルのための解決策ではありません。
• 理論的解析は損失関数の局所的な二次近似に依存していますが、実験結果は、この仮定からの逸脱に対しても本手法が頑健であることを示唆しています。
• パラメータノルムは複雑性のプロキシ（代理指標）として使用されていますが、あらゆるアーキテクチャにおけるモデルの汎化のあらゆる側面を捉えきれているわけではありません。