对神经切线核的初步理解
【摘 要】 想理解神经网络的训练动态过程,不妨从「神经切线核」入手。那么什么是神经切线核,核机制如何运行?就读于印度理工学院马德拉斯分校电气工程系的 Rajat Vadiraj Dwaraknath 撰文介绍了这一概念。
- 博客地址:https://rajatvd.github.io/NTK/
- 文章动图地址:https://github.com/rajatvd/NTK
- 神经切线核相关论文地址:https://arxiv.org/abs/1806.07572
1 引言
图 1: _这张动图展示了神经网络的训练动态过程。
最近一系列关于理论深度学习的论文讨论了在 无限宽 极限下分析神经网络的问题。这个极限最初似乎不切实际,甚至研究起来毫无意义;但事实证明,此时的神经网络可以简化为具有神经切线核的 线性模型,这使得梯度下降问题可以更好被研究。此研究虽然看起来很有希望,但实证结果表明,此时的神经网络性能并不比实际的超参数网络更好。无论如何,这仍然为神经网络训练的某些方面提供了理论见解,因此值得研究。 此外,神经切线核可在更宽泛的情况下出现,而无需无限宽极限的条件。
本文简单直观地介绍了神经切线核理论,它可作为梯度下降收敛性的证明。为方便可视化,本文使用了简单的一维示例,以帮助读者更轻松地理解。
2 模型设置
首先来看一个只有一维输入和一维输出的简单例子,在理解该示例之后我们再将其泛化。这个简单的 relu 网络有两个隐藏层,层宽为 。
图 2: 我们将使用这个简单的一维模型进行初始试验。
我们将该网络函数寻求随机初始化的过程画出来:
图 3: 100 个经过随机初始化的 relu 网络,该网络有两个隐藏层,宽为 100。
已经证明:当隐藏层宽度取到无穷大时,这些函数等同于从高斯过程中抽取的样本。
我们先来了解一些符号:
- 神经网络函数为 ,其中 是输入, 是权重的组合向量 (比如大小 )。
- 在这个一维示例中,数据集将只是点 。假设有 个数据点,则数据集是
为了学习网络,我们再次采取简单方法: 对最小二乘损失进行全批梯度下降训练。现在的损失为:
我们可以进一步使用向量符号来简化它:
- 首先,可以将所有输出数据值 堆叠成一个大小为 的向量,并记为 。
- 类似地,将每个输入的所有可能模型输出 堆叠成一个 预测向量,记为 。我们有 。这在风格上类似于将神经网络函数 视为 函数空间 中的一个向量。
所以,损失可以简化为:
现在,我们不会在任何地方更改数据集大小 ,它是损失表达式中不重要的常量。因此,可以在不影响进一步结果的情况下将其隐藏,同时使代数看起来不那么混乱(我们将保留 ,因为它会使导数表达更方便)。
图 4: 此图可以帮助读者厘清向量符号。我们对整个数据集执行向量化,在这个简单示例中数据集的大小为 2(即图中的两个蓝色点)。_
现在,训练网络即利用梯度下降最小化损失函数。
图 5: 基于平方损失利用梯度下降训练 100 个 ReLU 网络。
将训练过程中两个隐藏层之间的权重( 矩阵)做成动图,我们将看到一些有趣的事:
图 6: 训练过程中权重矩阵的动图。变化只显著发生在宽度小于 的网络中,而宽度大于 时权重看起来几乎完全静止。
上面三幅图实际上都是动图,但是对于更大的宽度,神经网络权重几乎不会改变!我们可以在初始化之后,通过观察权重向量范数的 相对变化,定量地给出分析:
当权重值翻倍时,无论隐藏层大小如何变化,相对变化都将是 。现在,如果我们为上面训练的神经网络绘制下图:
图 7: 损失曲线以及训练过程中权重的变化幅度。权重范数(weight norm)是将模型所有参数堆叠为一个向量而计算得到的。此外,所有其他超参数(如学习率)均保持不变。
相对变化情况表明,对于更大的隐层宽度,权重不会改变太多。
3 泰勒展开
我们可以在 初始化权重 周边,对神经网络函数做泰勒展开:
现在我们将一个非线性神经网络函数转化成了关于权重的简单线性函数。使用向量符号表示,我们可以将其重写为:
矩阵大小也很清晰:
图 8: 该图清晰地展示了所有向量和矩阵的大小。
这里的很多项都是常数,如初始输出 和模型的 Jacobian 矩阵 。我们重点关注对 的依赖,因为近似是权重的 线性模型,所以最小化最小二乘损失被简化为简单的 线性回归!但请注意:模型函数本身关于输入仍然是 非线性的。事实上,这只是一个使用了 特征映射 的线性模型,其初始化时的梯度向量:
该特征映射自然而然地导出输入上的核,我们称之为神经切线核。我们将在后面梯度流部分深入探讨神经切线核,现在先从理论角度证明该线性近似的合理性。
4 什么时候该近似是准确的?
线性化的模型非常适合分析,前提是它确实是非线性模型的精确近似。让我们尝试推导出一个定量准则,以判断该近似何时运行良好。这改自 Chizat 和 Bach 的论文 Lazy Training in Differentiable Programming( 惰性机制 是此线性近似成立时的另一个名称)。您可能需要打起精神来,因为我介绍的推导是直觉的,并且涉及到一些复杂的矩阵和运算符范数。
首先,我们只关心线性的模型,因为要使用梯度下降对其进行优化:
使用足够小的学习率 ,损失总是随着我们运行梯度下降而减少。因此,模型输出 总是更接近真实标签 。使用它,我们可以按如下方式给出净变化的界:
现在,我们可以用一阶近似来量化在参数空间中移动的 “距离”,
Jacobian 范数 是一个算子范数。如果您不熟悉算子范数,可以将其视为一种用于度量矩阵对向量所能实施的最大拉伸量的测度(或者更简单,只需将其视为一种对矩阵 “幅度” 的度量)。
为了获得 Jacobian 矩阵中的变化,我们可以将 “距离 ” 与 “模型输出相对于权重的 Hessian 矩阵 ” 一起使用。在我们的例子中,这个 Hessian 矩阵是一个 3 阶张量(仍然是一个线性算子),你可以把它看作是将每个输出的 Hessian 矩阵堆叠在一起,获得了一个 3-d 的 Hessian “长方体”。同样,将其范数视为 Hessian 矩阵的 “幅度” 或 Jacobian 矩阵变化率的某种度量。
当我们希望模型是线性的时候,我们真正关心的不仅仅是其雅可比行列式的变化,而是 相对变化。具体来说,我们希望这个相对变化 尽可能小():
我们将蓝色的量为 ,条件 可被直观地总结为:
权重 的变化量,为了能够在 上产生 , 的变化量,使得雅可比矩阵 产生细微变化
这意味着该模型非常 “接近其线性近似”。
我们现在需要了解 如何随着神经网络的隐藏层宽度 而变化。我们发现,随着宽度 , 。这里需要小心: 是一个随机变量,因为初始化 本身就是随机的。所以当 时,意味着其期望值趋于 。这仅适用于权重被正确初始化的情况。具体来说,它们需要是独立的零均值高斯随机变量,其方差与输入层的大小成反比(这种初始化被称为 LeCun 初始化)。此结果解释了为什么在训练更大宽度的神经网络时权重变化会比较小。
理解为何此结果的通用性是相当复杂的,所以在此将针对只有 一个 隐藏层的简单情况提供推导。
直观解释是:较大的神经网络宽度意味着更多会影响输出的神经元。所有这些神经元权重的微小变化都会导致输出发生非常大的变化,因此这些神经元需要移动非常小的距离才能拟合数据。权重移动幅度越小,线性近似越准确。随着网络宽度增加,神经元变化幅度减小,模型越来越接近其线性近似。如果您乐于相信这种直觉,则可以跳过证明并转到下一节。
在原文中点击左侧三角图标,即可查看相关证明,此处不再赘述。
4.1 梯度
4.2 Heissan
4.3 组合在一起
5 缩放输出
让我们再看看 的表达式:
我要做一件非常简单的事情,只需 将模型的输出乘以某个因子 。
这实际上只是 模型重缩放。 项很烦人,我们可以通过假设模型 在初始化时始终输出 来去除该项,即 (我们可以在初始化时减去输出的副本,使模型在初始化时的输出总是为 )。
我们可以通过简单地令 $\color{red} \alpha \alpha \to \infty $,使模型成为线性的! 《On Lazy Training in Differentiable Programming》的作者将 称为模型的 逆相对缩放。
需要注意的一件重要事情是:这适用于 任何非线性模型,并且不特定于神经网络(尽管需要二次可微)。因此,为了形象化这一点,我制作了一个一维示例模型。该模型的 单权重 初始化为 。
它还满足初始化时输出为 的需求。要查看 对模型线性度的影响,我们可以只查看函数在_一个特定值_ 处的值,比如 。改变 ,我们得到:
_在特定 x 值处评估该示例模型,不断更改 α 值。_随着缩放越来越大,线性(即该案例中的切线)逐渐接近真实函数。
我们也可以观察一下这两个模型单数据点的损失表面。线性化损失是一条漂亮的抛物线,我们期望实际的损失随着 color{red}\α 的增加而更接近于此抛物线。
_归一化损失。__只需除以α^2 即可完成归一化,这样实际损失不会随着 α 发生变化。_这两幅图展示的是同一件事,只不过下方的图在 w=0.4 周围进行了放大,以便我们更清楚地查看微小的差异。
从中,我们可以得到一些重要的观察结果:
-
线性化损失不会像我们期待的那样接近实际损失;
-
两个损失表面的极小值逐渐接近,更重要的是,它们逐渐靠近初始化值。这与上文中的观察「模型训练过程中权重很少发生改变」相符。
梯度流
我们解决了神经网络(和更广义的非线性模型)何时能够被其线性模型准确逼近的问题。现在,来看看梯度下降过程中的训练动态。
重写该公式,得到:
左侧的项看起来像是导数的有限差分近似。这个公式就变成了对微分方程进行差分方程近似。令学习率无限小,我们可以看到权重向量随着时间的变化,微分方程如下:
这就是梯度流(gradient flow)。本质上,它是标准梯度下降的连续时间变化。重点在于,梯度流在参数空间中的轨迹近似学习率足够小时该微分方程的解的轨迹。为了简化符号表示,本文用点 标注时间导数,梯度流为:
此外,由于时间变量可以从语境中推断出来,我们去掉该变量。使用梯度替代损失,得到:
现在,我们可以使用链式法则导出由该梯度流引起的模型输出 的动态(这基本上是 函数空间中的动态):
红色的量 称为神经切线核(简称NTK)。让我们给它一个符号 。
回到泰勒展开部分,我们看到线性化模型的特征映射 。这个特征映射对应的核矩阵是通过对所有数据点的特征映射取 逐对内积(pairwise inner products) 得到的。这正是 !
_初始化处的神经切线核包含数据点特征映射之间的成对内积。_注意,这也可以解释为数据点的外积(x¯_i 是数据集中的输入)。
如果模型接近其线性近似值 (),则模型输出的雅可比矩阵不会随着训练的进行而改变。换句话说,
这被称为 核机制(kernel regime),因为切线内核在训练期间保持不变。训练动力学现在简化为一个非常简单的线性常微分方程:
显然, 是这个 ODE 的平衡点,它对应于 的训练损失,这正是我们想要的。我们可以通过定义 来更改状态变量以去除 。梯度流简化为:
此 ODE 的解由矩阵指数给出:
因为我们的模型是 过度参数化的 (),NTK 总是 正定的(忽略数据集中会导致 不具备满列秩的任何退化)。通过对正定 NTK 执行谱分解,我们可以将梯度流的迹解耦为独立的一维分量(特征向量),这些分量以与相应特征值成比例的速率衰减。关键是它们都 衰减(因为所有特征值都是正的),这意味着梯度流 总是收敛 到训练损失为 的平衡点。
通过这一系列论证,我们已经证明:对于任意非线性模型,只要它接近其线性近似,其梯度下降就能收敛至 训练损失(这可以通过将尺度 来实现) )!这是近期证明梯度下降达到 训练损失的大多数论文中证明过程的核心。
核机制
梯度流的数学原理有点难以理解,所以让我们回到数据集中有两个样本的示例。对于 个数据点,NTK 只是一个 正定矩阵。我们可以将此矩阵可视化为二维平面中的椭圆,其中长轴和短轴是特征向量,其长度与特征值的平方根成反比。椭圆基本上是时变二次方程 。
除此之外,我们还可以在这个二维平面上可视化由梯度下降(近似梯度流)引起的 的轨迹。这个轨迹相当于对同一个时变或_瞬时_二次方进行梯度下降。这是它的样子:
不同模型大小在训练过程中的神经切线核(从不同初始化开始)。我对不同 α 的梯度流执行了归一化,以使时间尺度和切线核规模相匹配。我们可以清晰地看到 α 越大(粉色线),NTK 的变化幅度越小,而 α 较小时(红色线),NTK 在模型训练过程中的变化很大。下侧图展示了模型函数在训练过程中的变化。
- 我故意将椭圆集中在目标数据 上,以便轨迹接近中心。这样可以看到椭圆的短轴(特征值大)对应的分量比长轴(特征值小)收敛得快。
- 的较小值也收敛到 0 训练损失(这称为“深度”或“丰富”状态)。不过,我们对此没有任何理论依据。我们的证明仅在 大到足以使椭圆在整个训练过程中保持不变(核机制)时才有效。
- 最小的 收敛得如此之快以至于在动图中几乎无法察觉到。
- 模型运行中使用了宽度为 的 2-隐藏层网络。我还通过在初始化时减去一个副本来将这些网络的输出归零(这就是为什么所有轨迹都从原点开始的原因)。虽然我们无法想象这个网络的 多个权重在我们训练时如何移动,但我们可以想象网络输出,这就是这个动图的作用。
主要的收获是: 对于较大的 值,存在核机制,而对于较小 值的行为则缺乏理论依据。
此外,还有一点需要注意:
这些图中的 NTK 机制仅适用于较大的 值。我们还可以将宽扩展到无限大来使用该机制,而不仅是缩放输出。但是,该案例提供了一个特别的收获。注意,初始化处的 NTK 是不同的,因为初始化是随机的。但是,随着我们将隐藏层宽扩展至无限,初始化处的 NTK 变得确定了!也就是说,对于给定深度和激活函数,存在一个固定的 NTK。实际上,该 NTK 可以得到确切计算,高效模拟无限宽神经网络(参见论文《On Exact Computation in an Infinitely Wide Net》,此处不再赘述)。
泛化
到目前为止,我们只讨论了训练数据会发生什么。但是测试集样本之间的点呢?
首先,在核机制中训练模型等同于求解线性系统。过度参数化 () 只是意味着这个线性系统是 欠定的 所以它有无限个解。由于我们使用梯度下降来解决这个系统,它隐含地偏向最小范数的解。也就是说,梯度下降选择具有最小 的解(只要我们从低范数的初始化开始)。
如果我们现在考虑将 空间扩展到一个无限维度的 函数空间,那么 空间中的这个最小 范数转化为选择一个能够最小化特定函数范式的函数。那么决定这个范式的是什么呢?描述线性问题的核,在这种情况下是 NTK。这可以解释为一种正则化形式,你可以用它来谈论泛化。我不会深入研究更多细节,因为这篇文章的主要目的是了解优化结果,而不是泛化。如果您想了解更多相关信息,我建议您查找有关_再生核 Hilbert 空间_ 和核机器泛化性质的内容。
结论
NTK 理论很棒,但是上面的可视化显示神经网络不仅仅是 NTK 核机制。实验表明,实际上成功的神经网络并非 在核机制中运行。实际上,即使最优秀的线性化神经网络在标准基准数据集(如 MNIST 和 CIFAR)上的性能也比实际成功的神经网络差了大约 7%,不过近期发展或许弥补了这一差距,参见论文《Enhanced Convolutional Neural Tangent Kernels》。
然而,这种 NTK 机制理论并不是理解神经网络训练动力学的唯一理论。还有一系列工作(请参阅参考资料)使用最优传输和平均场理论中的思想来表征模型缩放比例不大()时通过非线性 PDE 的动力学,但是据我所知,这些理论并未扩展到具有单个隐藏层的网络之外。
尽管如此,这些 NTK 结果仍然非常有趣,因为它们提供了神经网络学习的新视角。在寻找更好的神经网络理论时,尝试了解核如何随着训练的进行而变化对于寻找更好的神经网络理论大有裨益。
参考文献
Related to NTK:
- This nice YouTube video by Arthur Jacot summarizes the NTK paper and also has some more visualizations which I didn’t present in this post.
- Jacot, Arthur, Franck Gabriel, and Clément Hongler. “Neural tangent kernel: Convergence and generalization in neural networks.” Advances in neural information processing systems. 2018.
- Chizat, Lenaic, and Francis Bach. “A note on lazy training in supervised differentiable programming.” arXiv preprint arXiv:1812.07956 (2018).
- Arora, Sanjeev, et al. “On exact computation with an infinitely wide neural net.” arXiv preprint arXiv:1904.11955 (2019).
- Li, Zhiyuan, et al. “Enhanced Convolutional Neural Tangent Kernels.” arXiv preprint arXiv:1911.00809 (2019).
Mean field dynamics and optimal transport point of view:
- Chizat, Lenaic, and Francis Bach. “On the global convergence of gradient descent for over-parameterized models using optimal transport.” Advances in neural information processing systems. 2018.
- Mei, Song, Theodor Misiakiewicz, and Andrea Montanari. “Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit.” arXiv preprint arXiv:1902.06015 (2019).