[论文解读] Feature Learning in Infinite-Width Neural Networks

Greg Yang, J. Edward Hu|arXiv (Cornell University)|Nov 30, 2020

Topic Modeling参考文献 34被引用 20

一句话总结

本文提出了最大更新参数化（µP），一种改进的神经网络参数化方法，可在无限宽度极限下实现特征学习——与标准NTK参数化（在初始化时固定特征）不同。通过张量程序（Tensor Programs）技术，作者推导出µP的精确无限宽度极限，表明其在Word2Vec和少样本学习任务上均优于NTK基线和有限宽度网络，且有限宽度模型的性能随宽度增加而趋近于µP性能。

ABSTRACT

As its width tends to infinity, a deep neural network's behavior under gradient descent can become simplified and predictable (e.g. given by the Neural Tangent Kernel (NTK)), if it is parametrized appropriately (e.g. the NTK parametrization). However, we show that the standard and NTK parametrizations of a neural network do not admit infinite-width limits that can learn features, which is crucial for pretraining and transfer learning such as with BERT. We propose simple modifications to the standard parametrization to allow for feature learning in the limit. Using the *Tensor Programs* technique, we derive explicit formulas for such limits. On Word2Vec and few-shot learning on Omniglot via MAML, two canonical tasks that rely crucially on feature learning, we compute these limits exactly. We find that they outperform both NTK baselines and finite-width networks, with the latter approaching the infinite-width feature learning performance as width increases. More generally, we classify a natural space of neural network parametrizations that generalizes standard, NTK, and Mean Field parametrizations. We show 1) any parametrization in this space either admits feature learning or has an infinite-width training dynamics given by kernel gradient descent, but not both; 2) any such infinite-width limit can be computed using the Tensor Programs technique. Code for our experiments can be found at github.com/edwardjhu/TP4.

研究动机与目标

• 为解决神经正切核（NTK）理论的关键局限性，即无法捕捉过参数化网络中的特征学习。
• 识别在无限宽度极限下支持特征学习的参数化方式，其中标准参数化和NTK参数化均无法实现。
• 开发一种系统化方法，利用张量程序技术精确计算神经网络的无限宽度极限。
• 证明所提出的µP参数化在特征学习任务（如Word2Vec和基于MAML的少样本学习）中表现更优。

提出的方法

• 提出abc参数化框架，通过可调缩放参数（a_l, b_l, c）泛化标准参数化、NTK参数化和平均场参数化。
• 将最大更新参数化（µP）定义为一种特定的abc参数化，其中a1 = 0，对所有l ≥ 2有al = 1/2，bl = 0，c = 1/2，从而在无限宽度极限下实现最大化的特征学习。
• 应用张量程序技术，严格计算µP参数化的无限宽度极限，实现对网络动态的精确解析推导。
• 推导出网络隐藏表示和logits在无限宽度极限下的显式公式，从而实现对训练动态的精确计算。
• 通过在Word2Vec和基于MAML的少样本学习任务上的实验，验证理论框架，将µP与NTK、GP及有限宽度模型进行比较。
• 利用主成分分析（PCA）可视化和词类比准确率评估特征空间质量与泛化性能。

实验结果

研究问题

• RQ1 无限宽度神经网络能否学习特征？若能，其在何种参数化条件下可实现？
• RQ2 为何标准参数化和NTK参数化在无限宽度极限下无法支持特征学习？
• RQ3 何种数学结构的参数化可在无限宽度条件下同时实现稳定训练与非平凡的特征演化？
• RQ4 如何对任意架构和训练动态的神经网络，系统性地精确计算其无限宽度极限？
• RQ5 所提出的µP参数化在特征学习任务上的表现是否优于NTK和有限宽度模型？

主要发现

• 与NTK极限中特征在初始化时固定不同，最大更新参数化（µP）的无限宽度极限可实现非平凡的特征学习。
• 在text8数据集上，µP无限宽度模型的词类比准确率达到43.31%，显著优于NTK基线（0.0%）和有限宽度网络。
• 有限宽度µP网络的性能随宽度增加而趋近无限宽度性能，准确率从宽度2^6时的33.35%提升至2^10时的42.56%。
• 在更大的fil9数据集上，µP无限宽度模型达到56.45%的准确率，再次优于NTK基线（0.0%）和有限宽度模型。
• 动力学二分定理（Dynamical Dichotomy theorem）证明：任何稳定且非平凡的abc参数化，其极限要么是特征学习极限，要么是核极限，但不会同时具备两者。
• PCA可视化结果表明，µP嵌入能有意义地区分城市与州，而NTK嵌入在无限宽度下基本保持随机。