🤖 machine learning

Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

本論文は、モジュラー算術で訓練されたトランスフォーマーにおいて、重み減衰が記憶、一般化（グロッキング）、および崩壊の間の遷移を支配する重要な制御パラメータとして機能することを示し、これらのダイナミクスをさまざまなモデルスケールおよびアーキテクチャにわたって追跡するための、アテンション活性化に基づく 2 つの計算効率的なオンライン診断法を導入する。

原著者： Lucky Verma

公開日 2026-05-21

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Lucky Verma

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、コンピュータの脳内の「アテンションヘッド」と呼ばれる生徒グループに、97 時間しかない時計で時刻を計算するなど、モジュラー算術という厄介な数学パズルを解く方法を教えていると想像してください。

長い間、これらの生徒は答えを丸暗記しているように見えました。練習テストでは満点を取りますが、実際の試験では失敗します。なぜなら、彼らはゲームの「ルール」を本当に学んでいないからです。すると、突然、魔法のようなことが起こります。彼らはその概念を「グロク（深く理解）」するのです。彼らは暗記をやめ、理解を始め、それによって瞬時に新しいどんな問題も解けるようになります。

この論文は、この突然の「ひらめき」の瞬間が「いつ」「どのように」起こり、生徒をやりすぎたり足りなかったりするとどうなるかを理解するためのフィールドガイドです。

ここでは、単純なアナロジーを用いて発見された内容を解説します。

1. 学習の「音量ノブ」（重み減衰）

研究者たちは、「重み減衰」と呼ばれる設定が学習プロセスのマスター音量ノブとして機能することを発見しました。これはトレーニングの「個性」を制御します。

静かすぎる（重み減衰が低い）： 生徒たちはリラックスしすぎています。彼らは答えを丸暗記するだけで、根本的な論理を理解することはありません。彼らは「暗記モード」に留まり続けます。
ちょうど良い（重み減衰が中程度）： これが「ジャスト・ゴールドィロックス」ゾーンです。生徒たちは特定の 2 段階のダンスを経由します。
1. 同期フェーズ： まず、すべての生徒が全く同じように考え始めます（同期します）。彼らは基本ルールで一致します。
2. 分化フェーズ： 次に、彼らは専門化し始めます。ある生徒は加算の専門家になり、他の生徒は乗算の専門家になるなどです。彼らはクローンであることをやめ、多様なチームになります。ここで「グロッキング」が起こります。
うるさすぎる（重み減衰が高い）： ノブを上げすぎると、生徒たちは圧倒されてしまいます。彼らはすべてが混同し、学習を完全に停止する単一の状態に崩壊します。

2. 「安価な温度計」（オンライン診断）

通常、生徒が「ひらめき」の瞬間を迎えようとしているかどうかを知るには、学期の終わりまで待って最終試験を行う必要があります。これには長い時間と多額のお金（計算資源）がかかります。

著者たちは、生徒がまだ授業中である間に数分おきに確認できる、2 つの「安価でリアルタイムな温度計」を発明しました。

「グループハグ」メーター（コサイン類似度）： これは生徒たちが互いにどの程度同意しているかを測定します。彼らが「グロッキング」し始めると、まず強くハグ（同意）し、その後、専門化に伴ってゆっくりと離れていきます。
「カオス」メーター（エントロピーの標準偏差）： これは生徒たちがどの程度異議を唱えたり、考えを変えたりしているかを測定します。「ひらめき」の瞬間が近づくと、このメーターは急上昇し、生徒たちが同期状態から抜け出し、独自の役割を見つけようとしていることを示します。

これらのツールにより、研究者たちは最終試験を待つことなく、リアルタイムで「相転移」が起きている様子を確認できます。

3. 「臨界の転換点」

この論文は、音量ノブに対する非常に特定の数字（「転換点」）を見つけました。

設定が0.0158未満の場合、生徒たちはほとんどルールを学ばず、単に暗記します。
設定が0.0158より上の場合、彼らはほぼ確実にルールを学びます。
彼らが学ぶのに要する時間は、ノブを上げるほど（ある点まで）速くなり、予測可能な数学的曲線に従います。

4. 「後期クラッシュ」（アンチ・グロッキング）

ひねりがあります。もしこれらの生徒を（通常の数千ステップではなく）20,000 ステップという途方もなく長い間トレーニングすると、彼らの一部は突然すべてを忘れ、ランダムに推測し始めます。この論文はこれを「アンチ・グロッキング」と呼びます。

良い知らせ： これは全員に起こるわけではありません。これは「シード」（ランダムな初期条件）に依存します。一部の生徒は脆弱でクラッシュしますが、他の生徒はタフで永遠に賢いままです。
原因： これは、生徒たちが独自の役割を失い、壊れた方法で再び似通ったものになってしまうことによるものと思われます。

5. これは他の「脳」でも機能するか？

研究者たちは、この「音量ノブ」のルールが、標準的なもの（トランスフォーマー）だけでなく、異なる種類のコンピュータ脳でも機能するかどうかをテストしました。

彼らはMLP（単純なフィードフォワードネットワーク）、LSTM（古典的なメモリネットワーク）、そしてMamba（より新しい効率的なアーキテクチャ）で試しました。
結果： はい！「音量ノブ」は依然として機能します。それらすべてに「静かすぎる」「ちょうど良い」「うるさすぎる」のゾーンがあります。ただし、「ちょうど良い」ゾーンに到達するために必要なノブの正確な数字は、脳の種類によって異なります。

彼らが「主張していない」ことのまとめ

彼らは、これが（論文を書くやあなたとチャットする）巨大な言語モデルに適用されるとは言っていません。彼らがテストしたのは小さな数学モデルだけです。
彼らは、すべての AI に対する究極の「物理法則」を見つけたとは主張していません。彼らはこの特定の数学パズルに対する特定のパターンを見つけ、それが他のタスクでは異なる可能性があると慎重に述べています。
彼らは、これが医学的な治療法や生物学的な事実であると言っているのではありません。彼らは「心拍」や「同期」といった言葉をメタファーとして使用しただけです。

要約すると： この論文は、単一のダイヤル（重み減衰）を操作することで、コンピュータの脳が暗記するか、深く学習するか、あるいはクラッシュするかを制御できることを示しています。また、このプロセスがリアルタイムで起こる様子を観察するための 2 つの簡単なツールを提供し、学習は通常、2 つの明確な段階で起こることを証明しました。つまり、まず全員が同意し、次に全員が専門化するという段階です。

技術的サマリー：Grokking トランスフォーマーにおける重み減衰レジーム

問題定義
Grokking（突然の理解）とは、アルゴリズム的タスク（特にモジュラー演算）で訓練されたトランスフォーマーが、訓練損失が飽和したはるか後に、暗記から一般化への急激な遷移を示す現象を指す。メカニズム的研究は、この遷移の駆動因子として内部回路（例えばフーリエ表現）やインダクションヘッドの形成を特定しているが、既存の文献には、これらのダイナミクスをリアルタイムで追跡するための「安価なオンライン診断法」が存在しない。さらに、これらのレジームを制御するパラメータとしての重み減衰（ $\lambda$ ）の役割が、モデルスケール全体にわたって体系的にマッピングされたことはなく、また、重みの大きさの効果からアテンションヘッドの協調への具体的な寄与を分離した研究も存在しない。

方法論
著者は、AdamW オプティマイザを用いて、モジュラー演算タスク（mod+, mod-, mod $\times$ , mod $\div$ ）上でデコーダーのみのトランスフォーマーを訓練する。本研究では、重み減衰値の密なスウィープを行い、モデルスケール（0.82M から 85M パラメータ）およびアーキテクチャ（トランスフォーマー、MLP、LSTM、Mamba）を変化させる。

主要な方法論的構成要素は以下の通りである：

オンライン秩序変数: 各訓練ステップにおいて、アテンション活性化のみを用いて計算される 2 つのスカラー診断指標：
- 平均ペアワイズアテンションヘッドコサイン類似度（ $\bar{s}(t)$ ）: アテンションヘッド間の協調（同期）を測定する。
- エントロピー標準偏差（ $\sigma_H(t)$ ）: ヘッド間のシャノンエントロピーの分散を測定し、分化を追跡する。
有限サイズスケーリング: 著者は、重み減衰に対するグロッキングの確率（ $P(\text{grok})$ ）に対してロジスティック関数を当てはめ、臨界閾値 $\lambda_c$ を推定する。また、「グロッキングまでの時間」（ $t_{\text{grok}}$ ）に対してべき乗則を当てはめ、臨界指数 $\nu$ を決定する。
因果的介入: ヘッドのパターン構造と重みの大きさを区別するため、著者は、分化フェーズのピーク時にアテンションヘッドを再初期化する実験と、重みのノルムをクリップする実験を対照的に行う。
形式的検証: 診断的な同一性（境界と代数的性質）は、数学的な妥当性を確保するために Lean 4 で機械検証される。

主要な貢献

定量的臨界閾値: 本研究は、mod+ タスクにおける標準的な 4 層 8 ヘッドのトランスフォーマーに対して、鋭い重み減衰の臨界閾値を特定した： $\lambda_c = 0.0158$ （95% 信頼区間 [0.0109, 0.0200]）。
レジーム図: 2 軸（ $\lambda$ $λ$ , $N$ $N$ ）のマップは、3 つの明確なレジームを明らかにする：
- 暗記: $\lambda < \lambda_c$ （グロッキングはほぼゼロ）。
- 発達的グロッキング: $\lambda \ge \lambda_c$ （約 100% のグロッキング率に到達し、 $\lambda$ が増加するにつれてグロッキングまでの時間が減少する）。
- 崩壊: $\lambda = 10$ （すべてのヘッドが同一のパターンに収束し、 $\bar{s}=1.0$ となる）。
2 フェーズダイナミクス: 訓練軌道は以下の特徴を持つ：
- フェーズ 1（同期）: ヘッドが収束し（ $\bar{s}$ が約 0.93 から 0.995 へ上昇）、グロッキングの開始と一致する。
- フェーズ 2（分化）: ヘッドが発散し（ $\bar{s}$ が約 0.88 へ低下、 $\sigma_H$ が増加）する一方、テスト精度は高いまま維持される。
ヘッドごとの次元としての振幅変調器: 分化の振幅（ $\sigma_H$ ）は、経験的にヘッドごとの次元（ $d/H$ ）によって支配され、 $d/H$ が増加するにつれて単調に飽和する。
アーキテクチャ横断的範囲: 重み減衰制御による遷移は、アテンションを持たないアーキテクチャ（MLP、LSTM、Mamba）でも観察されるが、臨界閾値 $\lambda_c$ はアーキテクチャに依存して大幅にシフトし（おおよそ 1 桁の範囲）、アーキテクチャによって異なる。

結果

スケーリング指数: グロッキングまでの時間は、 $t_{\text{grok}} \propto (\lambda - \lambda_c)^{-\nu}$ というべき乗則に従い、経験的指数は $\nu = 0.757$ （95% 信頼区間 [0.725, 0.799]）である。著者は、現在のグリッド密度の下では、この値が標準的な参照指数（例えば $\nu=1/2$ や 3 次元イジング模型の $\nu \approx 0.63$ ）と一致しないことを明示的に指摘し、普遍性クラスの同定は将来の研究に委ねている。
長期的な脆弱性: 長期（20,000 エポック）の視野において、一部の軌道でシード依存性の「アンチグロッキング」による崩壊が観察され、精度の低下で終わる 5 段階のパターンが特徴である。これは普遍性のあるサイクルではなく、脆弱性として特定される。
因果的隔離: アテンションヘッドを再初期化することは、対照群と比較してフェーズ 2 の振幅（ $\sigma_H$ ）を大幅に減少させるが、重みノルムのクリップはそうしない。これにより、その効果は重みの大きさではなく、ヘッドのパターン構造に起因することが隔離される。
マルチタスク複製: 重み減衰制御のパターンは、4 つのモジュラー演算と 2 つのモデルスケール全体で維持され、加法を超えたレジームマップの堅牢性を確認する。

意義と主張
本論文は、損失地形の診断（改良された局所学習係数 rLLC などより高価な手法）を補完するものとして、アテンション活性化のみを用いたグロッキングダイナミクスを追跡する安価でオンライン計算可能な診断フレームワークを提供すると主張する。また、重み減衰を、暗記、発達的グロッキング、崩壊の各レジームを分離する主要な経験的制御パラメータとして確立する。

著者はその範囲において慎重である：

観測された指数 $\nu \approx 0.76$ について普遍性クラスを主張せず、テストされた参照指数が信頼区間外にあることを指摘している。
重み減衰がすべてのアーキテクチャにおける一般化に必要であるとは主張せず、標準的なアテンションアーキテクチャ内において制御パラメータとして機能するのみであると述べている。
分化がなぜ起こるのかという因果的メカニズムに関する主張は、将来の研究に委ねており、代わりにそれがいつ、どのように現れるかに焦点を当てている。
知見は小規模トランスフォーマーアテンションモデルにおけるモジュラー演算に限定されており、大規模言語モデルやより広範な普遍性に関する主張は範囲外である。

本研究は、有限サイズスケーリング、リャプノフ収縮、ガラス緩和に関する同時進行中の理論的作業を補完するものとして、「レジームマップ」と「秩序変数」を提供する現象論的かつ診断的な貢献として位置づけられている。