[論文レビュー] Towards Understanding Grokking: An Effective Theory of Representation Learning

Ziming Liu, O. Kitouni+4 more|arXiv (Cornell University)|May 20, 2022

Neural dynamics and brain function被引用数 24

ひとこと要約

この論文は、効果理論と相図を用いてグロッキングを説明し、一般化を構造化表現と結びつけ、トランスフォーマーとおもちゃモデルにおいて四つの学習フェーズ（理解、グロッキング、記憶、混乱）を特定している。

ABSTRACT

We aim to understand grokking, a phenomenon where models generalize long after overfitting their training set. We present both a microscopic analysis anchored by an effective theory and a macroscopic analysis of phase diagrams describing learning performance across hyperparameters. We find that generalization originates from structured representations whose training dynamics and dependence on training set size can be predicted by our effective theory in a toy setting. We observe empirically the presence of four learning phases: comprehension, grokking, memorization, and confusion. We find representation learning to occur only in a "Goldilocks zone" (including comprehension and grokking) between memorization and confusion. We find on transformers the grokking phase stays closer to the memorization phase (compared to the comprehension phase), leading to delayed generalization. The Goldilocks phase is reminiscent of "intelligence from starvation" in Darwinian evolution, where resource limitations drive discovery of more efficient solutions. This study not only provides intuitive explanations of the origin of grokking, but also highlights the usefulness of physics-inspired tools, e.g., effective theories and phase diagrams, for understanding deep learning.

研究の動機と目的

• ニューラルネットワークが、小さなアルゴリズムデータセットに過剰適合した後に一般化する理由を調査する。
• 一般化を、構造化された入力表現の出現と結びつける。
• 臨界的な訓練データサイズと学習経路を予測する物理学に触発された効果理論を開発する。
• 学習フェーズを特徴づけ、ハイパーパラメータが遅延一般化に与える影響を分析する。

提案手法

• 入力が学習可能な埋め込みへ写像され、その和をデコーダが処理する簡略化されたおもちゃモデルを構築する。
• 埋め込みの平行四辺形構造に基づく Representation Quality Index (RQI) を定義する。
• 埋め込みダイナミクスを支配し、ヘッセ行列スペクトルからグロッキング速度を導出する効果的損失 ell_eff を定式化する。
• おもちゃモデルとトランスフォーマー様設定において、表現学習率とデコーダ学習率およびウェイト減衰を変化させて相図を導出する。
• モジュラー加算および非可換群へ分析を拡張し、タスク間の一般性を検証する。

実験結果

リサーチクエスチョン

• RQ1 モデルが訓練データに過剰適合した場合、グロッキングにおいて一般化が生じる原因は何か？
• RQ2 訓練データサイズは、構造化された表現の出現と一般化にどのように影響するか？
• RQ3 遅延一般化はどの条件下で発生し得るか、またハイパーパラメータはどのように遅らせたり加速させたりできるか？
• RQ4 相図と効果理論の予測は、 toyモデルを超えてトランスフォーマーアーキテクチャやMNISTにも適用できるか？
• RQ5 異なる演算において一般化を可能にする表現構造の役割は何か？

主な発見

• 一般化は、表現空間に平行四辺形を作る構造化埋め込みの出現と相関する。
• 臨界的な訓練データの割合が、線形で一意な表現が現れる時期を決定し、一般化を可能にする。
• グロッキングは理解と記憶の間のフェーズであり、適切なハイパーパラメータ調整で移動させたり除去したりできる。
• 相図は四つの学習フェーズを示し、デコーダを速くするとかデコーダ容量が過剰だと記憶を促進する一方で、バランスの取れた学習は理解とグロッキングを支える。
• トランスフォーマ設定では、デコーダ容量を減らすか正則化を適用すると一般化が速まり、グロッキングを抑制できる。MNIST実験の証拠あり。
• 効果理論はグロッキングの時間と相転移点を予測し、ニューラルネットの訓練経路と定性的に一致する。