[论文解读] How Two-Layer Neural Networks Learn, One (Giant) Step at a Time

Yatin Dandi, Florent Krząkała+3 more|arXiv (Cornell University)|May 29, 2023

Neural Networks and Applications被引用 7

一句话总结

本文通过早期梯度下降分析两层神经网络如何学习结构化目标函数，表明小批量大小在单步内决定学习目标函数的哪些方向。研究揭示，学习‘困难’方向（跳跃指数为 $\ell$）需要 $n = \mathcal{O}(d^\ell)$ 个样本，但通过 $n = \mathcal{O}(d)$ 的多步学习可分层学习多个方向，显著提升近似能力与泛化性能，超越懒惰参数化区域。

ABSTRACT

For high-dimensional Gaussian data, we investigate theoretically how the features of a two-layer neural network adapt to the structure of the target function through a few large batch gradient descent steps, leading to an improvement in the approximation capacity from initialization. First, we compare the influence of batch size to that of multiple steps. For a single step, a batch of size $n = \mathcal{O}(d)$ is both necessary and sufficient to align with the target function, although only a single direction can be learned. In contrast, $n = \mathcal{O}(d^2)$ is essential for neurons to specialize in multiple relevant directions of the target with a single gradient step. Even in this case, we show there might exist ``hard'' directions requiring $n = \mathcal{O}(d^\ell)$ samples to be learned, where $\ell$ is known as the leap index of the target. Second, we show that the picture drastically improves over multiple gradient steps: a batch size of $n = \mathcal{O}(d)$ is indeed sufficient to learn multiple target directions satisfying a staircase property, where more and more directions can be learned over time. Finally, we discuss how these directions allow for a drastic improvement in the approximation capacity and generalization error over the initialization, illustrating a separation of scale between the random features/lazy regime and the feature learning regime. Our technical analysis leverages a combination of techniques related to concentration, projection-based conditioning, and Gaussian equivalence, which we believe are of independent interest. By pinning down the conditions necessary for specialization and learning, our results highlight the intertwined role of the structure of the task to learn, the details of the algorithm, and the architecture, shedding new light on how neural networks adapt to the feature and learn complex task from data over time.

研究动机与目标

• 理解两层神经网络在早期训练中如何适应低维目标结构。
• 量化特征学习所需的样本复杂度，涵盖小批量大小、网络宽度与目标复杂度。
• 阐明近似能力与泛化性能方面，从懒惰参数化区域到特征学习区域的转变过程。
• 识别网络在高维设置下高效学习相关方向的条件。
• 基于埃尔米特分解与集中性技术，提供学习动态的精确理论刻画。

提出的方法

• 利用埃尔米特分解将目标函数表示为正交多项式展开，其中跳跃指数 $\ell$ 用于标识首个非零阶次。
• 应用高阶奇异值分解（HOSVD）以分解目标函数的张量结构，并识别相关学习方向。
• 借助测度集中与基于投影的条件化技术，分析网络神经元与目标子空间之间的重叠。
• 采用高斯等价技术，简化在高维随机初始化下梯度动态的分析。
• 研究聚焦于大批次的梯度下降，将学习过程建模为一系列巨步更新。
• 基于小批量大小 $n$、输入维度 $d$ 与目标函数跳跃指数 $\

实验结果

研究问题

• RQ1 两层网络在单步梯度下降中学习结构化目标函数的单个相关方向，所需的最小小批量大小是多少？
• RQ2 目标函数的跳跃指数 $\ell$ 如何影响学习困难方向的样本复杂度？
• RQ3 是否可在单步内学习多个相关方向？若可以，需满足何种小批量大小条件？
• RQ4 梯度步数如何影响网络学习日益复杂目标结构的能力？
• RQ5 与懒惰参数化区域相比，当发生特征学习时，近似能力与泛化误差的提升程度如何？

主要发现

• 小批量大小为 $n = \mathcal{O}(d)$ 时，网络可在单步内与目标函数的单个相关方向对齐，该条件既必要又充分。
• 在单步内学习多个方向需 $n = \mathcal{O}(d^2)$，而对于跳跃指数为 $\ell$ 的‘困难’方向，需 $n = \mathcal{O}(d^\ell)$ 个样本。
• 通过多步梯度下降且 $n = \mathcal{O}(d)$，网络可分层学习多个目标方向，满足阶梯状特性。
• 网络的近似能力在初始化后显著提升，尤其在发生特征学习时，与懒惰参数化区域形成明显分离。
• 即使使用大批次，当网络通过多步学习相关方向时，泛化误差仍被限制且显著改善。
• 基于埃尔米特分解与重叠集中性的理论框架，提供了学习动态的精确、非渐近界，无需依赖特定常数的样本复杂度。