[论文解读] Neural Networks Efficiently Learn Low-Dimensional Representations with SGD

Alireza Mousavi-Hosseini, Sejun Park+3 more|arXiv (Cornell University)|Sep 29, 2022

Stochastic Gradient Optimization Techniques被引用 6

一句话总结

该论文表明，即使输入维度 $d$ 较大，带有权重衰减的随机梯度下降（SGD）也能在两层 ReLU 神经网络中高效学习低维表示。它证明了第一层权重会收敛到由真实模型方向张成的 $k$-维主子空间，从而实现与网络宽度无关的泛化界，且单指数模型的样本复杂度为线性 $d$，优于核方法。

ABSTRACT

We study the problem of training a two-layer neural network (NN) of arbitrary width using stochastic gradient descent (SGD) where the input $\boldsymbol{x}\in \mathbb{R}^d$ is Gaussian and the target $y \in \mathbb{R}$ follows a multiple-index model, i.e., $y=g(\langle\boldsymbol{u_1},\boldsymbol{x} angle,...,\langle\boldsymbol{u_k},\boldsymbol{x} angle)$ with a noisy link function $g$. We prove that the first-layer weights of the NN converge to the $k$-dimensional principal subspace spanned by the vectors $\boldsymbol{u_1},...,\boldsymbol{u_k}$ of the true model, when online SGD with weight decay is used for training. This phenomenon has several important consequences when $k \ll d$. First, by employing uniform convergence on this smaller subspace, we establish a generalization error bound of $O(\sqrt{{kd}/{T}})$ after $T$ iterations of SGD, which is independent of the width of the NN. We further demonstrate that, SGD-trained ReLU NNs can learn a single-index target of the form $y=f(\langle\boldsymbol{u},\boldsymbol{x} angle) + ε$ by recovering the principal direction, with a sample complexity linear in $d$ (up to log factors), where $f$ is a monotonic function with at most polynomial growth, and $ε$ is the noise. This is in contrast to the known $d^{Ω(p)}$ sample requirement to learn any degree $p$ polynomial in the kernel regime, and it shows that NNs trained with SGD can outperform the neural tangent kernel at initialization. Finally, we also provide compressibility guarantees for NNs using the approximate low-rank structure produced by SGD.

研究动机与目标

• 理解当目标依赖于 $k \ll d$ 个方向时，宽两层神经网络中的 SGD 如何学习低维结构。
• 建立与网络宽度无关的 SGD 训练网络的泛化保证。
• 证明 SGD 训练的 ReLU 网络可以以与 $d$ 线性相关的样本复杂度学习某些多项式目标，优于核方法。
• 基于 SGD 所诱导的低秩结构，提供可压缩性保证。

提出的方法

• 使用在线 SGD 和权重衰减，分析任意宽度 $m$ 的两层 ReLU 神经网络的训练动态。
• 将目标建模为多指数函数 $y = g(\langle \mathbf{u}_1, \mathbf{x} \rangle, \dots, \langle \mathbf{u}_k, \mathbf{x} \rangle)$，其中输入为高斯分布且链接函数带噪声。
• 证明 SGD 的第一层权重会收敛到由 $\mathbf{u}_1, \dots, \mathbf{u}_k$ 张成的 $k$-维主子空间。
• 通过在低维子空间上的一致收敛性，推导出经过 $T$ 次迭代后的泛化误差界为 $\mathcal{O}(\sqrt{kd/T})$。
• 应用次高斯和次指数浓度不等式，以控制随机特征和梯度更新的行为。
• 利用 Rademacher 复杂度和对称化方法，界定泛化差距并建立样本复杂度结果。

实验结果

研究问题

• RQ1 当 $k \ll d$ 时，宽两层 ReLU 网络中的 SGD 是否会收敛到真实模型方向张成的主子空间？
• RQ2 SGD 训练网络的泛化误差是否可以独立于网络宽度进行界定？
• RQ3 SGD 训练的 ReLU 网络是否可以以与 $d$ 线性相关的样本复杂度学习 $p$-次多项式目标，而核方法则需要 $d^{\Omega(p)}$ 个样本？
• RQ4 SGD 训练网络的可压缩性如何？它与泛化性能有何关联？

主要发现

• SGD 训练的两层 ReLU 网络的第一层权重会收敛到由真实模型方向 $\mathbf{u}_1, \dots, \mathbf{u}_k$ 张成的 $k$-维主子空间。
• 经过 $T$ 次迭代后，泛化误差被界定为 $\mathcal{O}(\sqrt{kd/T})$，且与网络宽度 $m$ 无关。
• 对于单调 $f$ 且具有多项式增长的单指数模型，SGD 实现了与 $d$ 线性相关的样本复杂度（对数因子范围内），优于核方法所需的 $d^{\Omega(p)}$ 个样本。
• SGD 所诱导的低秩结构使得可压缩性保证成为可能，将模型复杂度与压缩表示的大小联系起来。
• SGD 训练的 ReLU 网络比初始化时的神经正切核更高效地学习某些 $p$-次多项式目标。
• 向主子空间的收敛性促进了更优的特征学习，从而实现了优于随机特征的泛化性能。