Grokking は代数で決まるのか：厳密に解ける極限モデル

導入

Grokking は、ニューラルネットワークが最初に訓練データを記憶し、その後しばらくして突然一般化規則を獲得する現象として知られている。特にモジュラー算術タスクでは、モデル容量がこの現象に大きく影響することが示されてきた。容量が小さすぎると記憶すら難しく、容量が大きいと一般化が早く起きる。

本論文が問うのは、そのさらに極端な場合である。もしアーキテクチャが表現できる関数全体が、有限次元の代数的多様体のような固定された対象にまで縮んだら、grokking は残るのだろうか。

核心ポイント

• 解析可能な二層モデル：著者らは、活性化関数に全純単項式 σ(z)=z^k を用いる二層ネットワークを扱う。モジュラー算術タスクは 1 の冪根で符号化され、離散フーリエ文字を通じて出力を記述できる。
• 隠れ幅を増やしても制約は消えない：ネットワーク出力は、隠れ層の幅に関係なく ((\mathbb{Z}_p)^2) 上の文字の ((k+1)) 次元部分空間に閉じ込められる。これは全関数空間のうち (O(k/p^2)) 程度のごく小さな部分である。
• 表現可能性の条件が厳密：タスクが表現可能であるための必要十分条件は、その離散フーリエ台が (u+v=k \pmod p) という対角線上にあること。線形位相ターゲットでは、これは (m+n=k) という算術条件に簡約される。
• 失敗は最適化不足ではない：条件を満たさないターゲットは、汎化できないだけでなく、訓練データにも適合できない。論文は、隠れ幅に依存しない正の訓練損失下限も証明している。
• 実験結果は二分される：585 回の実行で、代数的予測と観測結果は 99.8% の精度で一致した。記憶段階から一般化へ移るのではなく、即座に成功するか完全に失敗するかに分かれる。

意義と影響

この研究は、grokking を容量と表現可能性の観点から極限まで単純化する。通常の議論では、いつ記憶から一般化へ移るかが焦点になる。しかしこのモデルでは、問題は時間的な遷移ではなく、ターゲット関数がそもそもネットワークの代数的関数クラスに含まれるかどうかになる。

もちろん、標準的なニューラルネットワークが常にこのように二値的に振る舞うわけではない。論文はボトルネック除去実験を通じて、この極端な設定と通常のネットワークを接続している。表現不能、一般化しない記憶、そして容量増加に伴ってギャップが縮む grokking へと連続的に移るという見方である。

その意味で本研究は、訓練ダイナミクスによる遅延一般化と、アーキテクチャが持つ表現能力の限界を切り分ける理論的なレンズを提供している。

出典：arXiv