注意
转到末尾下载完整的示例代码。
神经切向核#
创建日期：2023年3月15日 | 最后更新：2023年6月16日 | 最后验证：未验证
神经切向核 (NTK) 是一种描述神经网络在训练过程中如何演变的核函数。近年来，围绕它已有很多研究，例如 此篇论文。本教程受到 JAX 中 NTK 实现的启发（详见《快速有限宽度神经切向核》），演示了如何使用 torch.func（PyTorch 的可组合函数变换）轻松计算这一量。
注意
本教程需要 PyTorch 2.0.0 或更高版本。
设置#
首先进行一些设置。让我们定义一个简单的 CNN，我们将计算它的 NTK。
import torch
import torch.nn as nn
from torch.func import functional_call, vmap, vjp, jvp, jacrev
device = 'cuda' if torch.cuda.device_count() > 0 else 'cpu'
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(3, 32, (3, 3))
self.conv2 = nn.Conv2d(32, 32, (3, 3))
self.conv3 = nn.Conv2d(32, 32, (3, 3))
self.fc = nn.Linear(21632, 10)
def forward(self, x):
x = self.conv1(x)
x = x.relu()
x = self.conv2(x)
x = x.relu()
x = self.conv3(x)
x = x.flatten(1)
x = self.fc(x)
return x
然后，让我们生成一些随机数据。
x_train = torch.randn(20, 3, 32, 32, device=device)
x_test = torch.randn(5, 3, 32, 32, device=device)
创建模型的函数版本#
torch.func 变换作用于函数。具体来说，为了计算 NTK，我们需要一个函数，它接受模型的参数和单个输入（而不是一批输入！），并返回单个输出。
我们将使用 torch.func.functional_call，它允许我们使用不同的参数/缓冲区调用一个 nn.Module，来帮助完成第一步。
请记住，该模型最初是为接受一批输入数据点而编写的。在我们的 CNN 示例中，没有批次间的操作。也就是说，批次中的每个数据点都与其他数据点无关。基于这个假设，我们可以轻松地生成一个函数，用于在单个数据点上评估模型。
net = CNN().to(device)
# Detaching the parameters because we won't be calling Tensor.backward().
params = {k: v.detach() for k, v in net.named_parameters()}
def fnet_single(params, x):
return functional_call(net, params, (x.unsqueeze(0),)).squeeze(0)
计算 NTK：方法 1（雅可比矩阵收缩）#
我们准备好计算经验 NTK 了。对于两个数据点 \(x_1\) 和 \(x_2\)，经验 NTK 定义为在 \(x_1\) 处评估的模型雅可比矩阵与在 \(x_2\) 处评估的模型雅可比矩阵之间的矩阵乘积。
在批处理情况下，\(x_1\) 是一批数据点，\(x_2\) 也是一批数据点，那么我们想要的是 \(x_1\) 和 \(x_2\) 中所有数据点组合的雅可比矩阵之间的矩阵乘积。
第一种方法就是这样做——计算两个雅可比矩阵，然后将它们收缩。以下是如何在批处理情况下计算 NTK：
def empirical_ntk_jacobian_contraction(fnet_single, params, x1, x2):
# Compute J(x1)
jac1 = vmap(jacrev(fnet_single), (None, 0))(params, x1)
jac1 = jac1.values()
jac1 = [j.flatten(2) for j in jac1]
# Compute J(x2)
jac2 = vmap(jacrev(fnet_single), (None, 0))(params, x2)
jac2 = jac2.values()
jac2 = [j.flatten(2) for j in jac2]
# Compute J(x1) @ J(x2).T
result = torch.stack([torch.einsum('Naf,Mbf->NMab', j1, j2) for j1, j2 in zip(jac1, jac2)])
result = result.sum(0)
return result
result = empirical_ntk_jacobian_contraction(fnet_single, params, x_train, x_test)
print(result.shape)
torch.Size([20, 5, 10, 10])
在某些情况下，您可能只需要这个数量的对角线或迹，特别是如果您事先知道网络架构产生的 NTK 中，非对角线元素可以近似为零。可以很容易地调整上面的函数来实现这一点。
def empirical_ntk_jacobian_contraction(fnet_single, params, x1, x2, compute='full'):
# Compute J(x1)
jac1 = vmap(jacrev(fnet_single), (None, 0))(params, x1)
jac1 = jac1.values()
jac1 = [j.flatten(2) for j in jac1]
# Compute J(x2)
jac2 = vmap(jacrev(fnet_single), (None, 0))(params, x2)
jac2 = jac2.values()
jac2 = [j.flatten(2) for j in jac2]
# Compute J(x1) @ J(x2).T
einsum_expr = None
if compute == 'full':
einsum_expr = 'Naf,Mbf->NMab'
elif compute == 'trace':
einsum_expr = 'Naf,Maf->NM'
elif compute == 'diagonal':
einsum_expr = 'Naf,Maf->NMa'
else:
assert False
result = torch.stack([torch.einsum(einsum_expr, j1, j2) for j1, j2 in zip(jac1, jac2)])
result = result.sum(0)
return result
result = empirical_ntk_jacobian_contraction(fnet_single, params, x_train, x_test, 'trace')
print(result.shape)
torch.Size([20, 5])
该方法的渐进时间复杂度为 \(N O [FP]\)（计算雅可比矩阵的时间）+ \(N^2 O^2 P\)（收缩雅可比矩阵的时间），其中 \(N\) 是 \(x_1\) 和 \(x_2\) 的批大小，\(O\) 是模型的输出大小，\(P\) 是参数总数，\([FP]\) 是单次前向传播通过模型的成本。详情请参见《快速有限宽度神经切向核》的 3.2 节。
计算 NTK：方法 2（NTK-向量积）#
我们接下来要讨论的方法是使用 NTK-向量积来计算 NTK。
该方法将 NTK 重新表述为一系列应用于大小为 \(O\times O\) 的单位矩阵 \(I_O\)（其中 \(O\) 是模型的输出大小）的列的 NTK-向量积的堆叠。
其中 \(e_o\in \mathbb{R}^O\) 是单位矩阵 \(I_O\) 的列向量。
设 \(\textrm{vjp}_o = J_{net}^T(x_2) e_o\)。我们可以使用向量-雅可比积来计算它。
现在，考虑 \(J_{net}(x_1) \textrm{vjp}_o\)。这是一个雅可比-向量积！
最后，我们可以使用
vmap对 \(I_O\) 的所有列 \(e_o\) 并行运行上述计算。
这表明我们可以结合使用反向模式自动微分（计算向量-雅可比积）和前向模式自动微分（计算雅可比-向量积）来计算 NTK。
让我们把它写成代码：
def empirical_ntk_ntk_vps(func, params, x1, x2, compute='full'):
def get_ntk(x1, x2):
def func_x1(params):
return func(params, x1)
def func_x2(params):
return func(params, x2)
output, vjp_fn = vjp(func_x1, params)
def get_ntk_slice(vec):
# This computes ``vec @ J(x2).T``
# `vec` is some unit vector (a single slice of the Identity matrix)
vjps = vjp_fn(vec)
# This computes ``J(X1) @ vjps``
_, jvps = jvp(func_x2, (params,), vjps)
return jvps
# Here's our identity matrix
basis = torch.eye(output.numel(), dtype=output.dtype, device=output.device).view(output.numel(), -1)
return vmap(get_ntk_slice)(basis)
# ``get_ntk(x1, x2)`` computes the NTK for a single data point x1, x2
# Since the x1, x2 inputs to ``empirical_ntk_ntk_vps`` are batched,
# we actually wish to compute the NTK between every pair of data points
# between {x1} and {x2}. That's what the ``vmaps`` here do.
result = vmap(vmap(get_ntk, (None, 0)), (0, None))(x1, x2)
if compute == 'full':
return result
if compute == 'trace':
return torch.einsum('NMKK->NM', result)
if compute == 'diagonal':
return torch.einsum('NMKK->NMK', result)
# Disable TensorFloat-32 for convolutions on Ampere+ GPUs to sacrifice performance in favor of accuracy
with torch.backends.cudnn.flags(allow_tf32=False):
result_from_jacobian_contraction = empirical_ntk_jacobian_contraction(fnet_single, params, x_test, x_train)
result_from_ntk_vps = empirical_ntk_ntk_vps(fnet_single, params, x_test, x_train)
assert torch.allclose(result_from_jacobian_contraction, result_from_ntk_vps, atol=1e-5)
我们为 empirical_ntk_ntk_vps 编写的代码看起来就像是上面数学公式的直接翻译！这展示了函数变换的强大功能：如果只使用 torch.autograd.grad，想写出上述代码的高效版本可就难了。
该方法的渐进时间复杂度为 \(N^2 O [FP]\)，其中 \(N\) 是 \(x_1\) 和 \(x_2\) 的批大小，\(O\) 是模型的输出大小，\([FP]\) 是单次前向传播通过模型的成本。因此，该方法比方法 1（雅可比收缩）执行了更多的网络前向传播（\(N^2 O\) 次而非 \(N O\) 次），但完全避免了收缩成本（没有 \(N^2 O^2 P\) 项，其中 \(P\) 是模型的参数总数）。因此，当 \(O P\)相对于 \([FP]\) 很大时，例如对于具有许多输出 \(O\) 的全连接（非卷积）模型，此方法更可取。在内存方面，两种方法应该相当。详情请参见《快速有限宽度神经切向核》的 3.3 节。
脚本总运行时间： (0 分 0.795 秒)
