技術要約：Grokking（グロッキング）を制御する：表現幾何学を制御信号として用いる方法

問題提起

Grokking（グロッキング）は、過剰パラメータ化されたニューラルネットワークにおいて、モデルが初期段階では訓練データを記憶（memorize）し、長時間の最適化を経て初めて強力な汎化を示すようになる現象である。近年の研究では、グロッキングとウェイト減衰、特徴量の圧縮、回路形成、および表現学習のダイナミクスとの関連性が示されているが、遅延した汎化が発生するタイミングや出現を支配する具体的な要因については、依然として完全には解明されていない。既存の文献における共通のテーマは、学習された表現が低次元構造へと漸進的に圧縮されることである。しかし、先行研究の多くは、表現の幾何学的性質を直接的な制御メカニズムとしてではなく、創発的な特性として扱っている。本論文では、隠れ層の表現の幾何学を明示的に操作することが、グロッキングの発生時期と速度に影響を与える制御信号として機能するかどうかを調査する。

手法

著者らは、訓練中の隠れ層の表現の有効次元数を直接制御するために設計された、表現レベルのスペクトル正則化手法であるGeometric Dimensionality Regularization (GeomDR) を導入する。

幾何学的特性

隠れ層の表現行列 $Z \in \mathbb{R}^{N \times d}$ に対し、手法はまずデータを中心化して $\tilde{Z}$ を得る。次に、経験共分散行列 $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$ を計算する。有効次元数 ($D_{eff}$) は、$C$ の固有値 ($\mu_i$) の参加比（participation ratio）を用いて以下のように算出される：
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
さらに、局所的な近傍距離 ($\rho$) は、正規化された表現の $k$ 近傍点への平均ユークリッド距離として定義される。

GeomDR の目的関数

GeomDR は、ターゲットとなる部分空間の次元 $d^*$ を超える方向に含まれる分散に対してペナルティを課す。降順にソートされた固有値を $\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$ とすると、層ごとの損失は以下の通りとなる：
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
全訓練目的関数は、タスク損失 ($L_{task}$) と幾何学的正則化項を組み合わせたものである：
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
極めて重要な点は、この正則化項は最初からアクティブではないことである。これは、初期の制約なしの訓練フェーズ（活性化ステップ $t_s$）の後に導入され、その強度 $\lambda(t)$ はコサインランプを用いて徐々に増加させる。この設計は、記憶が汎化に先行するという自然なグロッキングの軌跡を模倣している。

実験プロトコル

本フレームワークは、グロッキングを示すことが知られている3つのアルゴリズムタスクで評価された：

1. Modular Addition (加算): $(a + b) \mod 97$ を予測する。
2. Modular Division (除算): $(a \cdot b^{-1}) \mod 97$ を予測する。
3. Permutation Composition (置換合成): $S_5$ における2つの置換の合成を予測する。

実験には、MLP（多層パーセプトロン）とTransformerの2つのアーキテクチャを用いた。本研究では、介入スケジュール（強度 $\lambda_{max}$ およびランプ期間 $T_{ramp}$）、ターゲット次元数 ($d^*$)、および介入開始時刻 ($t_s$) を系統的に変化させた。

主な貢献

1. 先行する崩壊の観察: 著者らは、ベースラインの設定において、次元崩壊（dimensionality collapse）が記憶から汎化への転移に一貫して先行して起こることを示した。これは、幾何学的圧縮が単なる相関関係ではなく、転移のメカニズム的な駆動要因であることを示唆している。
2. GeomDR の導入: ターゲットとなる部分空間の外側にある共分散スペクトルを直接抑制することで、隠れ層の有効次元数を制御するスペクトル正則化手法である GeomDR を提案する。
3. ダイナミクスの系統的制御: 幾何学的な介入によって、遅延した汎化のダイナミクスを大幅に変化させられることを示す包括的な実証研究を提供する。介入スケジュールとターゲット次元数に応じて、GeomDR はグロッキングを加速させ、遅延させ、あるいは完全に抑制することができる。

結果

グロッキングの加速

評価されたすべての設定において、GeomDR は一貫してグロッキングのダイナミクスを変化させた：

• MLP の性能: Modular Addition において、GeomDR はグロッキングを最大 52倍 加速させた（ステップ数を約362Kから約7Kへ短縮）。Modular Division（約39倍）や Permutation Composition（約23倍）でも同様の改善が見られた。すべての MLP ケースにおいて、ベースラインが不安定または失敗した場面でも、GeomDR は100% の成功率を達成した。
• Transformer の性能: Transformer のベースラインのグロッキング・ダイナミクスは既に MLP よりも高速であったが、それでも GeomDR は一貫した加速（例：Modular Addition で約2倍）を提供し、100% の成功率を維持した。

ハイパーパラメータの影響

• ターゲット次元数 ($d^*$): 中程度から大きなターゲット次元 ($16 \le d^* \le 64$) が最も強力な加速をもたらした。過度な圧縮 ($d^* \le 8$) は、学習を遅延させたり不安定にしたりすることが多かった。
• 介入タイミング ($t_s$): 「決定的な時間窓」が特定された。有用な表現が形成される前（早すぎる）や、記憶が固定化された後（遅すぎる）に行われる介入は効果が低かった。中間的な活性化時間（例：MLP の場合 $t_s = 2000$）が、最も高速で安定したグロッキングを生み出した。
• スケジュール: 中程度の正則化強度とランプ期間が概して最も良好な結果を示した。強すぎる介入は訓練を不安定にする可能性がある。

幾何学的相関

本研究は、グロッキングを加速させる介入が、次元崩壊の発生も加速させることを確認した。GeomDR モデルでは、次元崩壊と成功した汎化との間の時間的ギャップがベースラインと比較して大幅に減少しており、これは幾何学的再編成と汎化の出現との間の緊密な結合を示唆している。

スケーラビリティと効率性

著者らは、共分散スペクトルが表現のランダムなサブセット（例：5%～25% のデータ）から推定される場合でも、GeomDR が効果的に機能することを示し、グロッキングの加速に実質的な影響を与えることなく、実行時間を 2.4倍～2.7倍高速化した。

意義と主張

本論文は、表現の幾何学がグロッキングに対する効果的な制御信号として機能することを主張している。隠れ層の活性化の共分散スペクトルを直接修正することで、研究者はネットワークがいつ、どのように迅速に記憶から汎化へと移行するかを系統的に制御できる。

著者らは、本研究が、幾何学的介入が遅延した汎化を研究するための実用的なアプローチを提供することを証明していると述べている。表現の変化を受動的に観察するのではなく、この手法を用いることで、学習ダイナミクスに関する仮説を検証するために幾何学的特性を直接操作することが可能となる。結果は、表現の幾何学とグロッキングとの関係が異なるタスクやアーキテクチャ間でロバストであることを示唆しているが、その改善の大きさは変動する（Transformer よりも MLP において顕著であり、これはおそらく Transformer の方が強力な帰納バイアスを持っているためである）。

結論として、本論文は、表現の次元数を唯一の因果的メカニズムとして確立したものではないが、幾何学的介入がグロッキングのタイミングを変えるのに十分であることを示し、最適化および汎化研究のための新しい実験的手法を提示している。