神经正切核（NTK）视角的缩放定律：无限宽网络极限下的学习动力学推导

大家好，今天我们来深入探讨一个深度学习领域的核心概念：神经正切核 (Neural Tangent Kernel, NTK)。我们将从 NTK 的基本原理出发，逐步推导出无限宽网络极限下的学习动力学，并最终解释缩放定律 (Scaling Laws) 如何在 NTK 框架下体现。

1. 神经网络与函数空间

首先，我们需要明确神经网络的本质：它是一个参数化的函数。给定输入 x 和网络参数 θ，神经网络 f(x; θ) 输出一个预测值。深度学习的目标是找到一组最优参数 θ，使得 f(x; θ)* 尽可能接近真实函数。

例如，一个简单的两层全连接神经网络可以表示为：

import numpy as np

def relu(x):
  return np.maximum(0, x)

def neural_network(x, W1, b1, W2, b2):
  """
  两层全连接神经网络。

  Args:
    x: 输入向量 (shape: (input_dim,)).
    W1: 第一层权重矩阵 (shape: (hidden_dim, input_dim)).
    b1: 第一层偏置向量 (shape: (hidden_dim,)).
    W2: 第二层权重矩阵 (shape: (output_dim, hidden_dim)).
    b2: 第二层偏置向量 (shape: (output_dim,)).

  Returns:
    输出向量 (shape: (output_dim,)).
  """
  z1 = np.dot(W1, x) + b1
  a1 = relu(z1)
  z2 = np.dot(W2, a1) + b2
  return z2

# 示例：初始化参数
input_dim = 10
hidden_dim = 50
output_dim = 1

W1 = np.random.randn(hidden_dim, input_dim)
b1 = np.random.randn(hidden_dim)
W2 = np.random.randn(output_dim, hidden_dim)
b2 = np.random.randn(output_dim)

# 示例：输入数据
x = np.random.randn(input_dim)

# 示例：计算输出
output = neural_network(x, W1, b1, W2, b2)
print(f"Output: {output}")

这个网络定义了一个从输入空间到输出空间的映射。网络的参数 θ = {W1, b1, W2, b2} 决定了这个映射的具体形式。训练神经网络的过程就是调整 θ 的过程，以最小化损失函数 L(f(x; θ), y)，其中 y 是真实标签。

2. 梯度下降与函数空间轨迹

我们通常使用梯度下降算法来训练神经网络：

θ_t+1 = θ_t – η ∇_θ L(f(x; θ_t), y)

其中 η 是学习率，∇_θ L 是损失函数关于参数 θ 的梯度。这个过程可以看作是在参数空间中沿着梯度方向迭代，最终收敛到局部最小值。

然而，从函数空间的角度来看，每次参数更新都会改变神经网络所表示的函数 f(x; θ)。因此，梯度下降也可以看作是在函数空间中寻找一个最优函数的轨迹。 NTK 理论的核心思想就是研究这种函数空间轨迹的性质。

3. 神经正切核 (NTK) 的定义

NTK 的定义基于神经网络输出关于参数的梯度： ∇_θ f(x; θ)。对于给定的两个输入 x 和 x’, NTK 定义为：

K(x, x’; θ) = ∇_θ f(x; θ)^T ∇_θ f(x’; θ)

换句话说，NTK 是神经网络在两个不同输入处的梯度向量的内积。它衡量了当参数发生微小变化时，两个输入对应的输出变化的相似程度。

更具体的，对于我们前面定义的两层全连接网络，我们可以计算其 NTK 如下：

def compute_ntk(x, x_prime, W1, b1, W2, b2):
  """
  计算两层全连接网络的 NTK。

  Args:
    x: 第一个输入向量 (shape: (input_dim,)).
    x_prime: 第二个输入向量 (shape: (input_dim,)).
    W1: 第一层权重矩阵 (shape: (hidden_dim, input_dim)).
    b1: 第一层偏置向量 (shape: (hidden_dim,)).
    W2: 第二层权重矩阵 (shape: (output_dim, hidden_dim)).
    b2: 第二层偏置向量 (shape: (output_dim,)).

  Returns:
    NTK 值 (float).
  """
  # 计算关于 W1, b1, W2, b2 的梯度
  def compute_gradients(x, W1, b1, W2, b2):
    z1 = np.dot(W1, x) + b1
    a1 = relu(z1)
    z2 = np.dot(W2, a1) + b2
    output = z2

    # 反向传播
    d_output = 1
    d_z2 = d_output
    d_W2 = np.outer(d_z2, a1)
    d_a1 = np.dot(W2.T, d_z2)
    d_z1 = d_a1 * (z1 > 0) # ReLU 的导数
    d_W1 = np.outer(d_z1, x)
    d_b1 = d_z1

    return d_W1, d_b1, d_W2, d_b2

  d_W1_x, d_b1_x, d_W2_x, d_b2_x = compute_gradients(x, W1, b1, W2, b2)
  d_W1_x_prime, d_b1_x_prime, d_W2_x_prime, d_b2_x_prime = compute_gradients(x_prime, W1, b1, W2, b2)

  # 将所有梯度向量化
  grad_x = np.concatenate([d_W1_x.flatten(), d_b1_x.flatten(), d_W2_x.flatten(), d_b2_x.flatten()])
  grad_x_prime = np.concatenate([d_W1_x_prime.flatten(), d_b1_x_prime.flatten(), d_W2_x_prime.flatten(), d_b2_x_prime.flatten()])

  # 计算 NTK
  ntk = np.dot(grad_x, grad_x_prime)
  return ntk

# 示例：计算 NTK
x = np.random.randn(input_dim)
x_prime = np.random.randn(input_dim)
ntk_value = compute_ntk(x, x_prime, W1, b1, W2, b2)
print(f"NTK Value: {ntk_value}")

这段代码演示了如何计算两层全连接网络的 NTK。关键步骤是计算神经网络输出关于所有参数的梯度，然后将梯度向量化，并计算内积。

4. 无限宽网络极限

NTK 理论最令人兴奋的结论之一是，当神经网络的宽度趋于无穷大时，NTK 会收敛到一个确定的核函数 K(x, x’)，并且在训练过程中保持不变。这意味着在无限宽网络极限下，神经网络的学习动力学可以完全由 NTK 描述。

为了理解这一点，我们需要考虑以下几点：

随机初始化： 神经网络的参数通常以均值为 0，方差为 1/W 的正态分布随机初始化，其中 W 是网络的宽度（例如，隐藏层的神经元数量）。
中心极限定理： 当 W 趋于无穷大时，神经网络的输出会收敛到一个高斯过程。
NTK 的稳定性： 在无限宽网络极限下，由于参数初始化的方差很小，并且梯度下降的步长也很小，参数的变化非常微弱，因此 NTK 在训练过程中几乎保持不变。

5. 线性化与 NTK 动力学

在 NTK 理论中，我们通常对神经网络的输出进行一阶泰勒展开：

f(x; θ_t+1) ≈ f(x; θ_t) + ∇_θ f(x; θ_t)^T (θ_t+1 – θ_t)

将梯度下降更新规则代入上式，可以得到：

f(x; θ_t+1) ≈ f(x; θ_t) – η ∇_θ f(x; θ_t)^T ∇_θ L(f(x; θ_t), y)

对于平方损失函数 L(f(x; θ), y) = (1/2) (f(x; θ) – y)²，我们可以进一步简化为：

f(x; θ_t+1) ≈ f(x; θ_t) – η ∇_θ f(x; θ_t)^T (f(x; θ_t) – y) ∇_θ f(x; θ_t)

现在，让我们考虑所有训练样本 x_i 的集合。我们可以将所有输出 f(x_i; θ_t) 堆叠成一个向量 f_t，将所有真实标签 y_i 堆叠成一个向量 y。那么，上述更新规则可以写成矩阵形式：

f_t+1 ≈ f_t – η K (f_t – y)

其中 K 是 NTK 矩阵，其元素为 K_ij = K(x_i, x_j; θ)。由于在无限宽网络极限下，K 在训练过程中保持不变，我们可以得到一个闭式解：

f_t = (I – *ηK)^t f₀ + ηK Σ_i=0^t-1 (I – η*K)ⁱ y

其中 I 是单位矩阵，f₀ 是神经网络的初始输出。这个公式表明，神经网络的输出 f_t 可以表示为初始输出 f₀ 和真实标签 y 的线性组合，系数由 NTK 矩阵 K 决定。

6. 缩放定律 (Scaling Laws) 的 NTK 解释

缩放定律是指，随着模型大小（例如，参数数量）、训练数据量和计算资源的增加，模型的性能会以某种幂律形式提升。 NTK 理论可以用来解释这种现象。

在 NTK 框架下，模型的性能取决于 NTK 矩阵 K 的性质。例如，K 的特征值分布决定了模型的泛化能力。当模型宽度增加时，K 会变得更加平滑，特征值分布会更加集中，从而提高模型的泛化能力。

此外，训练数据量也会影响 K 的性质。当训练数据量增加时，K 会更好地捕捉数据的内在结构，从而提高模型的性能。

最后，计算资源决定了我们可以训练多长时间。根据前面的闭式解，训练时间越长，f_t 越接近最优解。

因此，缩放定律可以看作是模型大小、训练数据量和计算资源对 NTK 矩阵 K 产生影响的结果。更具体地说，增加模型大小、训练数据量和计算资源可以使得 NTK 矩阵 K 更加稳定、更加平滑、更好地捕捉数据结构，从而提高模型的性能。

7. 代码示例：NTK 近似的线性回归

为了更好地理解 NTK 的应用，我们可以使用 NTK 近似来解决线性回归问题。假设我们有一个线性模型 f(x; θ) = θ^Tx，损失函数为平方损失。那么，NTK 为 K(x, x’) = x^Tx’。

import numpy as np

def ntk_linear_regression(X_train, y_train, X_test, learning_rate=0.01, epochs=1000):
  """
  使用 NTK 近似进行线性回归。

  Args:
    X_train: 训练数据 (shape: (n_train, input_dim)).
    y_train: 训练标签 (shape: (n_train,)).
    X_test: 测试数据 (shape: (n_test, input_dim)).
    learning_rate: 学习率.
    epochs: 训练轮数.

  Returns:
    预测标签 (shape: (n_test,)).
  """
  n_train = X_train.shape[0]
  n_test = X_test.shape[0]
  input_dim = X_train.shape[1]

  # 计算 NTK 矩阵
  K_train = np.dot(X_train, X_train.T)
  K_test_train = np.dot(X_test, X_train.T)

  # 初始化参数 (这里我们直接初始化为 0)
  f_train = np.zeros(n_train)

  # 梯度下降
  for epoch in range(epochs):
    f_train = f_train - learning_rate * np.dot(K_train, (f_train - y_train))

  # 预测
  y_pred = np.dot(K_test_train, (f_train - y_train))

  return y_pred

# 示例：生成数据
n_train = 100
n_test = 50
input_dim = 10

X_train = np.random.randn(n_train, input_dim)
X_test = np.random.randn(n_test, input_dim)
theta_true = np.random.randn(input_dim)
y_train = np.dot(X_train, theta_true) + 0.1 * np.random.randn(n_train) # 添加噪声
y_test = np.dot(X_test, theta_true)

# 示例：使用 NTK 线性回归
y_pred = ntk_linear_regression(X_train, y_train, X_test)

# 示例：评估性能 (均方误差)
mse = np.mean((y_pred - y_test)**2)
print(f"Mean Squared Error: {mse}")

这段代码演示了如何使用 NTK 近似来解决线性回归问题。关键步骤是计算 NTK 矩阵，并使用梯度下降来更新训练数据的输出。然后，我们可以使用 NTK 矩阵和训练数据的输出来预测测试数据的标签。

8. NTK 理论的局限性

虽然 NTK 理论提供了一个理解深度学习的重要视角，但它也有一些局限性：

无限宽网络假设： NTK 理论主要适用于无限宽网络。对于有限宽度的网络，NTK 在训练过程中可能会发生显著变化。
线性化假设： NTK 理论基于线性化假设，这意味着神经网络的输出可以用一阶泰勒展开来近似。对于某些非线性问题，这种近似可能不够准确。
计算复杂度： 计算 NTK 矩阵的复杂度很高，特别是对于大型数据集。

尽管存在这些局限性，NTK 理论仍然是深度学习领域的一个重要研究方向。它为我们理解深度学习的泛化能力、优化动力学和缩放定律提供了新的思路。

9. 总结：NTK视角的缩放定律

总结一下，我们从神经正切核 (NTK) 的基本原理出发，推导了无限宽网络极限下的学习动力学。在无限宽网络极限下，神经网络的学习过程可以被简化为线性模型，其行为完全由一个固定的核函数（NTK）决定。通过分析 NTK 的性质，我们可以理解缩放定律如何在深度学习中体现，即模型大小、训练数据量和计算资源的增加如何影响模型的性能。虽然 NTK 理论存在一些局限性，但它提供了一个强大的框架，帮助我们理解和解释深度学习的内在机制。

希望今天的讲座对大家有所帮助。谢谢！

神经正切核（NTK）视角的缩放定律：无限宽网络极限下的学习动力学推导

发表回复 取消回复

发表回复取消回复