A Stochastic--Geometric Theory of Scaling Laws in Grokking

あなたは、ビデオゲームのキャラクターが難解なパズルを解こうとしている様子を見ていると想像してください。最初、キャラクターは先ほど通った経路をすべて記憶しようとして、狂ったように走り回っています。彼らは瞬時にスコアをゼロにしますが、ルールを実際に理解したわけではなく、単に特定のマップを暗記しただけです。これは**「暗記（memorization）」**と呼ばれます。

次に、魔法のようなことが起こります。何も変化していないかのように見える長く退屈な期間を経て、キャラクターは突然動きを止め、盤面を見つめ、そして「理解」します。彼らは見たことのない新しいパズルをも完璧に解き始めるのです。この突然の「アハ体験（ひらめき）」こそが、研究者たちが**「グロッキング（grokking）」**と呼ぶ現象です。

長い間、なぜこの遅延が発生するのかは誰にも分かりませんでした。それは魔法だったのでしょうか？ それともグリッチ（不具合）だったのでしょうか？ 本論文において、著者らは、形状と距離のマップを用いて、コンピュータの「脳（ニューラルネットワーク）」の中で何が起きているのかを可視化する新しい方法を提案しています。

解のオニオン（玉ねぎの構造）

著者らは、コンピュータの「脳」が存在できる空間が、まるでオニオンやターゲット（的）のように、入れ子状になった球体の集合体であると提案しています。

1. 外殻（初期化）： コンピュータが開始されるとき、その設定は薄い外側の殻に散らばっています。それは、巨大な風船の非常に端の方にダーツを投げ、すべてが端に着地するようなものです。
2. 中間層（暗記）： 学習が進むにつれ、コンピュータは素早く中間層へと滑り落ちます。ここでは、訓練データを完璧に暗記しています。それは、ルールを理解せずにマップを暗記したばかりの、私たちのゲームのキャラクターのような状態です。コンピュータはここで長い間、足止めを食らいます。
3. 核（汎化）： オニオンの深い中心には、「汎化（generalization）」の核が存在します。ここは、コンピュータが実際にルールを理解し、新しい問題を解決できる「スイートスポット」です。

「グロッキング」という現象は、単に中間層から空隙を経て、核へと向かう旅なのです。謎は、なぜそこへ到達するのにこれほど時間がかかるのか？ という点でした。

酔っ払いの歩行と磁力による引き寄せ

著者らは、コンピュータの学習プロセス（Adamと呼ばれる最適化手法を使用）が、2つの力の混合であることを説明しています。

• 酔っ払いの歩行（拡散 / Diffusion）： コンピュータは小さなデータのバッチから学習するため、その経路は、直線を歩こうとする酔っ払いの足取りのように、少しふらついています。このふらつきはランダムなものです。
• 磁力による引き寄せ（ドリフト / Drift）： また、設定値を小さく保とうとするジェントルな磁石のように機能する、$\ell_2$ 正則化というルールによって、コンピュータを中央へと引き寄せる一定の力が働いています。

論文は、コンピュータが中間層で足止めを食らう理由について、「酔っ払いの歩行」が中間層から脱出するには弱すぎ、「磁力による引き寄せ」が直接中央へと引きずるほど強くないためであると示唆しています。コンピュータは、ランダムなふらつきによって、たまたま十分に遠くまで押し出され、核へと落ち込むまで、しばらくの間、中間層をさまよい続ける必要があるのです。

ゲームのルール（スケーリング則）

著者らは単に推測したのではなく、数学（具体的には停止時間理論と確率微分方程式）を用いて、この遅延がどのくらい続くかを正確に予測しました。彼らは、この旅の速度を制御する3つの主要な「つまみ（パラメータ）」を見つけました。

1. 学習率 ($\eta$): これはコンピュータが踏み出すステップの大きさです。論文では、ステップが小さすぎると中間層からの脱出に永遠に時間がかかり、大きすぎるとターゲットを通り過ぎてしまうことが示されています。そこには「ゴルディロックス（ちょうど良い）」ゾーンが存在します。
2. バッチサイズ ($b$): これはコンピュータが1ステップ進む前に見る例題の数です。論文は、バッチサイズが大きくなると「酔っ払いの歩行」のふらつきが減り、それが結果として中間層からの脱出を遅らせることを示唆しています。
3. 正則化 ($\lambda$): これは「磁力による引き寄せ」の強さです。論文は、より強い引き寄せがコンピュータの中間層からの脱出を早めることを発見しましたが、それはある一定の地点までです。

彼らは、これらの「つまみ」が待ち時間をどのように変化させるかについての具体的な公式（スケーリング則）を導き出しました。例えば、暗記から汎化へとジャンプするのにかかる時間は、おおよそ $1/(\eta \lambda)$ に比例します。これは、学習率や正則化を2倍にすれば、遅延時間が半分になることを意味します。

否定されたこと

この論文は、自分たちが何を主張していないかについても非常に慎重に述べています。彼らは、ネットワークが突然「回路を見つけた」とか、生物学的な目覚めのような神秘的なことが起きたと主張しているのではありません。むしろ、それは純粋に解空間の幾何学と、最適化プロセスのランダム性の結果であると論じています。また、これが単なる偶然の産物であるという考えも否定しています。彼らの数学は、これがAdamがこれらの特定のタイプの問題を最適化する際の、予測可能で構造的な特徴であることを示しています。

どの程度確かなのか？

著者らは自身の理論にかなりの自信を持っていますが、証明したものと測定したものを明確に区別しています。

• 数学： 彼らは、学習プロセスの連続時間モデルに基づいた厳密な数学的証明を用いて、これらのスケーリング則を導き出しました。彼らは、数式をチェックするためにシンボリック代数システム（数学計算用のプログラム）を使用しており、これにより方程式への高い信頼性を得ています。
• 証明： 彼らは、2つの特定の種類のパズルを用いて、自身のアイデアをテストしました。それは、群論的学習（具体的には対称群 $S_5$）と、剰余演算（具体的には整数 127 を法とする $Z_{127}$）です。
• 結果： これらの実験において、コンピュータの挙動は彼らの予測と一致しました。例えば、学習率やバッチサイズを変更すると、グロッキングにかかる時間は彼らの公式が予測した通りに変化しました。また、彼らの理論が起こると予測した「U字型」の暗記半径の曲線も確認されました。

しかし、彼らは、自身の数学が、小さな学習率や大きなバッチサイズといった特定の条件に依存していることも注記しています。彼らは、これがあらゆる可能なニューラルネットワークにおけるあらゆるインスタンスのグロッキングを説明すると主張しているのではなく、この「シェル・コア（殻と核）」の幾何学が存在する、これら特定の構造化されたタスクにおける現象を説明しているのだとしています。

結論

グロッキングは魔法ではなく、幾何学的な旅です。コンピュータは外側から始まり、「暗記の殻」に捕まり、学習に伴うランダムなノイズが自身を「汎化の核」へと押し出すまで、さまよい続けます。このジャンプにかかる時間は、ステップの速さ（学習率）、ステップの安定性（バッチサイズ）、そして中央へと引き寄せる力の強さ（正則化）によって決まります。著者らは数学によってこれらのルールをマッピングし、実験によって確認しました。これにより、なぜニューラルネットワークが、まるで長い眠りについているかのように、突然目覚める前に長い時間停滞することがあるのかについて、より明確な姿を描き出したのです。