[論文レビュー] Feature Learning in Infinite-Width Neural Networks

Greg Yang, J. Edward Hu|arXiv (Cornell University)|Nov 30, 2020

Topic Modeling参考文献 34被引用数 20

ひとこと要約

本稿では、無限幅極限における特徴の学習を可能にする修正されたニューラルネットワークのパrametrization、Maximal Update Parametrization (µP) を提案する。標準のNTKパラメータ化とは異なり、µPは初期化時の特徴を固定しない。著者らはTensor Programs手法を用いてµPの正確な無限幅極限を導出し、Word2Vecおよび少サンプル学習のタスクでNTKベースラインおよび有限幅ネットワークを上回ることを示した。有限幅モデルは幅が増加するにつれてµPの性能に近づく。

ABSTRACT

As its width tends to infinity, a deep neural network's behavior under gradient descent can become simplified and predictable (e.g. given by the Neural Tangent Kernel (NTK)), if it is parametrized appropriately (e.g. the NTK parametrization). However, we show that the standard and NTK parametrizations of a neural network do not admit infinite-width limits that can learn features, which is crucial for pretraining and transfer learning such as with BERT. We propose simple modifications to the standard parametrization to allow for feature learning in the limit. Using the *Tensor Programs* technique, we derive explicit formulas for such limits. On Word2Vec and few-shot learning on Omniglot via MAML, two canonical tasks that rely crucially on feature learning, we compute these limits exactly. We find that they outperform both NTK baselines and finite-width networks, with the latter approaching the infinite-width feature learning performance as width increases. More generally, we classify a natural space of neural network parametrizations that generalizes standard, NTK, and Mean Field parametrizations. We show 1) any parametrization in this space either admits feature learning or has an infinite-width training dynamics given by kernel gradient descent, but not both; 2) any such infinite-width limit can be computed using the Tensor Programs technique. Code for our experiments can be found at github.com/edwardjhu/TP4.

研究の動機と目的

• 過パラメータ化されたネットワークにおける特徴の学習を捉えられないというニューラルトランジットカーネル（NTK）理論の根本的限界を是正すること。
• 標準的およびNTKパラメータ化が機能しない無限幅極限においても特徴の学習を可能にするパラメータ化を同定すること。
• Tensor Programs手法を用いて、任意のアーキテクチャおよび学習ダイナミクスに対して、ニューラルネットワークの正確な無限幅極限を体系的に計算する方法を開発すること。
• 提案されたµPパラメータ化が、Word2VecやMAMLを用いた少サンプル学習といった特徴学習タスクで優れた性能を示すことを実証すること。

提案手法

• 調整可能なスケーリングパラメータ（a_l, b_l, c）を備えたabcパラメータ化フレームワークを提案し、標準的、NTK、およびマイルドフィールドパラメータ化を一般化する。
• 特徴学習を最大限に可能にする無限幅極限における、a1 = 0、l ≥ 2 で al = 1/2、bl = 0、c = 1/2 とする特定のabcパラメータ化としてMaximal Update Parametrization (µP) を導入する。
• Tensor Programs手法を適用して、µPパラメータ化の無限幅極限を厳密に計算し、ネットワークダイナミクスの正確な解析的導出を可能にする。
• ネットワークの隠れ表現およびログチの無限幅極限に対する明示的公式を導出し、トレーニングダイナミクスの正確な計算を可能にする。
• Word2VecおよびMAMLを用いた少サンプル学習における理論的枠組みの妥当性を実験的に検証し、µPをNTK、GP、および有限幅モデルと比較する。
• PCA可視化および語の類推精度を用いて、特徴空間の質および一般化性能を評価する。

実験結果

リサーチクエスチョン

• RQ1 無限幅ニューラルネットワークは特徴を学習可能か？ もしそうなら、どのようなパラメータ化条件下で可能か？
• RQ2 なぜ標準的およびNTKパラメータ化は無限幅極限において特徴学習をサポートできないのか？
• RQ3 安定したトレーニングと非自明な特徴の進化の両方を可能にするパラメータ化の数学的構造は何か？
• RQ4 任意のアーキテクチャおよび学習ダイナミクスに対して、ニューラルネットワークの無限幅極限を正確かつ体系的に計算する方法は何か？
• RQ5 提案されたµPパラメータ化は、NTKおよび有限幅モデルよりも特徴学習タスクで優れた性能を発揮するか？

主な発見

• Maximal Update Parametrization (µP) の無限幅極限は、特徴が初期化時に固定されるNTK極限とは異なり、非自明な特徴学習を可能にする。
• text8データセットでは、µPの無限幅モデルが43.31%の語の類推精度を達成し、NTKベースライン（0.0%）および有限幅ネットワークを著しく上回った。
• 有限幅のµPネットワークは幅が増加するにつれて無限幅性能に近づき、幅2^6では33.35%から幅2^10では42.56%まで精度が上昇した。
• より大きなfil9データセットでは、µPの無限幅モデルが56.45%の精度を達成し、再びNTKベースライン（0.0%）および有限幅モデルを上回った。
• 動的二分法則（Dynamical Dichotomy theorem）により、任意の安定的かつ非自明なabcパラメータ化は、特徴学習極限またはカーネル極限のいずれかを生成するが、両方を同時に生成することはできないことが証明された。
• PCA可視化により、µP埋め込みは都市と州を意味的に明確に分離している一方、NTK埋め込みは無限幅状態においてほとんどランダムなままであった。