[論文レビュー] High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation

Jimmy Ba, Murat A. Erdogdu+4 more|arXiv (Cornell University)|May 3, 2022

Stochastic Gradient Optimization Techniques被引用数 11

ひとこと要約

本稿は、高次元漸近的条件下における2層ニューラルネットワークの1層目重みに対する最初の勾配ステップを分析し、1ステップの更新ですでに重み行列にランク1のスパイクが生じ、教師モデルの線形成分と整合するようになると示している。この整合性のおかげで、更新された特徴量に対するリッジ回帰はランダム特徴量を上回り、十分に大きな学習率では入力に対する最適な線形モデルでさえも上回る。これは、比例的漸近的条件下における初期特徴量学習の強力さを示している。

ABSTRACT

We study the first gradient descent step on the first-layer parameters $\boldsymbol{W}$ in a two-layer neural network: $f(\boldsymbol{x}) = \frac{1}{\sqrt{N}}\boldsymbol{a}^ opσ(\boldsymbol{W}^ op\boldsymbol{x})$, where $\boldsymbol{W}\in\mathbb{R}^{d imes N}, \boldsymbol{a}\in\mathbb{R}^{N}$ are randomly initialized, and the training objective is the empirical MSE loss: $\frac{1}{n}\sum_{i=1}^n (f(\boldsymbol{x}_i)-y_i)^2$. In the proportional asymptotic limit where $n,d,N o\infty$ at the same rate, and an idealized student-teacher setting, we show that the first gradient update contains a rank-1 "spike", which results in an alignment between the first-layer weights and the linear component of the teacher model $f^*$. To characterize the impact of this alignment, we compute the prediction risk of ridge regression on the conjugate kernel after one gradient step on $\boldsymbol{W}$ with learning rate $η$, when $f^*$ is a single-index model. We consider two scalings of the first step learning rate $η$. For small $η$, we establish a Gaussian equivalence property for the trained feature map, and prove that the learned kernel improves upon the initial random features model, but cannot defeat the best linear model on the input. Whereas for sufficiently large $η$, we prove that for certain $f^*$, the same ridge estimator on trained features can go beyond this "linear regime" and outperform a wide range of random features and rotationally invariant kernels. Our results demonstrate that even one gradient step can lead to a considerable advantage over random features, and highlight the role of learning rate scaling in the initial phase of training.

研究の動機と目的

• 2層ニューラルネットワークにおける1層目重みに対する最初の勾配更新が表現学習をどのように改善するかを理解すること。
• 比例的漸近的領域（n, d, N → ∞ で同じ割合）におけるこの更新がカーネルリッジ回帰の性能に与える影響を特定すること。
• 更新された特徴量がランダム特徴量や線形モデルを上回るための学習率スケーリングを同定すること。
• 初期学習における非カーネル的挙動の理論的基盤を確立し、特徴量適応と結びつけること。

提案手法

• 2層のReLUまたはerf活性化関数を用いたニューラルネットワーク（初期化はランダム）における1層目重みWに対する最初の勾配降下ステップを分析する。
• 特徴量学習の効果を分離するために、単一インデックス教師モデルf*を用いた学生-教師フレームワークを採用する。
• 予測リスクの決定的同等物を導出するため、比例的漸近的領域（n, d, N → ∞ でn/d, N/d → 定数）における高次元的漸近的解析を用いる。
• 更新された重み行列の構造をランク1のスパイク＋ノイズとして導出し、教師モデルf*の線形成分と整合する様子を示す。
• スペクトル解析と確率的行列理論を用いて、1ステップ後の共役カーネルにおけるリッジ回帰の予測リスクを計算する。
• 学習率の2通りのスケーリングを検討する：小さなη（ガウス的同等性）と大きなη（線形領域を超える）。それぞれの状況でリスクの上限を導出する。

実験結果

リサーチクエスチョン

• RQ1 1層目重みに対する最初の勾配ステップは、ランダム特徴量を超える構造的変化を表現に引き起こすか？
• RQ2 この1回の更新が、高次元的領域においてランダム特徴量や線形モデルを上回る性能向上をもたらすか？
• RQ3 学習率のスケーリングは、更新された特徴量が線形領域を超えて一般化する能力にどのように影響するか？
• RQ4 1ステップ後のカーネルの正確な統計的挙動は何か？また、教師モデルの構造とどのように関係するか？

主な発見

• Wに対する最初の勾配更新により、重み行列にランク1のスパイクが生じ、教師モデルf*の線形成分と整合する。
• 小さな学習率では、訓練済み特徴量はガウス的同等性を達成し、予測リスクはランダム特徴量を上回るが、依然として線形モデルのリスクに上限を受ける。
• 十分に大きな学習率では、更新された特徴量に対するリッジ推定器は、入力に対するあらゆる線形モデルの性能を上回る。これは線形領域からの逸脱を示している。
• 1ステップ後の予測リスクは、10τ* + C′(√τ*·√(d/n) + d/n) で上界が与えられ、τ*はf*の非線形成分のノルムに依存しており、教師モデルの非線形性に依存することが示された。
• この改善は、カーネル行列の最大特異ベクトルが訓練ラベルと整合するようになることに起因し、これはランク1の更新によって駆動される。
• 結果として、1ステップの勾配更新でも、統計的に有意な特徴量学習が可能であり、固定カーネル手法を上回る可能性があることが確立された。