[论文解读] 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

一句话总结

本文在高维渐近条件下分析了两层神经网络中第一层权重的首个梯度步，表明即使仅一步即可在权重矩阵中诱导出一个秩-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.

研究动机与目标

• 理解第一层权重的首个梯度更新如何改进两层神经网络中的表征学习。
• 在比例渐近框架（n, d, N → ∞ 且速率相同）下，刻画该更新对核岭回归性能的影响。
• 确定在何种学习率缩放下，更新后的特征可超越随机特征和线性模型。
• 为早期训练中非核行为建立理论基础，将其与特征适应联系起来。

提出的方法

• 分析两层ReLU或erf激活神经网络中第一层权重W的首个梯度下降步，网络随机初始化。
• 采用学生-教师框架，使用单指标教师模型f*以隔离特征学习的影响。
• 在比例极限（n, d, N → ∞ 且 n/d, N/d → 常数）下使用高维渐近分析，推导预测风险的确定性等价形式。
• 推导出更新后权重矩阵的结构为一个秩-1尖刺加噪声，表明其与f*的线性分量对齐。
• 利用谱分析和随机矩阵理论，计算在一次梯度步后对共轭核进行岭回归的预测风险。
• 考虑学习率的两种缩放：小η（高斯等价性）与大η（超越线性区域），分别推导风险上界。

实验结果

研究问题

• RQ1 第一层权重的首个梯度步是否在表征中诱导出超越随机特征的结构化变化？
• RQ2 这一单步更新是否可在高维条件下实现对随机特征和线性模型的性能超越？
• RQ3 学习率缩放如何影响更新后特征在超越线性区域外泛化的能力？
• RQ4 一次梯度步后核的精确统计行为是什么？其与教师模型结构有何关联？

主要发现

• 第一层权重的首个梯度更新在权重矩阵中诱导出一个秩-1尖刺，导致与教师模型f*的线性分量对齐。
• 对于小学习率，训练后的特征实现高斯等价性，预测风险优于随机特征，但受线性模型风险的上界约束。
• 对于足够大的学习率，更新后特征上的岭估计器可超越输入空间上任意线性模型的性能，表明其已脱离线性区域。
• 一步后的预测风险上界为10τ* + C′(√τ*·√(d/n) + d/n)，其中τ*为f*的非线性分量的范数，表明其依赖于教师模型的非线性程度。
• 性能提升源于核矩阵的主奇异向量与训练标签对齐，该对齐由秩-1更新驱动。
• 结果表明，即使仅一个梯度步即可实现统计上显著的特征学习，并可超越固定核方法。