技術要約：To Grok Grokking: Ridge Regressionにおける証明可能なGrokking

問題提起

本論文は、モデルが完璧な訓練性能（過学習）を達成した後、かなり時間が経過してから汎化が始まる現象である**grokking（グロッキング）**について調査している。grokkingはこれまで、主に複雑な非線形ニューラルネットワークやアルゴリズム的なデータセットにおいて観察されてきたが、著者らは、この現象に対して厳密かつエンドツーエンドの理論的保証を確立することを目指している。具体的には、既存の文献では「過学習」と「汎化」の間の「遅延」を形式的に証明できていないか、あるいは特定の（モジュロ加算のような）一般化しにくい問題構造に依存しているというギャップに対処している。核心となる問いは、過剰パラメータ化されたリッジ回帰において、重み減衰（weight decay）を伴う勾配降下法（GD）を用いて、grokkingを定量的に証明できるか、そしてハイパーパラメータがいかにして「grokking time（grokkingが発生するまでの時間）」を制御するかである。

手法

著者らは、問題を**教師・生徒フレームワーク（teacher-student framework）**内で定式化している：

• 設定: 過剰パラメータ化された線形回帰において、生徒モデル $N(x; \theta) = \langle \theta, \phi(x) \rangle$ が、実現可能な教師関数 $N^*(x) = \langle \theta^*, \phi(x) \rangle$ を学習する。
• 最適化: 生徒は、固定ステップサイズ $\eta$ と $\ell_2$ 正則化項（重み減衰） $\lambda$ を用いた**勾配降下法（GD）**を用いて訓練される。目的関数は、正則化された平均二乗誤差である：
  $$L_n(\theta; \lambda) = \frac{1}{2n}\sum_{i=1}^n (N(x_i; \theta) - N^*(x_i))^2 + \frac{\lambda}{2}\|\theta\|_2^2$$
• 分析戦略: 著者らは、パラメータベクトル $\theta$ を、データのスパン（span）に対する2つの直交する成分に分解する：
  1. $\theta_{\parallel}$: 特徴行列 $\Phi$ の行空間（データが占める部分空間）内の成分。
  2. $\theta_{\perp}$: $\Phi$ の零空間（補空間）内の成分。
• 理論的アプローチ: 証明は、訓練損失（経験誤差）と汎化損失（母集団誤差）の収束率を個別に分析することに基づいている。核心となるメカニズムは、$\theta_{\parallel}$ が訓練データを適合させるために急速に収束する一方で、$\theta_{\perp}$ は重み減衰によって緩やかに減衰し、その結果として汎化性能が低い状態が長く続くことを示すことにある。

主な貢献

1. 初のエンドツーエンドで証明可能なGrokking: 本論文は、線形設定におけるgrokkingの3つの明確な段階を示す、厳密かつ定量的な境界（bounds）を提供する最初の研究である：
   • (i) 訓練データへの急速な過学習。
   • (ii) 訓練誤差がゼロであるにもかかわらず、汎化誤差が高いまま維持される持続的な遅延。
   • (iii) 最終的に、極めて小さな汎化誤差を持つモデルへの収束。
2. Grokking Timeに関する定量的境界: 著者らは、「grokking time」（過学習までの時間 $t_1$ と汎化までの時間 $t_2$ の差）の明示的な下限値を導出した。これらの境界は、以下のハイパーパラメータによって表現される：
   • 重み減衰 ($\lambda$): grokking timeは $\lambda$ に対して反比例する（$t_2 \propto 1/\lambda$）。$\lambda$ が小さいほど、汎化が大幅に遅れる。
   • 初期化スケール ($\nu^2$): 初期化スケールを大きくすると、grokkingは対数的に増幅される（$t_1, t_2 \propto \ln(\nu^2)$）。
   • 過剰パラメータ化: この現象は、高次元の設定（$m \gg n$）を必要とすることが示されている。
3. ハイパーパラメータによる制御: 著者らは、grokkingを誘発または排除するための、ハイパーパラメータ（サンプルサイズ $n$、特徴量次元 $m$、重み減衰 $\lambda$、および初期化スケール $\nu^2$）に関する十分条件を確立した。
4. 非線形モデルにおける実証的検証: 理論は線形回帰に基づいて導出されているが、著者らは、ハイパーパラメータの定量的な依存関係（例：$\lambda$ が遅延に与える影響）が、非線形な2層ReLUニューラルネットワークで見られる挙動と定性的に一致することを実証的に示している。

主要な結果

• 収束率: 訓練誤差は、経験的特徴共分散の最小の非ゼロ固有値（$\lambda^+_{\min}(\Phi^\top \Phi)$）によって支配される速度で収束し、これは高速である。対照的に、汎化誤差は $(1 - \eta\lambda)^t$ の速度で減衰する項によって下限付けられ、$\lambda$ が小さい場合には極めて緩慢になる。
• Grokkingのメカニズム: 遅延は、小さな重み減衰を伴うGDが、訓練誤差を最小化するためにデータのスパンへの重みの投影のみを効果的に更新する一方で、直交成分（汎化誤差に寄与するもの）は重み減衰の項によってのみ緩やかに減衰することによって生じる。
• 定理:
  • 定理 4.1: ゼロ教師（zero teacher）におけるgrokkingを確立し、損失収束の一致する上限および下限を示す。
  • 定理 4.2: 結果を任意の実現可能な教師へと一般化し、$t_1$ と $t_2$ の明示的な境界（式6および7）を提供する。
  • 定理 4.4 & 4.5: 訓練損失の高速な収束と、不十分な汎化の持続性を証明する。
  • 定理 4.6: GDが最終的に、良好な汎化保証を持つグローバルな最小値に収束することを証明する。

意義と主張

本論文は、grokkingを理解するための基礎的な理論的枠組みを提供すると主張している。線形回帰という単純かつ透明な設定に現象を分離することで、著者らは以下のことを述べている：

• Grokkingは、ディープラーニングのアーキテクチャや非線形性に固有の失敗モードではなく、むしろ特定の訓練条件（具体的には、過剰パラメータ化と小さな重み減衰の相互作用）の結果である。
• この現象は、モデルのアーキテクチャや学習アルゴリズムに根本的な変更を加えることなく回避可能であり、ハイパーパラメータのチューニングを通じて制御できる。
• 本結果は、従来の文献で見られるような（「lazy-to-rich」レジームの遷移といった）経験的な観察やヒューリスティックな説明を超えて、汎化の遅延に関する最初の厳密な定量的境界を提供するものである。
• これらの知見は、分類問題でよく見られる「テスト誤差のプラトー（停滞）」が、閾値ベースの精度指標を用いて評価した場合、回帰においても存在することを示唆しており、損失の種類を問わないこの現象の普遍性を裏付けている。

著者らは、自身の線形解析の範囲について謙虚な姿勢を保っており、線形解析は厳密な証明を提供するものの、非実現的な設定への拡張や、深層ニューラルネットワークにおける「lazy-to-rich」の遷移の証明には、新たな理論的ツールを必要とする未解決の課題が残っていると述べている。