谷歌开源Neural Tangents：简单快速训练无限宽度神经网络

无限宽的神经网络可以训练了？神经正切核果然可以发挥作用了。

近日，谷歌一项研究成果在社区引发了关注，该研究宣称可通过神经正切核（Neural Tangent Kernel）使用贝叶斯推理或梯度下降分析式地训练无限宽度的神经网络。使用谷歌开源的软件库 Neural Tangents，这个过程不仅简单且快速，而且效果非常好，甚至只需 5 行代码就能一步到位地构建并训练这种无限宽度网络的集成模型！该论文已被 ICLR 2020 接收为 Spotlight 论文，相关链接见文末。

在自然语言处理、会话智能体和连接组学等许多领域，深度学习都已取得了广泛的成功，机器学习领域的研究图景也已经发生了变革。不过，仍还有一些有趣而又重要的问题有待解答，比如：为什么即使在过度参数化时，深度神经网络（DNN）也能取得非常好的泛化能力？深度网络的架构、训练和性能之间有何关系？如何提取出深度学习模型中的显著特征？

我们何以在近些年里取得了如此的进步？一大关键理论见解是：增加 DNN 的宽度能使 DNN 的行为更有规律可循，也就人更容易理解它们。近来的许多研究已经表明，宽度可无限扩增的 DNN 可以收敛成另一类更为简单的名为「高斯过程（Gaussian processes）」的模型。因此，贝叶斯推理与卷积神经网络的梯度下降情况等复杂现象便可约简为简单的线性代数方程。从这些无限宽度网络所得到的见解通常也适用于它们对应的有限版本。因此，无限宽度网络可用作研究深度学习的透镜，而且它们本身也可用作有用的模型。

左图：深度神经网络变得无限宽时推导输入/输出映射关系的示意图。右图：随着神经网络宽度的增长，可以看到在网络的不同随机初始化上得到的输出分布逐渐变成高斯分布。

不幸的是，推导有限网络的无限宽度极限需要大量数学专业知识，而且必须针对每种不同架构专门研究。一旦推导得到了无限宽度模型，如果要得到一种高效且可扩展的实现，还进一步需要大量工程开发专业技能。总的来说，将有限宽度模型扩展成对应的无限宽度网络的这个过程可能耗费数月时间，而且这个研究本身就足以支撑起一篇研究论文。

为了解决这个问题，谷歌开源了一个用 Jax 写的开源软件库 Neural Tangents。谷歌宣称这个软件库可让无限宽度神经网络的构建与训练变得与有限宽度神经网络一样简单。究其核心，Neural Tangents 提供了一个易于使用的神经网络软件库，其可同时为神经网络构建有限和无限宽度的版本。

举个实际使用 Neural Tangents 的例子：假设要在某些数据上训练一个全连接神经网络，通常的做法是随机初始化一个神经网络，然后使用梯度下降来训练它。初始化并训练许多这样的神经网络可以得到一个集成模型。这个集成模型中各个模型成员都会给出一个预测结果，而为了得到最终结果，研究者和实践者通常的做法是求取这些结果的平均。这种做法通常能取得更好的表现。此外，集成模型中各个模型成员的预测结果还可用于估计不确定度。但训练集成模型也有缺点，即计算成本的预算很高，因此人们通常会避免这样操作。但是，如果神经网络的宽度变得无限大，就可以用一个高斯过程来描述这个集成模型了，而且这个高斯过程的均值和方差都可在训练过程中计算得到。

谷歌表示，使用 Neural Tangents，只需 5 行代码就能一步到位地构建并训练这种无限宽度网络的集成模型！下面展示了所得到的训练过程，而且谷歌还发布了一个可以尝试这一实验的交互式 Colab 笔记：

https://colab.sandbox.google.com/github/google/neural-tangents/blob/master/notebooks/neural_tangents_cookbook.ipynb

尽管事实上这种无限宽度的集成模型受一个简单的闭式表达式的控制，但它却能表现出与有限宽度集成模型的一致性。而且因为这种无限宽度的集成模型是一个高斯过程，所以它自然就能提供闭式的不确定度估计（上图中填充颜色的区域）。这些不确定度估计与在训练有限网络的许多不同副本（虚线）时所观察到的预测变化范围是紧密一致的。

上面的例子表明，无限宽度神经网络有习得并获取训练动态的能力。不仅如此，使用 Neural Tangents 构建的网络适用于任意可应用常规神经网络解决的问题。谷歌又给出了一个例子。下面比较了使用 CIFAR-10 数据集训练的三个不同的无限宽度神经网络架构。Neural Tangents 表现出色：不管是用梯度下降还是用完全贝叶斯的推理，都能以闭式形式评估无限宽残差网络等高度复杂的模型所组成的集成模型。而如果用有限宽度模型，这个任务根本无法解决。