Grokking

グロッキングとは、汎化が遅れて起こる現象です。ニューラルネットワークが早い段階で学習データに対して完璧な精度を達成し、長い間検証データで偶然と同程度の精度にとどまった後、過学習しているように見える時期を過ぎてから、突然汎化する現象です。OpenAIのPowerら（2022）は、モジュラー算術などのアルゴリズムタスクで小規模なTransformerを学習させ、この現象に名前を付けました。学習精度が飽和したかなり後になって、検証精度がランダムな水準から完璧な水準まで上昇しました。この言葉は、Robert Heinleinの小説『異星の客』に由来します。作中の「grok」は、何かを完全に理解することを意味します。

グロッキングが進行する仕組み#

モジュラー加算は、グロッキングを調べる古典的なタスクです。ネットワークは(a, b)のペアを受け取り、限られた割合のペアだけで学習し、(a + b) mod pを出力します。以下の実行例では、p = 97における9,409ペアの30%を使い、PyTorchのAdamWオプティマイザーで小さな多層パーセプトロンを学習させます。CPUでも1分以内に完了します。

import torch
import torch.nn.functional as F

torch.manual_seed(0)
p = 97  # task: predict (a + b) mod 97 from one-hot encodings of a and b
a, b = torch.cartesian_prod(torch.arange(p), torch.arange(p)).T
x, y = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float(), (a + b) % p
idx = torch.randperm(p * p)
train, val = idx[: p * p * 3 // 10], idx[p * p * 3 // 10 :]  # train on 30% of all pairs

model = torch.nn.Sequential(torch.nn.Linear(2 * p, 256), torch.nn.ReLU(), torch.nn.Linear(256, p))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1.0)  # strong weight decay
for step in range(1, 20001):
    loss = F.cross_entropy(model(x[train]), y[train])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 4000 == 0:
        train_acc, val_acc = [(model(x[i]).argmax(1) == y[i]).float().mean().item() for i in (train, val)]
        print(f"step {step:5d}: train accuracy {train_acc:.2f}, validation accuracy {val_acc:.2f}")

ステップ4,000では、学習精度がすでに1.00なのに対し、検証精度は0.00です。その後、検証精度はステップ8,000で0.16、12,000で0.54、16,000で0.84、20,000で0.97まで上昇します。weight_decay=0.0を使うと、同じ実行でも20,000ステップすべてで検証精度は0.00のままです。

Nandaら（2023）は、モジュラー加算で学習した1層Transformerをリバースエンジニアリングして、この停滞を説明しました。ネットワークは離散フーリエ特徴量と三角関数の恒等式を使い、加算を円周上の回転に変換することを学習します。学習は、記憶、回路形成、整理という3つの連続した段階に分かれます。汎化を担う回路は、検証精度が横ばいの間に徐々に成長します。重み減衰によって、出力を依然として支配していた記憶解が取り除かれると、精度が急上昇します。

グロッキングが起きる要因#

• 正則化：元の実験では重み減衰が最も効果的な介入策であり、上記の実行では重み減衰がないとまったく汎化しません。
• データセットのサイズ：学習データの割合が小さいほど、汎化までに必要な最適化の量が大幅に増えます。一定のしきい値を下回ると、実用的な計算予算内ではネットワークが汎化しません。
• 重みノルム：Omnigrokは、汎化解の重みノルムが記憶解より小さいことから、グロッキングを重みの大きさと関連付けています。大きな初期重みから開始することで、著者らは画像、言語、分子のタスクでグロッキングを引き起こし、重みノルムを制約することでアルゴリズムタスクからグロッキングを取り除きました。
• 数値安定性：Prietoら（2025）は、正則化がない場合、交差エントロピー学習がロジットを拡大し続け、最終的に浮動小数点誤差によってsoftmaxの学習が停止することを示しました。彼らはこれをSoftmax Collapseと呼んでいます。StableMax活性化関数と⊥Gradオプティマイザーを使うと、重み減衰なしでグロッキングが起こります。

グロッキングと関連概念の比較#

• 過学習：学習時と検証時の性能差が持続または拡大する現象です。グロッキングでは、学習を続けることでこの差が自然に縮まります。
• 二重降下：モデルのサイズや学習時間の増加に伴い、テスト誤差が低下、上昇、再び低下します。エポック単位の二重降下も、改善の遅延が起きる現象の一種です。学習中にテスト誤差がしばらく上昇した後、再び低下します。
• メカニスティック解釈可能性：学習済みネットワーク内部のアルゴリズムを研究する分野です。モジュラー算術におけるグロッキングは、学習されたアルゴリズム全体を読み解けるため、この分野の主要なケーススタディの1つになりました。

実務におけるグロッキング#

グロッキングは、単一の厳密なルールに従う小規模なアルゴリズムデータセットで最も明確に見られます。それでもグロッキングは実験設計においても重要であり、早期停止は検証指標の改善が止まった時点で実行を終了するため、グロッキングの実行が停滞期の途中で打ち切られることになります。Ultralytics YOLOのトレーニングは、patience エポックにわたって検証フィットネスが改善しないと停止し（デフォルトは100です）、patience=0の場合は毎エポック実行されます。