特征学习的高维渐近:单步梯度如何改善表示

我们研究两层神经网络中第一层参数 $\boldsymbol{W}$ 的首个梯度下降步:$f(\boldsymbol{x}) = \frac{1}{\sqrt{N}}\boldsymbol{a}^\top\sigma(\boldsymbol{W}^\top\boldsymbol{x})$,其中 $\boldsymbol{W}\in\mathbb{R}^{d\times N}, \boldsymbol{a}\in\mathbb{R}^{N}$ 被随机初始化,训练目标为经验 MSE 损失:$\frac{1}{n}\sum_{i=1}^n (f(\boldsymbol{x}_i)-y_i)^2$。在 $n,d,N\to\infty$ 以相同速率增长的比例渐近极限下,以及一个理想化的师生设定中,我们证明首个梯度更新包含一个秩为 1 的“尖峰”,其导致第一层权重与教师模型 $f^*$ 的线性分量之间的对齐。为刻画该对齐的影响,当 $f^*$ 为单指标模型时,我们计算了在 $\boldsymbol{W}$ 上以学习率 $\eta$ 执行一步梯度后,基于共轭核的岭回归的预测风险。我们考虑首步学习率 $\eta$ 的两种缩放。对于小的 $\eta$,我们建立了训练特征映射的高斯等价性质,并证明所学核优于初始随机特征模型,但无法击败输入上的最佳线性模型。而对于足够大的 $\eta$,我们证明对于某些 $f^*$,在训练特征上的同一岭估计器可以超越该“线性 regime”并优于广泛的随机特征与旋转不变核。我们的结果表明,即使单步梯度也能带来相对于随机特征的显著优势,并凸显了训练初始阶段中学习率缩放的作用。