神经正切核

神经正切核（英語：neural tangent kernel，简称NTK）是一种核函数，用于描述深度人工神经网络在梯度下降训练过程中的动态演变。NTK将核方法的理论工具引入人工神经网络的研究之中。

通常而言，核函数是一类正半定的对称函数，衡量两个输入之间的某种相似性。而NTK则是从一个具体的神经网络中推导而来的特殊核函数。一般情况下，当神经网络的参数在训练过程中更新时，NTK也会随之变化。然而，当神经网络的层宽趋于无穷大时，NTK会收敛成为一个固定不变的函数。这揭示了宽神经网络的训练与核方法之间的对偶性：无限宽度极限下的梯度下降完全等价于使用NTK的核梯度下降。这意味着，通过梯度下降来最小化神经网络的最小二乘损失，与使用NTK的无脊核回归所得到的均值估计是一致的。这种对偶性使得我们能够用简单的闭式方程来描述宽神经网络的训练动态、泛化能力和预测结果。

NTK最早由Arthur Jacot、Franck Gabriel和Clément Hongler于2018年提出^[1]，他们利用NTK研究了全连接神经网络的收敛和泛化特性。随后的一些研究^[2][3]则将NTK的分析框架推广到了其他类型的神经网络架构中。实际上，NTK所揭示的现象并非神经网络所特有，在更一般的非线性模型中通用适当的缩放也能观察到类似的行为。^[4]

主要结论

令${\displaystyle f(x;\theta )}$表示一个通过给定神经网络计算的标量函数，其中${\displaystyle x}$是输入，${\displaystyle \theta }$是网络中的所有参数，则相应的神经正切核可以定义为^[1]

${\displaystyle \Theta (x,x';\theta )=\nabla _{\theta }f(x;\theta )\cdot \nabla _{\theta }f(x';\theta ).}$

由于NTK可以表示为经特征映射后两个输入之间的点积（此处神经网络函数的梯度充当了将输入映射到高维空间的特征映射函数），这保证了NTK具有对称性与半正定性。因此，NTK是一个有效的核函数。

现在我们考虑一个全连接神经网络，其参数${\displaystyle \theta }$从任意均值为零的分布中进行独立同分布采样。这种对函数的随机初始化，会相应地使${\displaystyle f(x;\theta )}$形成一个特定的函数分布。我们的目标便是分析该分布在初始化及整个梯度下降训练过程中的统计特性。为了直观地展现该分布，我们可以构建一个神经网络的集成（ensemble），即从${\displaystyle f(x;\theta )}$的初始分布中多次采样，并按照完全相同的流程来训练每一个神经网络实例。

初始化时，宽神经网络的集成是一个零均值的高斯过程。在基于均方误差的梯度下降训练期间，该集成则会根据NTK进行演化。收敛后的集成也是一个高斯过程，其分布的均值为无脊核回归的解，方差则在所有训练数据点上降为零。此处的神经网络是一个标量函数，其训练数据从单位圆上采样得到。

神经网络中每层的神经元数量称为该层的宽度。假设我们将每个隐藏层的宽度都设为无穷大，并使用梯度下降法（并配合一个足够小的学习率）来训练这个神经网络。在此无限宽度极限之下，神经网络会出现一些非常理想的特性：

• 在训练开始前的初始化阶段，神经网络集成是一个零均值高斯过程（GP）。^[5]这意味着函数分布是一种最大熵分布，其均值为${\displaystyle \mathbb {E} _{\theta }[f(x;\theta )]=0}$，协方差为${\displaystyle \mathbb {E} _{\theta }[f(x;\theta )f(x';\theta )]=\Sigma (x,x')}$，其中${\displaystyle \Sigma (x,x')}$由神经网络的架构所决定。换言之，神经网络函数的分布在初始化时除了一阶矩（均值）和二阶矩（协方差）之外没有其他结构。这一结论遵循中心极限定理。

• NTK具有确定性^[1][6]，即其独立于参数的随机初始化。
• NTK在整个训练过程中保持不变。^[1][6]
• 尽管每个参数在训练过程中的变化都微乎其微，但它们的共同作用却足以对网络的最终输出产生显著的改变，从而实现有效的学习。^[6]
• 正由于在训练过程中每个参数的变化都可以忽略不计，神经网络可以被线性化，即可以通过关于初始参数${\displaystyle \theta _{0}}$的一阶泰勒展开式来近似：${\displaystyle f(x;\theta _{0}+\Delta \theta )=f(x;\theta _{0})+\Delta \theta \cdot \nabla _{\theta }f(x;\theta _{0})}$。^[6]（但神经网络对于输入仍保持非线性的特征）
• 神经网络的训练动态等价于使用NTK作为核函数的核梯度下降。^[1]如果损失函数是均方误差，那么${\displaystyle f(x;\theta )}$的最终分布仍然是高斯过程，但均值和协方差则会发生变化。^[1][6]具体而言，其均值会收敛到以NTK为核的无脊核回归所给出的解，而协方差可以用NTK和初始协方差来表示。同时可以证明，集成的方差在所有训练数据点处都会降为零。换句话说，无论如何随机初始化，神经网络最终总能完美地拟合所有训练数据。

从物理学的角度来看，NTK可以被视为一种哈密顿量，因为当神经网络通过无穷小的步长（即连续时间极限）进行梯度下降训练时，它描述了网络中可观测量随时间演变的过程。^[7]

应用

无脊核回归和核梯度下降

核方法是一类机器学习算法，其特性是仅使用输入点之间的成对关系进行计算，而不依赖于输入的具体值。所有这些成对关系都可以完全由一个核函数所描述。核函数是一个对称、半正定的函数，它接收两个输入，并返回一个表示它们之间某种相似度的值。一个与此完全等效的定义是，存在某种特征映射${\displaystyle {\mathbf {x}}\mapsto \psi ({\mathbf {x}})}$，使得核函数可以表示为经映射后输入值的点积

${\displaystyle K({\mathbf {x} },{\mathbf {x} }')=\psi ({\mathbf {x} })\cdot \psi ({\mathbf {x} }').}$

核方法的性质取决于核函数的选择。我们以经典的线性回归为例来说明。假设有${\displaystyle N}$个训练样本${\displaystyle ({\mathbf {x}}_{i},y_{i})}$，它们都由一个线性函数${\displaystyle y^{*}({\mathbf {x}})={\mathbf {w}}^{*}\cdot {\mathbf {x}}}$生成，而我们的任务则是通过这些样本来学习权重${\displaystyle \mathbf {\hat {w}} }$，使其尽可能地接近真实权重${\displaystyle {\mathbf {w}}^{*}}$。为实现这一目标，我们可以最小化模型预测值与真实训练值之间的均方误差：

${\displaystyle {\mathbf {\hat {w}} }=\arg \min _{\mathbf {w} }{\frac {1}{N}}\sum _{i=0}^{N}||y^{*}({\mathbf {x} }_{i})-{\mathbf {w} }\cdot {\mathbf {x} }_{i}||^{2}}$

对于上述的最小化问题，${\displaystyle \mathbf {\hat {w}} }$存在显式解：${\displaystyle {\mathbf {\hat {w}}}=({\mathbf {X}}{\mathbf {X}}^{T})^{-1}{\mathbf {X}}{\mathbf {y}}}$。其中，${\displaystyle {\mathbf {X}}}$是由所有训练输入作为列向量构成的矩阵，而${\displaystyle {\mathbf {y}}}$则是所有训练输出构成的向量。求得${\displaystyle \mathbf {\hat {w}} }$后，便可以对任意新的输入做出预测：

${\displaystyle {\hat {y}}({\mathbf {x}})={\mathbf {\hat {w}}}\cdot {\mathbf {x}}}$

我们可以将上述结果改写为${\displaystyle {\hat {y}}({\mathbf {x}})=({\mathbf {x}}^{T}{\mathbf {X}})({\mathbf {X}}^{T}{\mathbf {X}})^{-1}{\mathbf {y}}}$。^[8]在这种形式下，问题的解可以完全通过输入之间的点积来表达。这一发现意味着我们得以将线性回归进行推广。此时我们不再直接计算输入之间的点积，而是先用一个给定的特征映射对输入进行变换，然后再计算变换后的点积。如前文所述，这可以通过一个核函数${\displaystyle K({\mathbf {x}},{\mathbf {x}}')}$来实现。于是，我们能够得到无脊核回归的预测公式：

$${\displaystyle {\hat {y}}({\mathbf {x}})=K({\mathbf {x}},{\mathbf {X}})\;K({\mathbf {X}},{\mathbf {X}})^{-1}\;{\mathbf {y}}.}$$

如果核矩阵${\displaystyle K({\mathbf {X}},{\mathbf {X}})}$为奇异阵，可使用穆尔-彭罗斯广义逆来代替逆矩阵。该回归方程被称为“无脊”回归，是因为其公式中缺少了脊正则化项。

从这个角度来看，线性回归可以看作是核回归的一个特例，它对应于使用恒等特征映射${\displaystyle \psi ({\mathbf {x}})={\mathbf {x}}}$的核回归。反过来，核回归也可看作是在特征空间中的线性回归。但其通常在输入空间中是非线性的，这也是核算法的核心优势所在。

正如我们可以使用梯度下降这类等迭代优化算法来来解线性回归问题，我们同样也可以使用核梯度下降来求解核回归问题。这等价于在特征空间中使用标准的梯度下降。对于线性回归而言，如果权重向量的初始值接近于零，那么最小二乘梯度下降会收敛到最小范数解，即在所有能够完美拟合训练数据的解中欧几里得范数最小的解。类似地，核梯度下降会得到再生核希尔伯特空间（英语：Reproducing kernel Hilbert space）范数最小的解。这一现象被则为梯度下降的隐式正则化。

NTK在无限宽神经网络与核方法之间建立了严格的对应关系：当使用最小二乘损失训练时，无限宽神经网络给出的预测，其期望与以NTK作为核函数的无脊核回归所得到的预测是一致的。这表明，对于采用NTK参数化的大规模神经网络，其性能可以通过选取合适核函数的核方法加以复现。^[1][2]

过参数化、插值和泛化

在过参数化（overparametrization）的模型中，可调参数的数量大于训练样本的数量。此种情况下，模型能够记忆（完美拟合）所有的训练数据。因此，过参数化的模型会对训练数据进行插值，使其在训练集上基本实现零误差。^[9]

现代过参数化的模型尽管具有对训练集进行插值（记忆）的能力，但仍能实现了较低的泛化误差。^[9]通过研究高维核回归的泛化特性可以解释这一现象。

核回归通常被视为一种非参数机器学习算法，因为一旦选定了核函数，模型就不再有需要学习的显式参数了。另外核回归还可以看作是特征空间中的线性回归，因此“有效”的参数数量相当于特征空间的维数。因此，研究具有高维特征映射的核方法，可以让我们对严重过参数化的模型有更深的了解。

我们以泛化问题为例。根据经典统计学，记忆会导致模型拟合训练数据中的噪声信号，从而损害其在未见数据上的预测能力。为了避免这种情况，传统机器学习算法通常会引入正则化来抑制这种过拟合噪声的倾向。然而令人惊讶的是，往往严重过度参数化的现代神经网络，即便在没有显式正则化的情况下，依然表现出了优异的泛化能力。^[9][10]于是我们可以利用无脊核回归来研究过参数化神经网络的泛化特性。有研究^[11][12][13]推导出了描述高维核回归期望泛化误差的方程，而这些结果可以用来解释足够宽的神经网络在经过最小二乘损失训练后的泛化能力。

全局最小值

对于具有全局最小值的凸损失泛函${\displaystyle {\mathcal {C}}}$，如果NTK在训练期间始终保持正定，那么当${\displaystyle t\to \infty }$时，神经网络的损失${\displaystyle {\mathcal {C}}\left(f\left(\cdot ;\theta \left(t\right)\right)\right)}$就能保证收敛到该全局最小值。这种正定性已在一些情况下得到证实，从而首次证明了足够宽的神经网络在训练过程确实能够收敛到全局最小值。^{[1][14][15][16][17][18]}

扩展与局限

NTK可用于研究各种神经网络架构^[2]，如卷积神经网络（CNN）^[19]、循环神经网络(RNN) 和Transformer等。^[20]在分析这些架构时，无限宽度极限的含义也相应有所变化，通常指在保持网络层数固定的情况下增加参数的数量。以CNN为例，“宽度”对应的是卷积层的通道数。

在标准的核机制下，宽神经网络中的各个参数在训练过程中几乎不会发生变化。然而，这意味着无限宽神经网络无法进行特征学习。而特征学习被广泛认为是实际应用中深度神经网络的一个重要特性。值得注意的是，这一点并不是无限宽神经网络的普遍特征，而很大程度上是源于我们将宽度推向无穷时采取的一种特定的参数缩放方式。事实上，已有研究^{[21][22][23][24]}探索不同缩放极限下的无限宽神经网络，在这些情形下神经网络与核回归之间的对偶性不复存在，而特征学习则得以出现。此外，还有研究^[25]引入了“神经正切层级”（neural tangent hierarchy）的概念来描述有限宽度效应及其与特征学习间的联系。

Neural Tangents是一个由Google开发的免费开源Python库，可以用于计算和推断与各常见神经网络架构相对应的无限宽度NTK和神经网络高斯过程。^[26]此外，还有一个名为scikit-ntk的库提供与scikit-learn兼容的NTK工具。^[27]

细节

当通过梯度下降法来优化一个神经网络的参数${\displaystyle \theta \in \mathbb {R} ^{P}}$并最小化其经验损失时，NTK在整个训练过程中主导着网络输出函数${\displaystyle f_{\theta }}$的动态变化。

情形一：标量输出

一个输出为标量的神经网络可以被看作是由参数${\displaystyle \theta \in \mathbb {R} ^{P}}$定义的一系列函数${\displaystyle f\left(\cdot ,\theta \right):\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$。

其对应的NTK是一个核函数${\displaystyle \Theta :\mathbb {R} ^{n_{\mathrm {in} }}\times \mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$，定义为

${\displaystyle \Theta \left(x,y;\theta \right)=\sum _{p=1}^{P}\partial _{\theta _{p}}f\left(x;\theta \right)\partial _{\theta _{p}}f\left(y;\theta \right).}$

在核方法的框架下，${\displaystyle \Theta }$是与特征映射${\displaystyle \left(x\mapsto \partial _{\theta _{p}}f\left(x;\theta \right)\right)_{p=1,\ldots ,P}}$相关联的核函数。为了理解这一个核函数如何驱动神经网络的训练动态，我们可以考虑一个数据集，其中${\displaystyle \left(x_{i}\right)_{i=1,\ldots ,n}\subset \mathbb {R} ^{n_{\mathrm {in} }}}$是输入，${\displaystyle \left(z_{i}\right)_{i=1,\ldots ,n}\subset \mathbb {R} }$是对应的标量标签，${\displaystyle c:\mathbb {R} \times \mathbb {R} \to \mathbb {R} }$则是损失函数。于是，定义在${\displaystyle f:\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$上的经验损失可表示为

${\displaystyle {\mathcal {C}}\left(f\right)=\sum _{i=1}^{n}c\left(f\left(x_{i}\right),z_{i}\right).}$

当通过连续时间梯度下降来训练神经网格${\displaystyle f\left(\cdot ;\theta \right):\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$以拟合数据集（即最小化${\displaystyle {\mathcal {C}}}$）时，其参数${\displaystyle \left(\theta \left(t\right)\right)_{t\geq 0}}$的演化遵循以下常微分方程：

${\displaystyle \partial _{t}\theta \left(t\right)=-\nabla {\mathcal {C}}\left(f\left(\cdot ;\theta \right)\right).}$

在训练过程中，网络输出函数本身的演化则可以由一个以NTK表示的微分方程来描述：

${\displaystyle \partial _{t}f\left(x;\theta \left(t\right)\right)=-\sum _{i=1}^{n}\Theta \left(x,x_{i};\theta \right)\partial _{w}c\left(w,z_{i}\right){\Big |}_{w=f\left(x_{i};\theta \left(t\right)\right)}.}$

这个方程揭示了训练期间NTK是如何驱动${\displaystyle f\left(\cdot ;\theta \left(t\right)\right)}$在函数空间${\displaystyle \mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} }$中的变化的。

情形二：向量输出

当神经网络的输出是维度为${\displaystyle n_{\mathrm {out} }}$的向量时，它可以看作是由参数${\displaystyle \theta \in \mathbb {R} ^{P}}$定义的一系列函数${\displaystyle f\left(\cdot ;\theta \right):\mathbb {R} ^{n_{\mathrm {in} }}\to \mathbb {R} ^{n_{\mathrm {out} }}}$。

此种情况下，NTK${\displaystyle \Theta :\mathbb {R} ^{n_{\mathrm {in} }}\times \mathbb {R} ^{n_{\mathrm {in} }}\to {\mathcal {M}}_{n_{\mathrm {out} }}\left(\mathbb {R} \right)}$是一个具有矩阵值的核函数，将输入映射到${\displaystyle n_{\mathrm {out} }\times n_{\mathrm {out} }}$的矩阵上，定义为

${\displaystyle \Theta _{k,l}\left(x,y;\theta \right)=\sum _{p=1}^{P}\partial _{\theta _{p}}f_{k}\left(x;\theta \right)\partial _{\theta _{p}}f_{l}\left(y;\theta \right).}$

其经验风险最小化的过程与标量情形类似，主要的区别在于损失函数改为采用向量输入${\displaystyle c:\mathbb {R} ^{n_{\mathrm {out} }}\times \mathbb {R} ^{n_{\mathrm {out} }}\to \mathbb {R} }$，而其在连续时间梯度下降训练过程中函数空间的演化动态${\displaystyle f_{\theta \left(t\right)}}$也同样由NTK所主导。对应的演化方程为：

${\displaystyle \partial _{t}f_{k}\left(x;\theta \left(t\right)\right)=-\sum _{i=1}^{n}\sum _{l=1}^{n_{\mathrm {out} }}\Theta _{k,l}\left(x,x_{i};\theta \right)\partial _{w_{l}}c\left(\left(w_{1},\ldots ,w_{n_{\mathrm {out} }}\right),z_{i}\right){\Big |}_{w=f\left(x_{i};\theta \left(t\right)\right)}.}$

这一方程是情形一中标量输出方程的直接推广。

解释

在训练的每一步，每个训练数据点${\displaystyle x_{i}}$都会对任意一个输入${\displaystyle x}$所对应的网络输出${\displaystyle f\left(x;\theta \right)}$的演化产生影响。更具体地说，对于第${\displaystyle i}$个训练样本，其产生的梯度损失${\displaystyle \partial _{w}c\left(w,z_{i}\right){\big |}_{w=f\left(x_{i};\theta \right)}}$会对${\displaystyle f\left(x;\theta \right)}$的更新作出贡献，而NTK的值${\displaystyle \Theta \left(x,x_{i};\theta \right)}$则决定了该贡献的大小。在标量输出的情形下，这一过程可以用以下离散时间的梯度下降更新公式来直观地表示：

${\displaystyle f\left(x;\theta \left(t+\epsilon \right)\right)-f\left(x;\theta \left(t\right)\right)\approx \epsilon \sum _{i=1}^{n}\Theta \left(x,x_{i};\theta \left(t\right)\right)\partial _{w}c\left(w,z_{i}\right){\big |}_{w=f\left(x_{i};\theta \right)}.}$

训练中保持不变的确定性NTK

考虑一个包含全连接层${\displaystyle \ell =0,\ldots ,L}$的神经网络，其各层宽度为${\displaystyle n_{0}=n_{\mathrm {in} },n_{1},\ldots ,n_{L}=n_{\mathrm {out} }}$。该网络可以表示成多层函数的复合${\displaystyle f\left(\cdot ;\theta \right)=R_{L-1}\circ \cdots \circ R_{0}}$ ，其中每一层的${\displaystyle R_{\ell }=\sigma \circ A_{\ell }}$都是先进行一次仿射变换${\displaystyle A_{i}}$，再逐点应用一个非线性激活函数${\displaystyle \sigma :\mathbb {R} \to \mathbb {R} }$。网络的全部可训练参数${\displaystyle \theta \in \mathbb {R} ^{P}}$定义了这些仿射变换${\displaystyle A_{0},\ldots ,A_{L-1}}$，并在训练开始前以独立同分布的方式进行随机初始化。

随着网络宽度的增加，NTK的尺度会受到${\displaystyle A_{i}}$及参数初始化的影响。为得到一个性质良好的极限，有研究提出了所谓的NTK参数化：${\displaystyle A_{\ell }\left(x\right)={\frac {1}{\sqrt {n_{\ell }}}}W^{\left(\ell \right)}x+b^{\left(\ell \right)}}$ 。采用这种参数化并以标准正态分布来初始化所有参数${\displaystyle \theta \in \mathbb {R} ^{P}}$时，可以确保当网络宽度趋于无穷时，NTK会收敛到一个有限且非平凡的极限${\displaystyle \Theta _{\infty }}$。${\displaystyle \Theta _{\infty }}$不仅是一个确定性（非随机）的函数 ，而且在训练过程中保持不变。

可以将${\displaystyle \Theta _{\infty }}$表示为${\displaystyle \Theta _{\infty }=\Theta ^{\left(L\right)}}$， 而${\displaystyle \Theta ^{\left(L\right)}}$则能通过以下递归方程组计算得到：

${\displaystyle {\begin{aligned}\Theta ^{\left(1\right)}\left(x,y\right)&=\Sigma ^{\left(1\right)}\left(x,y\right),\\\Sigma ^{\left(1\right)}\left(x,y\right)&={\frac {1}{n_{\mathrm {in} }}}x^{T}y+1,\\\Theta ^{\left(\ell +1\right)}\left(x,y\right)&=\Theta ^{\left(\ell \right)}\left(x,y\right){\dot {\Sigma }}^{\left(\ell +1\right)}\left(x,y\right)+\Sigma ^{\left(\ell +1\right)}\left(x,y\right),\\\Sigma ^{\left(\ell +1\right)}\left(x,y\right)&=L_{\Sigma ^{\left(\ell \right)}}^{\sigma }\left(x,y\right),\\{\dot {\Sigma }}^{\left(\ell +1\right)}\left(x,y\right)&=L_{\Sigma ^{\left(\ell \right)}}^{\dot {\sigma }},\end{aligned}}}$

在上述方程中，${\displaystyle L_{K}^{f}}$定义了由旧核${\displaystyle K}$和函数${\displaystyle f}$生成新核的方式，其计算依赖于高斯期望：

${\displaystyle L_{K}^{f}\left(x,y\right)=\mathbb {E} _{\left(X,Y\right)\sim {\mathcal {N}}\left(0,{\begin{pmatrix}K\left(x,x\right)&K\left(x,y\right)\\K\left(y,x\right)&K\left(y,y\right)\end{pmatrix}}\right)}\left[f\left(X\right)f\left(Y\right)\right].}$

式中的${\displaystyle \Sigma ^{\left(\ell \right)}}$被称为神经网络的激活核。^[28][29][5]

训练中参数的线性化

NTK描述了神经网络在函数空间中的动态演化，而与相对应的另一个视角则是网络在参数空间中的演化。在无限宽度极限下，这两个视角之间的关联变得十分有趣。与NTK在训练中保持恒定这一现象同时发生的，是神经网络在整个训练过程中的变化都可以由其在初始化参数处的一阶泰勒展开很好地近似：^[6]

${\displaystyle f\left(x;\theta (t)\right)=f\left(x;\theta (0)\right)+\nabla _{\theta }f\left(x;\theta (0)\right)\left(\theta (t)-\theta (0)\right)+{\mathcal {O}}\left(\min \left(n_{1}\dots n_{L-1}\right)^{-{\frac {1}{2}}}\right).}$