Grokking: 訓練精度が満点になってから27,400ステップ、何も起きなかった。そこから突然「わかった」
約7分で読めます

Grokking: 訓練精度が満点になってから27,400ステップ、何も起きなかった。そこから突然「わかった」


「訓練データを完璧に覚えた時点で、学習は終わり」——普通はそう考える。ところが機械学習には、そこからさらに何万ステップも回し続けると、ある日突然テストデータでも解けるようになるという奇妙な現象がある。grokking(グロッキング、「腑に落ちる」の意)と呼ばれている。

陰陽の記事で目玉が450エポック無視され続けたのを見て以来、この本家の現象を自分の手で再現したかった。mod加算 (a+b)mod53(a+b) \bmod 53 を小さなMLPに学習させ、weight decayの強さを振って追いかけた。

結果: 27,400ステップの沈黙

grokking曲線。訓練精度(青)はstep 200で1.0に到達し、以後ずっと満点のまま。テスト精度(赤)は0付近に張り付き続け、step 27,600で急激に立ち上がって1.0に達する。図内注釈「step 200 訓練は満点に」「step 27,600 ここで突然わかる」、その間に「何も起きていないように見える 27,400 ステップ」の両矢印

  • step 200: 訓練精度が99%に到達。1124問の答えを完全に暗記した
  • step 200〜27,600: テスト精度は0.000のまま微動だにしない。訓練損失も既にゼロなので、外から見える数字は何ひとつ動かない
  • step 27,600: テスト精度が急上昇し、1.000(未知の問題を全問正解)に到達

グラフの横軸は対数目盛りだが、それでもこの平坦区間の長さは異常に見える。訓練が終わってから、理解が始まるまでに137倍の時間がかかった

失敗談: weight decayを1行間違えて、重みを全て0に潰した

正直に書くと、1回目の実験は完全な失敗だった。weight decayを勾配に足す形で実装したのだ。

g = g + wd * p       # ← これが間違い(結合型L2正則化)
mom = B1*mom + (1-B1)*g
p -= lr * mom_hat / (sqrt(vel_hat) + eps)

結果、wd=1.0では重みノルムが即座に0に崩壊し、訓練精度2.6%(ほぼランダム)のまま4万ステップが終わった。原因はAdamの正規化にある。Adamは勾配を分散で割って正規化するので、wd * p の項を勾配に混ぜると、その項も正規化されて重みの大きさに関係なく一定の力で0に引っ張るようになる。

正しくは、AdamWの「分離型(decoupled)」——Adamの更新とは別に、最後に独立して縮める。

p -= lr * mom_hat / (sqrt(vel_hat) + eps)
p -= lr * wd * p     # ← 正解(分離型)。Adamの正規化を通さない

AdamWの論文タイトルが「Decoupled Weight Decay Regularization」であることの意味を、頭ではなく壊して理解した。1行の違いで、モデルは学習不能にも、grokkingするようにもなる。

weight decayが弱すぎると、永遠に暗記のまま

修正版で、weight decayを0から1.0まで5段階振った。

左: テスト精度の推移。wd=0.0と0.03は最後まで0のまま、wd=0.1は0.317まで、wd=0.3はstep 27,600で1.0に、wd=1.0はstep 10,400で1.0に到達する。右: 重みの大きさの推移。すべて一度1400ステップ付近で最大になった後、wdが強いほど大きく減少する

weight decay訓練99%到達テスト90%到達最終テスト精度
0.0step 200到達せず0.000
0.03step 200到達せず0.000
0.1step 200到達せず0.317
0.3step 200step 27,6001.000
1.0step 200step 10,4001.000

訓練精度はどの条件でも同じくstep 200で満点になる。違うのはその後だけだ。weight decayがないと、6万ステップ回しても汎化は一度も訪れない(テスト精度0.000のまま)。強くするほど「わかる」のが早くなる。

重みが縮んだ先に、理解がある

なぜweight decayが必要なのか。重みの大きさとテスト精度を重ねると、答えが見える。

テスト精度(赤)と重みの大きさ(青)の重ね描き。重みはstep 1400付近で最大74まで増加した後、減少に転じる。重みが減少していく途中のstep 27,000付近で、テスト精度が急上昇する

学習は2つの局面に分かれていた。

  1. 暗記の局面(〜step 1,400): 重みがぐんぐん大きくなる。個々の問題を覚え込むために、大きな重みで無理やり答えを合わせている
  2. 整理の局面(step 1,400〜): 訓練精度は満点のまま、weight decayが重みを削り続ける。暗記を維持できるギリギリまで、余計なものが削ぎ落とされていく
  3. その先で、削ぎ落としきった構造がたまたま「一般的な規則」と一致した瞬間に、汎化が起きる

暗記を保ったまま重みを減らし続けると、最後には「mod加算という規則そのもの」——覚えるより計算した方が安上がりな表現——にたどり着く。汎化は、正解を覚える競争ではなく、同じ答えをより小さく表現する競争の副産物だった。

手を動かして意外だったこと

いちばん怖いと思ったのは、27,400ステップの間、どの指標を見ても進捗が観測できないことだ。訓練損失はゼロ、訓練精度は満点、テスト精度は0.000。もし「テスト精度が5,000ステップ改善しなければ打ち切り」というearly stoppingを入れていたら、この実験は永遠に失敗で終わっていた。しかも打ち切った時点のモデルは、訓練データに関しては完璧なのだ。

陰陽の記事の「目玉が450エポック無視される」も、Q学習の「価値が1歩ずつしか遡らない」も、同じ形をしている。学習の内部では何かが進んでいるのに、外から見える数字は動かない期間がある。数字が動かないことは、何も起きていないことの証明にはならない。

そして、汎化が「重みを削る過程」の産物だったという点。伸び悩んでいる期間にやるべきことは、新しい知識を足すことではなく、すでに知っていることをより少ない要素で説明し直すことなのかもしれない。覚えた量が同じでも、表現が小さくなるほど、まだ見ぬ問題に届くようになる。

まとめ

  • mod加算タスク((a+b)mod53(a+b) \bmod 53、訓練40%)で grokking を再現。訓練99%到達はstep 200、テスト90%到達はstep 27,600——137倍の遅れ
  • その間の27,400ステップ、訓練損失も訓練精度もテスト精度も動かない。early stoppingを入れていたら永遠に到達できなかった
  • 実装の失敗: weight decayを勾配に足す結合型で書くとAdamの正規化に飲まれ、重みが0に崩壊してモデルが学習不能になる。AdamWの分離型(更新後に独立して縮める)が正解
  • weight decayスイープ: 0.0/0.03では6万ステップ経っても汎化ゼロ、0.1で0.317、0.3でstep 27,600に到達、1.0でstep 10,400に到達——強いほど早く「わかる」
  • 重みの推移を見ると、step 1,400までは増加(暗記)、以降は減少(整理)。汎化は重みが削られていく途中で突然起きる——同じ答えをより小さく表現できたときに一般化が生まれる