- 博客地址：https://rajatvd.github.io/NTK/
- 文章动图地址：https://github.com/rajatvd/NTK
- 神经正切核相关论文地址：https://arxiv.org/abs/1806.07572
这张动图展示了神经网络的训练动态过程。阅读本文，你将了解其背后原理。
近期大量理论深度学习论文从无限宽度极限（infinite-width limit）入手分析神经网络。首先，无限宽度极限可能看起来不切实际，甚至没有研究意义。但是，无限宽度极限下的神经网络可以简化为具备神经正切核（neural tangent kernel，NTK）的线性模型。这样我们就可以很容易地研究梯度下降了。起初它看起来很有前景，但是实验结果证明，无限宽度极限下的神经网络效果不如实际的超参数网络。不过，这仍然为神经网络训练的某些方面提供了理论视角，值得我们继续探索。
此外，神经正切核可在更宽泛的标准下出现（这取决于模型的规模），且无需无限宽度极限条件。本文简单直观地介绍了该理论，它可作为梯度下降收敛性证明。为方便可视化，本文使用简单的一维示例，帮助读者更轻松地理解该理论。
设置
我们首先来看以下一个非常简单的例子，理解该示例之后再将其泛化。我们看到一个一维输入和一维输出网络。这个简单的 relu 网络有两个隐藏层，宽度为 m。
我们将使用这个简单的一维模型进行初始试验。
我们将该网络函数寻求随机初始化的过程画出来：
100 个经过随机初始化的 relu 网络，该网络有两个隐藏层，宽度为 100。
我们会发现这些函数等价于无限宽度极限下的高斯过程。
我们先来了解一些符号：
该神经网络函数为 f(x,w)，其中 x 为输入，w 为权重组合向量（大小为 p）。
在这个一维示例中，数据集为点 (x,y) 的集合。假设我们有 N 个数据点，则数据集为：
为了学习该网络，我们再使用一个简单的方法：基于最小二乘损失执行* *full-batch 梯度下降。损失函数为：
我们可以用一些向量符号简化该公式。
首先，将所有输出数据集的值 y¯_i 堆叠为一个大小为 N 的向量 y¯；
类似地，将模型对每个输入的输出 f(x¯_i,w) 堆叠为一个预测向量 y(w)∈R^N，于是就有了 y(w)_i=f(x¯_i,w)。这相当于将神经网络函数 f(⋅,w) 看作函数空间中的一个向量。
因此，上述损失函数可简化为：
现在，我们不用改变数据集的大小 N，它在损失函数公式中是一个不必要的常量。因此，我们可以在不影响任何结果的前提下删除它，这样上述公式会看起来更加简洁（为方便求导，公式保留了 1/2）。
该图可以帮助读者厘清向量符号表示。我们对整个数据集高效执行向量化，在这个简单示例中数据集的大小为 2（即图中的两个蓝色点）。
现在，训练网络即利用梯度下降最小化损失函数。
基于平方损失利用梯度下降训练 100 个 ReLU 网络。
将训练过程中两个隐藏层之间的权重（m×m 矩阵）做成动图，我们将看到一些有趣的事：
训练过程中权重矩阵的动图。你可能发现了，有些变化只发生在宽度小于 10 的情况，而宽度大于 10 时权重看起来是完全静止的。
上面三幅图都是动图，但宽度较大时，权重竟然没有变化！看起来这些模型太「懒惰」了！我们来定量地看一下，权重向量与初始化权重之间的相对变化：
当权重值翻倍时，该相对变化与隐藏层大小的关系也远了一倍。现在，我们再来绘制训练得到的网络：
损失曲线以及训练过程中权重的变化幅度。权重范数（weight norm）是将模型所有参数堆叠为一个向量而计算得到的。此外，所有其他超参数（如学习率）均保持不变。
当隐藏层宽度较大时，权重的变化幅度并不大。
泰勒展开
对网络函数及其初始化周围的权重执行泰勒展开：
我们将这个非线性网络函数转化成了简单的权重线性函数。使用更准确的向量符号表示，我们可以将其重写为：
矩阵大小也很清晰：
该图清晰地展示了所有向量和矩阵的大小。
这里很多项（如初始输出 y(w_0) 和模型雅可比矩阵 ∇_w y(w_0)）都是常量。我们主要来看 w 的依赖关系，即该近似是权重的线性模型，因此最小化最小二乘损失被简化为线性回归！不过需要注意的是，模型函数仍然是非线性的，因为寻找模型梯度绝非线性运算。事实上，这只是一个使用特征图 ϕ(x) 的线性模型，其在初始化时的梯度向量为：
该特征图自然而然地导出输入上的核，即神经正切核。之后在介绍梯度流时我们会深入探讨神经正切核，现在先从理论角度证明该线性近似的合理性。
什么时候该近似是准确的？
该线性化模型非常适合分析，但这只出现在它确实是非线性模型的准确近似的情况下。在该近似运行良好时，我们尝试推导出其定量标准。这基于 Chizat 和 Bach 的论文《On Lazy Training in Differentiable Programming》（「懒惰」模式是线性近似的别称）进行。我要做的推导非常直观，包含少量 handwaving 以及矩阵和算子范数。
首先，我们利用梯度下降优化模型，使之变为线性模型：
使用足够小的学习率 η，我们知道在运行梯度下降时，损失通常会下降。因此，模型输出 y(w_n) 通常接近真值标签 y¯。基于此，我们将网络变化的边界设定为：
现在，我们可以利用一阶近似量化网络在参数空间中移动的「距离」：
雅可比矩阵的范数 ‖∇_w y(w_0)‖ 是算子范数。如果你不熟悉算子范数，可以把它想像成衡量矩阵在向量上的最大伸展的指标（或者更简单些，将它作为衡量矩阵「大小」的指标）。
为了得到雅可比矩阵的变化，我们可以利用距离 d 和模型输出与权重的海森矩阵 ∇_w^2 y(w_0)。此处，海森矩阵是一个秩为 3 的张量（仍是线性算子），你可以将它看作将每个输出的海森矩阵堆叠在一起获得三维海森「立方体」。将其范数作为衡量海森矩阵大小或雅可比矩阵变化率的指标。
当我们想让模型成为线性时，我们在意的不只是雅可比矩阵的变化，还有相对变化。具体来说，我们想让雅可比矩阵的相对变化尽可能小（≪1）：
我们将蓝色项叫做 κ(w_0)，条件 κ(w_0)≪1 可被直观地总结为：
为得到 y 的变化 ‖(y(w_0)−y¯)‖，w 需要变化的量将导致雅可比矩阵∇_w y(w) 的细微变化。
这意味着该模型非常接近其线性近似。
现在，我们需要了解神经网络隐藏层宽度 m 对 κ(w_0) 变化的影响。我们发现随着 m 趋向无穷大，k 趋向于 0。这里我们需要谨慎一些，由于初始化权重 w_0 是随机的，κ(w_0) 也是随机变量。也就是说，当 k 趋向于 0 时，其期望值也趋向于 0。这仅适用于权重得到恰当初始化的情况。具体来说，它们必须是独立的零均值高斯随机变量，方差与输入层大小成反比（这叫做 LeCun initialization）。该结果解释了，神经网络宽度越大权重变化越小的原因。
为什么这一结果是通用的？理解这一点比较复杂，因此我将在仅具备一个隐藏层的更简单示例中，展示推导过程。
直观解释是：较大的网络宽度意味着影响输出的神经元很多。这些神经元权重发生一点细微变化都可能导致输出出现非常大的改变，因此这些神经元需要移动非常小的距离来拟合数据。如果权重移动幅度较小，则线性近似更加准确。随着网络宽度增加，神经元变化幅度降低，模型越来越接近其线性近似。如果你理解了这一直观解释，你可以跳过接下来的证明部分，直接看下一章节。
在原文中点击左侧三角图标，即可查看相关证明，此处不再赘述。
缩放输出
我们再来看一下 κ(w_0) 的公式：
我将模型输出乘以因子 α：
这就是模型重缩放。‖(y(w_0)−y¯)‖ 项很繁琐，我们可以假设模型在初始化时的输出总是为 0，即 y(w_0)=0，从而去除该项。（我们可以在初始化时减去输出的副本，使模型在初始化时的输出总是为 0。）
我们令 α→∞，使模型成为线性（即 κ(w_0)→0）。因此，《On Lazy Training in Differentiable Programming》的作者将 κ(w_0) 称为模型的逆相对缩放（inverse relative scale）。
有一件重要的事需要注意：这适用于任意非线性模型，且并不特定于神经网络（尽管需要两次微分）。因此，为了方便可视化，我制作了一个一维示例模型。该模型有一个权重 w，初始化权重为 w_0=0.4。
它满足初始化处输出为 0 的要求。要想看 α 对模型线性情况的影响，我们只需查看函数在特定值 x 处的值（假设 x=1.5）。通过不断改变 α 的值，我们可以得到如下可视化图：
在特定 x 值处评估该示例模型，不断更改 α 值。随着缩放越来越大，线性（即该案例中的切线）逐渐接近真实函数。
我们还可以查看单个数据点在这两个模型上的损失表面。线性化损失是一条漂亮的抛物线，我们希望随着 α 值的增加，实际损失逐渐接近这条抛物线。
归一化损失。只需除以α^2 即可完成归一化，这样实际损失不会随着 α 发生变化。这两幅图展示的是同一件事，只不过下方的图在 w=0.4 周围进行了放大，以便我们更清楚地查看微小的差异。
从中，我们可以得到一些重要的观察结果：
线性化损失不会像我们期待的那样接近实际损失；
两个损失表面的极小值逐渐接近，更重要的是，它们逐渐靠近初始化值。这与上文中的观察「模型训练过程中权重很少发生改变」相符。
梯度流
我们解决了神经网络和更通用的非线性模型何时被其线性模型准确逼近的问题。现在，我们来看看梯度下降过程中的训练动态。
重写该公式，得到：
左侧的项看起来像是导数的有限差分近似。这个公式就变成了对微分方程进行差分方程近似。令学习率无限小，我们可以看到权重向量随着时间的变化，微分方程如下：
这就是梯度流（gradient flow）。本质上，它是标准梯度下降的连续时间变化。重点在于，梯度流在参数空间中的轨迹近似学习率足够小时该微分方程的解的轨迹。为了简化符号表示，本文用点标注时间导数：
此外，由于时间变量可以从语境中推断出来，我们去掉该变量。使用梯度替代损失，得到：
现在，我们可以利用链式法则基于该梯度流推导出模型输出 y(w) 的动态过程（这基本上就是函数空间中的动态变化）：
红色项 ∇y(w)^T∇y(w) 就是神经正切核（NTK），我们用符号 H(w) 表示它。
回到泰勒展开那一节，我们看到线性化模型的特征图 ϕ(x)=∇_w f(x,w_0)。该特征图对应的核矩阵是通过对所有数据点的特征图取成对内积得到的。这就是 H(w_0)！
初始化处的神经正切核包含数据点特征图之间的成对内积。注意，这也可以解释为数据点的外积（x¯_i 是数据集中的输入）。
如果模型接近其线性近似（κ(w_0)≪1），则模型输出的雅可比矩阵不会随着训练过程发生变化。也就是说：
这就是核机制（kernel regime），因为正切核在训练过程中始终为常量。现在，训练动态过程被简化为简单的线性常微分方程（ODE）：
很明显，y(w)=y¯ 是该常微分方程的平衡点，它对应于训练损失为 0 的情况，这正是我们想要的。我们可以定义 u=y(w)−y¯，来改变状态变量，从而去除 y¯。则梯度流公式被简化为：
该 ODE 的解由矩阵指数运算得出：
由于模型过参数化（p>n），则 NTK ∇y(w_0)^T∇y(w_0) 总是正定的（忽略导致∇y(w_0) 不具备满列秩的任意数据集退化）。对正定 NTK 执行谱分解后，我们将梯度流的轨迹分解为独立的一维组件（特征向量），它以正比于对应特征值的速率衰减。其关键在于，它们全都衰减（因为所有特征值均为正），这意味着梯度流总是收敛至训练损失为 0 的平衡点。
通过这一系列论证，我们证明了，对于任意非线性模型，只要它接近其线性近似，其梯度下降就能收敛至 0 训练损失。这是近期证明梯度下降达到 0 训练损失的大多数论文中大部分证明的本质。
核机制
梯度流的数学原理有点难懂，现在我们回到数据集中有两个样本的示例。对于两个数据点，NTK 是 2x2 正定矩阵。我们可以将该矩阵可视化为二维平面中的椭圆，其中长轴和短轴是特征值，其长度与特征值成反比。基于此，我们可以在此二维平面中看到梯度下降（逼近梯度流）导致的 y(w(t)) 的轨迹。如下所示：
不同模型大小在训练过程中的神经正切核（从不同初始化开始）。我对不同 α 的梯度流执行了归一化，以使时间尺度和正切核规模相匹配。我们可以清晰地看到 α 越大（粉色线），NTK 的变化幅度越小，而 α 较小时（红色线），NTK 在模型训练过程中的变化很大。下侧图展示了模型函数在训练过程中的变化。
我特意将椭圆的中心设置为目标数据 y¯，以使其轨迹接近中心。这样，你可以看到椭圆短轴（较大特征值）对应的组件比长轴（较小特征值）收敛得更快；
α 的较小值也会收敛至 0 训练损失，不过我们对此没有任何理论依据。我们的证明仅对 α 足够大能使椭圆在训练过程中保持不变时有效（核机制）；
最小的 α 值收敛得很快，快到动图中几乎看不出来；
我在模型运行中使用的是两层网络，其宽度为 100。我还在初始化处减去了其副本，使网络输出为 0（正因如此，所有轨迹均从原点开始）。尽管我们无法看到该网络有 10000 个权重时，权重在训练过程中的变化，但是我们可以看到网络输出，如上图所示。
重点在于：对于较大的 α 值，存在核机制，而对于较小 α 值的行为则缺乏理论依据。
此外，还有一点需要注意：
这些图中的 NTK 机制仅适用于较大的 α 值。我们还可以将宽度扩展到无限大来使用该机制，而不仅是缩放输出。但是，该案例提供了一个特别的收获。注意，初始化处的 NTK 是不同的，因为初始化是随机的。但是，随着我们将隐藏层宽度扩展至无限，初始化处的 NTK 变得确定了！也就是说，对于给定深度和激活函数，存在一个固定的 NTK。实际上，该 NTK 可以得到确切计算，高效模拟无限宽度神经网络（参见论文《On Exact Computation in an Infinitely Wide Net》，此处不再赘述）。
泛化
目前，我们仅探讨了训练数据的情况。那么测试集呢？
首先，在核机制中训练模型等价于解决线性系统。过参数化（p>n）意味着该线性系统是欠定的（under-determined），因此它具备无限个解。由于我们使用梯度下降来解决该系统，那么它隐性地偏向于极小范数解。即，梯度下降选择具备极小 ‖w‖_2 的解（前提是，我们从具备低范数的初始化开始）。
我们现在考虑将 y 空间扩展为无限维度的函数空间，则 w 空间中的极小 ℓ2 范数被转换为选择一个函数，使特定函数范数最小化。决定该范数的是什么呢？描述该线性问题的核，即 NTK。这可以解读为一种正则化形式，你可以使用它谈论泛化。本文不再赘述。
结论
NTK 理论很伟大，而上述可视化图证明了神经网络不止 NTK 机制。实验证明，实际成功的神经网络并非在 NTK 机制中运行。实际上，即使最优秀的线性化神经网络在标准基准数据集（如 MNIST 和 CIFAR）上的性能也比实际成功的神经网络差了大约 7%（不过近期发展或许弥补了这一差距，参见论文《Enhanced Convolutional Neural Tangent Kernels》）。
但是，NTK 理论不是理解神经网络训练动态过程的唯一理论。还有一系列工作利用最优传输和平均场理论描述训练动态过程，这些案例中通过非线性 PDE 实现的模型缩放并不大（κ∼1），但是据我所知，这些理论都无法扩展至具备一个隐藏层的网络。
不过，这些 NTK 结果仍然非常有趣，它们提供了探索神经网络学习的新思路。尝试理解核在训练过程中的变化对于寻找更好的神经网络理论大有裨益。
