Grokking as the Transition from Lazy to Rich Training Dynamics

要約

私たちは、ニューラル ネットワークのトレーニング損失がテスト損失よりもはるかに早く減少するグロッキング現象が、ニューラル ネットワークが遅延トレーニング ダイナミクスから豊富な特徴学習体制に移行するために発生する可能性があると提案します。
このメカニズムを説明するために、既存の理論では説明できない方法で正則化なしでグロッキングを示す 2 層ニューラル ネットワークを使用した多項式回帰問題に対するバニラ勾配降下の簡単な設定を研究します。
このようなネットワークのテスト損失について十分な統計を特定し、トレーニングを通じてこれらを追跡すると、ネットワークが最初にカーネル回帰解をその初期特徴に適合させようとしたとき、この設定でグロッキングが発生し、その後、遅延時間の特徴学習が行われたことが明らかになります。
一般化する解決策は、列車損失がすでに低くなった後に特定されます。
グロッキングの主な決定要因は、特徴学習率 (ネットワーク出力をスケールするパラメーターによって正確に制御可能) と、初期特徴量とターゲット関数 $y(x)$ の調整であることがわかりました。
私たちは、この一般化の遅れは、(1) 初期ニューラル タンジェント カーネルの上位固有ベクトルとタスク ラベル $y(x)$ の位置がずれているが、(2) データセットのサイズが十分に大きいため、ネットワークがそれを可能にする場合に発生すると主張します。
最終的には一般化しますが、トレーニング損失がすべてのエポックでテスト損失を完全に追跡するほど大きくはなく、(3) ネットワークは遅延体制でトレーニングを開始するため、特徴をすぐには学習しません。
最後に、遅延 (線形モデル) からリッチ トレーニング (特徴量学習) への移行により、MNIST、1 層トランスフォーマー、生徒と教師のネットワークなど、より一般的な設定でグロッキングを制御できるという証拠を示します。

要約(オリジナル)

We propose that the grokking phenomenon, where the train loss of a neural network decreases much earlier than its test loss, can arise due to a neural network transitioning from lazy training dynamics to a rich, feature learning regime. To illustrate this mechanism, we study the simple setting of vanilla gradient descent on a polynomial regression problem with a two layer neural network which exhibits grokking without regularization in a way that cannot be explained by existing theories. We identify sufficient statistics for the test loss of such a network, and tracking these over training reveals that grokking arises in this setting when the network first attempts to fit a kernel regression solution with its initial features, followed by late-time feature learning where a generalizing solution is identified after train loss is already low. We find that the key determinants of grokking are the rate of feature learning — which can be controlled precisely by parameters that scale the network output — and the alignment of the initial features with the target function $y(x)$. We argue this delayed generalization arises when (1) the top eigenvectors of the initial neural tangent kernel and the task labels $y(x)$ are misaligned, but (2) the dataset size is large enough so that it is possible for the network to generalize eventually, but not so large that train loss perfectly tracks test loss at all epochs, and (3) the network begins training in the lazy regime so does not learn features immediately. We conclude with evidence that this transition from lazy (linear model) to rich training (feature learning) can control grokking in more general settings, like on MNIST, one-layer Transformers, and student-teacher networks.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google