偏微分方程与机器学习（三）：变分原理与优化

训练神经网络的本质是什么？当我们在高维参数空间中运行梯度下降时，背后是否存在某种更深刻的连续时间动力学？当网络宽度趋于无穷时，离散的参数更新是否会收敛到某个优雅的偏微分方程？这些问题的答案，正位于变分法、最优传输与 PDE 理论的交汇处。

过去十年深度学习的成功主要依赖工程直觉。然而近年来，数学家们提出了一个引人注目的观点：将神经网络视为概率测度空间上的粒子系统，并在 Wasserstein 几何下研究其演化，能够揭示训练过程的全局结构——包括收敛性保证、过参数化的作用以及初始化的意义。使这一视角成为可能的核心工具正是变分原理：从物理学中的最小作用量原理，到现代最优传输中的 JKO 格式，再到神经网络的平均场极限。

本文将从基础出发，逐步构建这一图景。我们首先回顾经典变分法（泛函、第一变分、Euler-Lagrange 方程），接着引入 Wasserstein 几何与梯度流理论（例如热方程可视为熵的梯度流），最终推导两层神经网络的平均场方程，并由此解读其全局收敛行为。

图 1. 在固定起点 $$A=(0,0)$$ 与终点 $B=(\pi, 2)$ 之间，几条候选曲线的下降时间 $T[y] = \int \sqrt{(1+y'^2)/(2gy)}\,dx$ 。摆线给出最小值——这是变分法的经典范例。

训练神经网络这件事，到底在做什么？

我们在高维参数空间里跑梯度下降——这背后有没有更深刻的连续时间动力学？当网络宽度趋于无穷时，离散的参数更新会不会收敛到某个优雅的偏微分方程？这些问题听起来像是哲学，但其实有非常具体的数学回答，它们坐落在变分法、最优传输和 PDE 理论的交汇处。

过去十年深度学习靠工程直觉赢了很多，但近年来数学家们提出了一个让人眼前一亮的视角：把神经网络看成概率测度空间上的粒子系统，在 Wasserstein 几何下研究它的演化。这个视角能给出收敛性保证、解释过参数化的作用、说明初始化为什么重要。撑起这个视角的核心工具是变分原理——从物理里的最小作用量原理，到现代最优传输里的 JKO 格式，再到神经网络的平均场极限，是同一根线索的不同长度。

本文按照“先经典再现代”的顺序展开：先回顾经典变分法（泛函、第一变分、Euler-Lagrange），再引入 Wasserstein 几何与梯度流（你会看到热方程其实是熵的梯度流），最后推导两层神经网络的平均场方程，从中读出全局收敛行为。

变分法基础：从泛函到 Euler-Lagrange 方程#

泛函与第一变分#

变分法的基本对象是泛函——一种以函数为输入、输出实数的映射。如果说普通函数“吃”数字，那么泛函“吃”的则是整个函数。

定义（泛函）：设 $$X$$ 是一个函数空间（例如 $$C^1([a,b])$$ ），泛函 $J : X \to \mathbb{R}$ 将每个 $y \in X$ 映射为一个实数 $$J[y]$$ 。

本文将反复用到以下三个典型例子：

弧长：曲线 $$y(x)$$ 在区间 $$[a,b]$$ 上的长度为 $L[y] = \int_a^b \sqrt{1 + y'(x)^2}\, dx.$
旋转曲面面积：将 $$y$$ 绕 $$x$$ 轴旋转所得曲面的面积为 $A[y] = 2\pi \int_a^b y(x)\sqrt{1 + y'(x)^2}\, dx.$
作用量：质点轨迹 $$q(t)$$ 的作用量为 $S[q] = \int_{t_0}^{t_1} L(q(t), \dot q(t), t)\, dt$ ，其中 $$L$$ 是拉格朗日量。

核心问题是：在满足给定边界条件的所有函数中，哪一个使泛函 $$J$$ 取得极值？

\delta J[y; \eta] = \lim_{\varepsilon \to 0} \frac{J[y + \varepsilon \eta] - J[y]}{\varepsilon}.

\delta J[y; \eta] = \int \frac{\delta J}{\delta y}(x)\, \eta(x)\, dx,

并称 $\delta J / \delta y$ 为 $$J$$ 的泛函导数。

J[y] = \int_a^b F(x, y(x), y'(x))\, dx

\frac{\partial F}{\partial y} - \frac{d}{dx}\frac{\partial F}{\partial y'} = 0.

\int_a^b \left(\frac{\partial F}{\partial y}\eta + \frac{\partial F}{\partial y'}\eta'\right) dx = 0.

\int_a^b \left(\frac{\partial F}{\partial y} - \frac{d}{dx}\frac{\partial F}{\partial y'}\right)\eta\, dx = 0

对任意容许的 $\eta$ 成立。由变分法基本引理可知，括号内的表达式必须恒等于零，这正是 Euler-Lagrange 方程。 $\square$

Euler-Lagrange 几何图像：左图为扰动族 $y+\varepsilon\eta$
（$\eta$
在端点为零）；右图为 $\Phi(\varepsilon)$
的剖面，极值表现为 $\varepsilon=0$
处水平切线。

图 2. 第一变分论证的几何呈现。极值曲线是唯一一条这样的曲线：沿任何容许扰动方向的方向导数都为零。

最速降线问题（Brachistochrone）#

问题：在均匀重力场中，从点 $$A=(0,0)$$ 到点 $$B$$ 的无摩擦曲线上，质点下滑所需时间最短的路径是什么？

T[y] = \int_0^{x_B} \frac{\sqrt{1 + y'(x)^2}}{\sqrt{2 g\, y(x)}}\, dx.

y(1 + y'^2) = \frac{1}{2gC^2} =: 2R.

x(\theta) = R(\theta - \sin\theta), \qquad y(\theta) = R(1 - \cos\theta) ,

即半径为 $$R$$ 的圆沿直线滚动时其圆周上一点的轨迹。图 1 显示，直线、抛物线、圆弧等其他候选曲线的下降时间均严格大于摆线。

实现细节
最速降线是变分法教科书中的经典案例，但通过数值方式直观观察其行为，更能加深理解。我们采用梯形数值积分对泛函 $$T[y]$$ 进行离散化，并对比五条候选路径的下降时间：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import numpy as np

g = 9.81
x_B, y_B = np.pi, 2.0

def descent_time(y_vals, x_vals):
    # 计算 T[y] = ∫ sqrt((1 + y'²) / (2*g*y)) dx
    dx = x_vals[1] - x_vals[0]
    yp = np.gradient(y_vals, dx)
    integrand = np.sqrt((1 + yp**2) / (2 * g * np.maximum(y_vals, 1e-12)))
    return np.trapz(integrand, x_vals)

N = 200
x = np.linspace(0, x_B, N)
theta = np.linspace(0, np.pi, 1000)
R = y_B / 2
x_cyc = R * (theta - np.sin(theta))
y_cyc = R * (1 - np.cos(theta))
y_cycloid = np.interp(x, x_cyc, y_cyc)

y_line = y_B / x_B * x
y_parabola = y_B * (x / x_B)**2
y_cubic = y_B * (x / x_B)**3

for name, y in [("Cycloid", y_cycloid), ("Line", y_line),
                ("Parabola", y_parabola), ("Cubic", y_cubic)]:
    print(f"{name:12s}  T = {descent_time(y, x):.4f} s")

摆线以约 10–15% 的显著优势胜出——这绝非微小差距。注意：初始段陡峭的路径（如三次曲线）确实有助于加速，但因其末端路径过于平缓，整体仍非最优。而摆线则完美权衡了重力驱动的加速度与路径总长度，恰好满足欧拉-拉格朗日方程所要求的极值条件。

这对机器学习的意义何在？
最速降线问题本质上是对一个泛函进行标量优化；而神经网络训练，正是同一类问题在高维空间中的推广：在所有从初始化到收敛的权重轨迹中，梯度下降所遵循的那条路径，（近似地）最小化了一个类时间型泛函。此时，欧拉-拉格朗日方程退化为训练过程所满足的动力学常微分方程（ODE）；而贝尔特拉米恒等式则对应于该动力学流中能量的耗散机制。从单条曲线跃迁到概率测度空间——这正是后续章节将要展开的核心内容。

Hamilton 原理与辛几何视角#

\frac{d}{dt}\frac{\partial L}{\partial \dot q} - \frac{\partial L}{\partial q} = 0

\dot q = \frac{\partial H}{\partial p}, \qquad \dot p = -\frac{\partial H}{\partial q} .

辛二形式 $\omega = dp \wedge dq$ 沿流保持不变。本文在第 7 节再次看到这一结构，并将其与梯度流进行对比。

从泛函导数到梯度流#

\frac{\delta J}{\delta y} = \frac{\partial F}{\partial y} - \frac{d}{dx}\frac{\partial F}{\partial y'} .

例（Dirichlet 能量）：对 $J[u] = \tfrac12 \int |\nabla u|^2\, dx$ ，有 $\delta J/\delta u = -\Delta u$ ，因此极值条件即为Laplace 方程 $\Delta u = 0$ 。

\partial_t u = -\frac{\delta J}{\delta u} .

对 Dirichlet 能量而言，这就是热方程 $\partial_t u = \Delta u$ 。物理学中许多重要 PDE 都可如此解释。而可以看到，即使状态不再是函数，而是一个概率测度，这种结构依然成立。

实现：狄利克雷能量的梯度流
热方程 $\partial_t u = \Delta u$ 正是能量泛函 $J[u] = \frac{1}{2}\int |\nabla u|^2\,dx$ 的梯度流。我们可以通过显式时间步进求解热方程，并观察 $$J$$ 随时间单调递减来数值验证这一性质：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
import numpy as np

N, dx = 256, 1.0 / 256
dt = 0.4 * dx**2  # 显式欧拉格式的稳定性条件
x = np.linspace(0, 1, N)
u = np.sin(3 * np.pi * x) + 0.5 * np.sin(7 * np.pi * x)  # 初始条件

def dirichlet_energy(u, dx):
    grad = np.diff(u) / dx  # 一阶差分近似梯度
    return 0.5 * np.sum(grad**2) * dx  # 狄利克雷能量离散化

energies = []
for step in range(5000):
    energies.append(dirichlet_energy(u, dx))
    laplacian = (np.roll(u, -1) - 2*u + np.roll(u, 1)) / dx**2  # 二阶中心差分近似拉普拉斯算子
    u = u + dt * laplacian  # 显式欧拉更新
    u[0] = u[-1] = 0  # 施加齐次狄利克雷边界条件

print(f"能量: {energies[0]:.2f} → {energies[-1]:.4f}")

狄利克雷能量在每一步时间迭代中都严格递减——既不会上升，也不会因浮点数舍入误差而出现任何微小反弹。这并非数值巧合，而是梯度流的本质特征：能量沿其负梯度方向演化，从而保证能量单调下降。本文后续介绍的所有偏微分方程均具备这一结构：先识别对应的物理或几何能量，再沿其最速下降方向构造演化方程，所得 PDE 自然具有明确的物理意义。

梯度流与 Wasserstein 几何#

$\mathbb{R}^n$ 中的梯度流#

\dot x(t) = -\nabla f(x(t)),

即最速下降法的连续时间类比。以下三点对后续讨论至关重要：

能量耗散： $\frac{d}{dt} f(x(t)) = -\|\nabla f(x(t))\|^2 \leq 0.$
平衡点即 $$f$$ 的临界点；在（强）凸性条件下，流会收敛到唯一极小值。
隐式欧拉离散化： $x_{k+1} = \arg\min_x \{ f(x) + \tfrac{1}{2\tau}\|x - x_k\|^2 \}$ 。本文用 Wasserstein 距离替代欧氏距离，从而得到 JKO 格式。

Wasserstein 距离#

当系统状态是一个概率密度时，欧氏度量不再自然。合适的替代是Wasserstein-2 距离，它衡量将一个分布转换为另一个所需的“搬运成本”。

W_2^2(\rho_0, \rho_1) = \inf_{\gamma \in \Pi(\rho_0, \rho_1)} \int_{\mathbb{R}^d \times \mathbb{R}^d} |x - y|^2 \, d\gamma(x, y),

其中 $\Pi(\rho_0, \rho_1)$ 是边缘分布为 $\rho_0, \rho_1$ 的耦合集合。

Brenier 定理：若 $\rho_0$ 绝对连续，则最优耦合集中在某个映射 $T = \nabla \varphi$ 的图像上，其中 $\varphi$ 是凸函数，且 $\rho_1 = T_\# \rho_0$ 。

W_2^2 = \|\mu_0 - \mu_1\|^2 + \mathrm{tr}\!\left(\Sigma_0 + \Sigma_1 - 2 (\Sigma_0^{1/2}\Sigma_1\Sigma_0^{1/2})^{1/2}\right).

实现：一维情形下计算 $$W_2$$ 。
在一维情况下，Wasserstein 距离具有优美的闭式解：只需对两组样本分别排序，再逐点配对即可。这是因为最优传输映射恰好对应于分位数函数（quantile function）。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
import numpy as np

def wasserstein_2_1d(samples_a, samples_b):
    # 计算两个一维经验分布之间的精确 $W_2$ 距离
    a_sorted = np.sort(samples_a)
    b_sorted = np.sort(samples_b)
    if len(a_sorted) != len(b_sorted):
        n = max(len(a_sorted), len(b_sorted))
        # 使用等间距分位点（含偏移以避免端点问题）
        q = np.linspace(0, 1, n, endpoint=False) + 0.5/n
        a_sorted = np.quantile(samples_a, q)
        b_sorted = np.quantile(samples_b, q)
    return np.sqrt(np.mean((a_sorted - b_sorted)**2))

rng = np.random.default_rng(42)
p = rng.normal(0, 1, 10000)      # 样本来自 N(0, 1)
q = rng.normal(3, 2, 10000)      # 样本来自 N(3, 2)
print(f"W_2(N(0,1), N(3,2)) = {wasserstein_2_1d(p, q):.3f}")

这种一维排序技巧可在 $O(n \log n)$ 时间内给出精确的最优传输解。而在高维情形下，计算 $$W_2$$ 需要求解线性规划问题，或采用熵正则化方法（如 Sinkhorn 算法），计算开销显著更高。

概率度量对比。
定义在 $\mathcal{P}(\mathbb{R}^d)$ 上的不同度量，刻画的是概率分布间不同层面的差异。在算法设计中，度量的选择至关重要：

度量	对什么敏感	一维计算复杂度	$$d$$ 维计算复杂度	典型应用场景
KL 散度	密度比 $\frac{p}{q}$	$$O(n)$$	$$O(n)$$	变分自编码器（VAE）、变分推断（VI）、信息论
$$W_2$$ （Wasserstein 距离）	质量迁移（mass transport）	$O(n\log n)$	$$O(n^3)$$ 或 Sinkhorn 迭代	最优传输、WGAN、JKO 梯度流
Fisher-Rao 度量	分数函数（score） $\nabla\log p$	$$O(n)$$	$O(n \cdot d)$	自然梯度法、摊销推断（amortised inference）
TV（全变差距离）	支撑集（support）重叠程度	$$O(n)$$	$$O(n)$$	假设检验、收敛性证明
MMD（核均值嵌入）	RKHS 中的矩（moments）	$$O(n^2)$$	$$O(n^2)$$	两样本检验、核方法

关键洞见：
当两个分布的支撑集不重叠时，KL 散度为无穷大——这在生成模型训练初期极易发生，直接导致优化失败。而 Wasserstein 距离始终有限，且能度量弱收敛（weak convergence），因此 WGAN 的训练稳定性远优于原始 GAN。Fisher-Rao 度量则天然尊重参数化分布族所构成的流形结构，这也是自然梯度法相比标准随机梯度下降（SGD）收敛更快的根本原因。

JKO 格式： $\mathcal{P}_2(\mathbb{R}^d)$ 上的梯度流#

如何在概率测度空间上定义梯度流？Jordan、Kinderlehrer 和 Otto（1998）给出了答案：只需将隐式欧拉中的欧氏距离替换为 Wasserstein 距离。

\rho_{k+1}^\tau \in \arg\min_{\rho} \left\{\mathcal{E}[\rho] + \frac{1}{2\tau} W_2^2(\rho, \rho_k^\tau)\right\}.

当 $\tau \to 0$ 时（若极限存在），该离散过程的连续时间极限称为 $\mathcal{E}$ 的Wasserstein 梯度流。

\mathcal{F}[\rho] = \int V\rho + \int \rho\log\rho,

我们可以采用一种简洁的基于粒子的近似方法实现：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import numpy as np

def jko_particles(particles, V_grad, tau, n_steps):
    # 基于粒子的 JKO 格式，用于求解 Fokker-Planck 方程
    # particles: (N,) 形状的一维位置数组
    # V_grad: 可调用对象，表示势能函数 V 的梯度
    history = [particles.copy()]
    for _ in range(n_steps):
        N = len(particles)
        bw = 0.5 * N**(-0.2)  # 核密度估计（KDE）带宽
        diffs = particles[:, None] - particles[None, :]
        kernel = np.exp(-diffs**2 / (2 * bw**2))
        # 利用 KDE 梯度估计经验密度的得分函数（score function）
        score = np.sum(diffs * kernel, axis=1) / (bw**2 * np.sum(kernel, axis=1))
        # JKO 步骤：漂移项（由势能驱动）+ 扩散项（由熵产生的排斥效应）
        particles = particles - tau * V_grad(particles) + tau * score
        history.append(particles.copy())
    return history

V_grad = lambda x: 4 * x * (x**2 - 1)
rng = np.random.default_rng(0)
x0 = rng.normal(2.0, 0.3, 2000)  # 初始分布远离平衡态（集中在 x ≈ 2 附近）
traj = jko_particles(x0, V_grad, tau=0.01, n_steps=500)
print(f"均值: {traj[0].mean():.2f} → {traj[-1].mean():.2f}")
print(f"标准差: {traj[0].std():.2f} → {traj[-1].std():.2f}")

\delta\mathcal{F}/\delta\rho = V + \log\rho + 1

的负梯度：其中 $-\nabla V$ 表示外部势场施加的力，而 $-\nabla\log\rho$ 则代表由熵驱动的压力（即“熵斥力”）。

热方程作为熵的梯度流#

\mathcal{H}[\rho] = \int \rho \log \rho\, dx

的 Wasserstein 梯度流即为热方程 $\partial_t \rho = \Delta \rho$ 。

简要解释：JKO 的一阶最优性条件为 $\delta \mathcal{H}/\delta \rho + \varphi/\tau = \text{常数}$ ，其中 $\varphi$ 是将 $\rho_k$ 推至 $\rho_{k+1}$ 的 Brenier 势。由于 $\delta \mathcal{H}/\delta\rho = \log\rho + 1$ ，诱导出的速度场为 $v = -\nabla \log \rho$ 。代入连续性方程 $\partial_t \rho + \nabla\cdot(\rho v) = 0$ ，即得 $\partial_t \rho = \nabla\cdot(\rho \nabla \log \rho) = \Delta \rho$ 。

同一框架还能导出其他经典 PDE。例如，自由能 $\mathcal{F}[\rho] = \int V \rho\, dx + \int \rho\log\rho\, dx$ 对应Fokker-Planck 方程 $\partial_t \rho = \nabla\cdot(\rho \nabla V) + \Delta \rho$ ，其稳态解为 Gibbs 分布 $\rho_\infty \propto e^{-V}$ ；内能 $\int \rho^m \, dx$ 对应多孔介质方程 $\partial_t \rho = \Delta(\rho^m)$ ；加入吸引相互作用项后，可得Keller-Segel 趋化系统。

自由能 $\mathcal{F}[\rho]$
的 Wasserstein 梯度流：左图为密度快照，右图为沿流的能量耗散。

图 3. 一维 Fokker-Planck 数值仿真。密度 $\rho_t$ 在 $\mathcal{P}_2(\mathbb{R})$ 中沿自由能最速下降方向移动，最终收敛到 Gibbs 平衡 $\rho_\ast \propto e^{-V}$ 。右图显示自由能差 $F[\rho_t] - F[\rho_\ast]$ 和均值漂移 $|\mathbb{E}\rho_t - \mu_\ast|$ 单调递减——梯度流的直接数值证据。

模拟揭示了什么？ 图 3 展现出三个值得关注的特征：

$\frac{d}{dt}\mathcal{E}[\rho_t] = -\|\nabla_{W_2}\mathcal{E}\|^2 \leq 0.$

这一耗散恒等式，正是变分结构“显露锋芒”的直接体现。
$F[\rho_t] - F[\rho_\ast] \leq e^{-2\lambda t}(F[\rho_0] - F[\rho_\ast]),$

其中 $\lambda$ 为对数 Sobolev 常数。对于方差为 $\sigma^2$ 的高斯分布，有 $\lambda = 1/\sigma^2$ 。
密度整体平移运动 —— 与 $$L^2$$ 空间中的逐点收敛不同，Wasserstein 收敛刻画的是质量的输运过程，而非仅函数值的局部调整：密度峰沿水平方向滑动，而非仅仅在原地被平滑抹平。正因如此，Wasserstein 几何才是神经网络训练的天然舞台——在此场景中，权重分布确实在参数空间中发生真实的整体移动。

神经网络训练的平均场理论#

从有限宽度到无限宽度#

f_\theta(x) = \frac{1}{m} \sum_{i=1}^m a_i\, \sigma(w_i^\top x + b_i),

其中参数为 $\theta_i = (a_i, w_i, b_i)$ ，宽度为 $$m$$ 。给定数据 $\{(x_k, y_k)\}_{k=1}^n$ ，训练损失为经验风险 $\hat R(\theta) = \frac{1}{n}\sum_k \ell(f_\theta(x_k), y_k)$ ，梯度下降按 $-\eta \nabla_{\theta_i}\hat R$ 更新每个 $\theta_i$ 。

粒子图像：将每个神经元视为参数空间中的粒子，训练过程即一个相互作用的粒子系统： $$m$$ 个粒子通过 $f_\theta$ 耦合在一起。

\rho_t^m := \frac{1}{m} \sum_{i=1}^m \delta_{\theta_i(t)}

f_{\rho_t^m}(x) = \int a\,\sigma(w^\top x + b)\, d\rho_t^m(\theta) .

平均场极限：在适当缩放下（学习率 $\propto 1/\sqrt{m}$ 或 $$1/m$$ ，取决于参数化方式），若初始参数独立同分布于 $\rho_0$ ，则当 $m \to \infty$ 时， $\rho_t^m \xrightharpoonup{} \rho_t$ ，其中 $\rho_t$ 满足平均场方程（一种 Vlasov / 连续性 PDE）。

两层 ReLU 网络的 Mean-Field 极限：不同宽度下，第一层权重直方图随训练演化。

图 4. 经验权重分布 $\rho_t^m$ 随宽度增大趋于光滑。到 $$m = 2000$$ 时已接近连续密度，符合 Vlasov 型理论预言。最右列显示损失曲线：宽度越大收敛越快（过参数化效应）。

实现：观察平均场极限的涌现过程
以下代码训练了三个不同宽度的两层 ReLU 网络，并记录第一层权重直方图在训练过程中的演化：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import torch
import torch.nn as nn

def train_and_record(width, n_epochs=1500, n_data=64):
    # 训练一个两层 ReLU 网络；返回关键训练轮次的权重快照
    torch.manual_seed(42)
    x = torch.linspace(-1, 1, n_data).unsqueeze(1)
    y = torch.sin(torch.pi * x) + 0.5 * torch.sin(3 * torch.pi * x)

    model = nn.Sequential(
        nn.Linear(1, width), nn.ReLU(), nn.Linear(width, 1)
    )
    lr = 0.1 / width**0.5  # 平均场缩放规则：学习率按 $1/\sqrt{m}$ 缩放
    opt = torch.optim.SGD(model.parameters(), lr=lr)

    snapshots = {}
    for epoch in range(n_epochs):
        loss = ((model(x) - y)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        if epoch in [0, 200, 500, n_epochs - 1]:
            w = model[0].weight.data.squeeze().numpy().copy()
            snapshots[epoch] = w

    return snapshots

for m in [20, 200, 2000]:
    snaps = train_and_record(m)
    w_final = snaps[max(snaps.keys())]
    print(f"m={m:4d}: 最终权重标准差={w_final.std():.3f}, "
          f"取值范围=[{w_final.min():.2f}, {w_final.max():.2f}]")

当宽度 $$m = 2000$$ 时，权重直方图已几乎无法与一条光滑的核密度估计（KDE）曲线区分——平均场极限在此尺度下已清晰显现。标准差持续收窄的现象，正反映了平均场机制的本质：整体权重幅值以 $O(1/\sqrt{m})$ 的速率衰减，而单个权重相对于其初始值的移动量仍保持 $$O(1)$$ 量级。

平均场方程的推导#

\mathcal{L}[\rho] = \frac{1}{n} \sum_{k=1}^n \ell(f_\rho(x_k), y_k), \quad f_\rho(x) = \int a \, \sigma(w^\top x + b)\, d\rho(\theta) .

\dot \theta_i = -\nabla_{\theta_i} \frac{\delta \mathcal{L}}{\delta \rho}\bigg|_{\rho = \rho_t^m}(\theta_i).

\partial_t \rho + \nabla_\theta\cdot(\rho\, v_t) = 0, \qquad v_t(\theta) = -\nabla_\theta \frac{\delta \mathcal{L}}{\delta \rho}[\rho_t](\theta) .

这就是平均场 PDE：一个确定性的、非线性的神经元密度演化方程。

全局收敛性#

定理（Mei-Montanari-Nguyen 2018, Chizat-Bach 2018，精神上）：在以下条件下，平均场方程可驱动损失趋于零： (a) 充分过参数化（或在极限下 $\rho_0$ 的支撑足够大）； (b) 数据上的神经正切核（NTK）正定； (c) 合理的初始化（例如高斯分布）。

三行证明：在适当 regime 下，网络在初始化附近线性化，残差 $r_t(x) = f_{\rho_t}(x) - y(x)$ 满足 $\dot r = -K r$ ，其中 $K \succeq \lambda_{\min}(K) I$ ，因此 $\|r_t\|^2 \leq e^{-\lambda_{\min}(K) t}\|r_0\|^2$ 。

NTK 与平均场对比：两层网络存在两种缩放 regime：

NTK / lazy regime（Jacot-Gabriel-Hongler 2018）： $m \to \infty$ 且学习率固定；参数几乎不动；网络行为类似核方法。
平均场 regime：学习率随 $$m$$ 缩放，使参数移动量为 $$O(1)$$ ；动力学本质上是非线性的，并表现出特征学习。

Wasserstein 梯度流形式#

Wasserstein 梯度流：密度向 Gibbs 平衡态演化及 KL 散度衰减

\mathcal{L}[\rho] = \tfrac12 \int K(\theta, \theta')\, d\rho(\theta)\, d\rho(\theta') + \int g(\theta)\, d\rho(\theta) ,

平均场方程恰好是 $\mathcal{L}$ 在 $\mathcal{P}_2$ 上的 Wasserstein 梯度流。这是“训练即 $\mathcal{P}_2$ 上梯度流”这一口号最清晰的理论依据：尽管损失在参数空间 $\theta$ 中是非凸的，但在测度空间中可能是位移凸的，从而保证全局收敛。

图 5. 能量是非凸的且存在多个盆地，但梯度流仍单调下降。这正是高维神经网络训练应有的直观图像——只不过在平均场极限下，参数 $\theta$ 的角色被整个密度 $\rho_t$ 所取代。

深度网络的连续时间极限#

h_{\ell + 1} = h_\ell + \tau\, F(h_\ell, \theta_\ell)

\min_F \int c(x, F(x))\, d\rho_X(x), \quad \text{约束 } F_\# \rho_X = \rho_Y ,

即学习一个最优映射，逐步将数据分布展平，使其易于被分类器头读取。

变分推断与 ELBO#

\mathrm{ELBO}(q) = \mathbb{E}_{q(z|x)}[\log p(x|z)] - \mathrm{KL}\!\left(q(z|x) \,\Vert\, p(z)\right) ,

这等价于最小化 $\mathrm{KL}(q(z|x)\,\Vert\,p(z|x))$ ——即到真实后验的 KL 距离。这是一个带约束的变分问题；当 $\mathcal{Q}$ 包含所有密度时，最优解即为精确后验。

在前文的 Wasserstein 视角下，KL 散度 $\mathrm{KL}(\cdot\,\Vert\,p)$ 的梯度流正是第 2.4 节所述的 Fokker-Planck 方程，这也解释了为何扩散生成模型与 VAE 可用相同工具分析。

小型 VAE 在四模高斯混合上的 ELBO 分解：左图为重构损失 vs KL 正则项，中图为数据与重构对照，右图为隐空间编码与标准正态先验等高线。

图 6. 重构项和 KL 项作用方向相反：前者要求 $\hat x \approx x$ ，后者要求 $q(z|x) \approx p(z)$ 。二者之和（即负 ELBO）构成变分目标。隐空间图显示，四个数据模态已被组织成标准正态先验内的四个清晰分离的簇。

实现：ELBO 的分解
ELBO 在重构质量与后验分布正则化之间进行权衡。下面是一段极简的代码，直观展示了这一权衡过程：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleVAE(nn.Module):
    def __init__(self, data_dim=2, latent_dim=2, hidden=64):
        super().__init__()
        # 编码器主干网络：输入 → 隐藏层 → 隐藏层
        self.enc = nn.Sequential(nn.Linear(data_dim, hidden), nn.ReLU(),
                                 nn.Linear(hidden, hidden), nn.ReLU())
        # 输出隐变量均值 μ
        self.mu = nn.Linear(hidden, latent_dim)
        # 输出隐变量对数方差 log σ²
        self.logvar = nn.Linear(hidden, latent_dim)
        # 解码器：隐变量 → 隐藏层 → 隐藏层 → 重构数据
        self.dec = nn.Sequential(nn.Linear(latent_dim, hidden), nn.ReLU(),
                                 nn.Linear(hidden, hidden), nn.ReLU(),
                                 nn.Linear(hidden, data_dim))

    def forward(self, x):
        h = self.enc(x)
        mu, logvar = self.mu(h), self.logvar(h)
        # 重参数化采样：z = μ + σ × ε，其中 ε ~ N(0, I)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)
        recon = self.dec(z)
        # 重构损失（均方误差，按 batch 归一化）
        recon_loss = F.mse_loss(recon, x, reduction='sum') / x.shape[0]
        # KL 散度项（相对于标准正态先验 N(0, I)）
        kl = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp()) / x.shape[0]
        return recon_loss, kl, recon_loss + kl

centers = torch.tensor([[2,2],[-2,2],[2,-2],[-2,-2]], dtype=torch.float)
data = centers[torch.randint(4, (2000,))] + 0.3 * torch.randn(2000, 2)

vae = SimpleVAE()
opt = torch.optim.Adam(vae.parameters(), lr=1e-3)
for epoch in range(2000):
    recon, kl, loss = vae(data)
    opt.zero_grad(); loss.backward(); opt.step()
    if epoch % 500 == 0:
        print(f"Epoch {epoch:4d}: recon={recon:.3f}, KL={kl:.3f}, ELBO={loss:.3f}")

数值变化清晰地揭示了这种权衡：随着重构误差持续下降，KL 散度项却不断上升——这是因为编码器正主动将隐变量分布“推离”标准正态先验，以在隐空间中显式区分不同数据簇（即编码聚类身份）。因此，ELBO 初期快速下降（重构提升带来的收益压倒 KL 惩罚），但后期 KL 项逐渐占据主导。这种张力正是变分原理的本质体现：最优的后验近似 $$q(z|x)$$ 必须在保真度（忠实重建输入）与正则性（贴近先验约束）之间取得平衡——恰如 Ritz 泛函在边界拟合精度与解的光滑性之间寻求最优折衷。

数值验证#

我们现在汇总支持上述理论的实验结果。所有实验均可通过脚本 scripts/figures/pde-ml/03-variational.py 复现。

最速降线（图 1）#

五条候选下降曲线连接 $$A=(0,0)$$ 和 $B=(\pi, 2)$ 。我们用梯形法则离散时间泛函 $T[y] = \int \sqrt{(1+y'^2)/(2gy)}\, dx$ ，并按下降时间排序。摆线以明显优势胜出，击败了“最短路径”（直线）和更陡峭的“先降后滑”方案。这张图浓缩了变分法的全部精髓：不同的容许曲线对应不同的泛函值，而只有使第一变分为零的曲线才是极小值点。

第一变分论证（图 2）#

我们选取区间 $$[0, 1]$$ 上 Dirichlet 能量的极值函数 $y(x) = \sin(\pi x)$ （满足 $$y(0) = y(1) = 0$$ ）和一个在边界为零的光滑扰动 $\eta(x) = \sin(2\pi x)\, x(1-x)$ 。绘制 $J(\varepsilon) = J[y + \varepsilon \eta]$ 关于 $\varepsilon$ 的曲线，可见其为一条在 $\varepsilon = 0$ 处切线水平的抛物线——这正是 Euler-Lagrange 方程的几何内涵。

Wasserstein 梯度流（图 3）#

\partial_t \rho = \nabla\cdot(\rho \nabla V) + \Delta \rho, \qquad V(x) = \tfrac12 (x - \mu_\ast)^2 ,

初始密度为远离目标的尖锐高斯分布。模拟同时展示了三点：

$\rho_t$ 在测度空间中向 Gibbs 目标 $\rho_\ast = \mathcal{N}(\mu_\ast, 1)$ 移动；
自由能 $F[\rho_t]$ 单调递减（梯度流的能量耗散特性）；
均值相对于 $\mu_\ast$ 的漂移也单调衰减——可视为“Wasserstein 距离的代理指标”。

平均场极限（图 4）#

一个两层 ReLU 网络（宽度 $m \in \{20, 200, 2000\}$ ）被训练以拟合 $f_\ast(x) = \sin(\pi x) + 0.5 \sin(3\pi x)$ （基于 64 个采样点）。在训练 0、200 和 1500 轮后的第一层权重直方图随 $$m$$ 增大而明显更光滑；当 $$m = 2000$$ 时，直方图已是对连续密度的良好采样。右列显示相应损失的衰减——网络越宽收敛越快，符合平均场过参数化理论。学习率按 $\propto 1/\sqrt{m}$ 缩放，以保证每个粒子的漂移量为 $$O(1)$$ ，这正是平均场 regime 的正确缩放。

能量景观（图 5）#

一个刻意设计的二维非凸能量（含倾斜和三角函数“凸起”）展示了普适的耗散恒等式。用 RK45 积分得到的梯度流轨迹最终落入最低能量极小值的吸引域；能量时间序列证实了单调递减。将此图景扩展到百万维，便是神经网络训练的心理模型——只不过真实损失景观通常具有更平坦的极小值（这一特性既解释了泛化能力，也说明了随机梯度方法的成功）。

另外两种几何：Fisher-Rao 与 Adam#

Fisher-Rao 与自然梯度#

I(\theta) = \mathbb{E}_{p_\theta}\!\left[\nabla_\theta \log p_\theta(x)\, \nabla_\theta \log p_\theta(x)^\top\right]

定义了另一种度量——Fisher-Rao 度量——对应的自然梯度流 $\dot\theta = -I(\theta)^{-1} \nabla_\theta J(\theta)$ （Amari 1998）通常比普通梯度下降收敛快得多，因为它考虑了参数空间的内蕴曲率。Wasserstein 与 Fisher-Rao 的视角截然不同：

Wasserstein 衡量传输成本；适用于分布支撑移动的情形。
Fisher-Rao 衡量信息几何距离；适用于分布在原地变形的情形。

近期关于 Kernel Approximation of Fisher-Rao Gradient Flows 的工作展示了如何数值逼近 Fisher-Rao 流，并据此设计新的 Langevin 式采样算法。

Adam 作为重参数化的梯度流#

\dot\theta = -\frac{m_t}{\sqrt{v_t} + \epsilon}, \qquad \dot m_t = \alpha_1\big(\nabla J(\theta) - m_t\big), \qquad \dot v_t = \alpha_2\big(\|\nabla J(\theta)\|^2 - v_t\big) ,

这实际上是 $$J$$ 在一个坐标相关的黎曼度量 $g_{ii}(\theta) = (\sqrt{v_i(\theta)} + \epsilon)$ 下的梯度流。这是一种对自然梯度的对角近似——自适应学习率本质上是一种度量变换。

Hamilton 流 vs. 梯度流#

一个微妙但关键的点：并非所有流都是梯度流。同一个能量函数可能关联两种性质截然相反的动力学。

同一个能量 $H$
的两种流：左为 Hamilton 流，沿闭轨保持能量守恒；右为梯度流，耗散能量并收敛到唯一极小值。

图 7. 同一能量下的两种流。左侧，辛向量场 $\dot q = H_p, \dot p = -H_q$ 旋转相空间并保持能量不变——形成闭合轨道，无收敛性。右侧，梯度场 $\dot q = -H_q, \dot p = -H_p$ 将所有轨迹收缩至原点。两种图景都很重要：梯度流描述优化过程；辛流描述保守动力学，并启发了“保结构”的 Neural ODE 架构（本系列第五部分的主题）。

串联起来：从最小作用量到机器学习#

变分思想贯穿全文。让我们再梳理一遍，看看各个部分是如何紧密咬合的：

层级	变分问题	状态空间	度量	导出的偏微分方程（PDE）/常微分方程（ODE）
经典力学	极小化作用量 $$S[q]$$	轨迹空间	—	欧拉-拉格朗日方程 / 哈密顿方程
偏微分方程理论	极小化狄利克雷能量 $$J[u]$$	函数空间	$$L^2$$ 范数	拉普拉斯方程 / 热传导方程
最优传输	在能量泛函 $\mathcal{E}[\rho]$ 上执行 JKO 邻近步	概率测度空间 $\mathcal{P}_2$	二阶 Wasserstein 距离 $$W_2$$	福克-普朗克方程 / 多孔介质方程
神经网络训练	极小化经验风险 $\hat{R}[\rho]$	神经元分布空间 $\mathcal{P}_2$	二阶 Wasserstein 距离 $$W_2$$	平均场偏微分方程
变分推断	极小化 KL 散度 $\mathrm{KL}(q \mid p)$	近似后验分布族 $\mathcal{Q}$	KL 散度或 $$W_2$$ 距离	朗之万随机微分方程（SDE）/ 福克-普朗克方程
Adam 优化器	函数 $$J$$ 的梯度流	参数空间 $\mathbb{R}^n$	自适应对角度量	黎曼流形上的梯度流

每一行都对应一个变分问题；而每一行又通过在相应几何结构下沿能量泛函的最速下降方向演化，导出一个动力学系统。真正变化的，仅是状态空间、所采用的度量以及待极小化的能量泛函。正因如此，PDE 视角才如此强大：一旦你识别出背后的变分结构，收敛性证明、数值求解方案乃至算法设计原则，便能在所有这些场景中自然迁移与复用。

面对任何新算法，你只需问三个关键问题：

它在极小化什么能量？（目标函数、证据下界 ELBO、经验风险、自由能等）
在何种度量下进行？（欧氏度量、Wasserstein 度量、Fisher-Rao 度量、自适应度量等）
其对应的梯度流满足哪一类 PDE？（热方程、福克-普朗克方程、Vlasov 方程……）

只要你能清晰回答这三个问题，收敛性理论几乎就已水到渠成。

展望#

PDE / 变分视角当前最活跃的几个方向包括：

理论：非凸损失的更精细收敛速率；有限宽度修正；注意力层的 PDE 理论。
算法：高阶 ODE/PDE 求解器作为新型优化器；基于控制理论的超参数调度；Wasserstein 感知的 MCMC 方法。
应用：扩散生成模型是反向 Fokker-Planck 方程；策略梯度是策略空间上的梯度流；Deep Ritz 和 PINN 将 PDE 求解转化为变分问题。
交叉融合：自旋玻璃类比、用于端到端优化的 Pontryagin 最大值原理、深度学习中的信息几何与辛几何。

总结#

从变分法出发，我们构建了一个以 PDE 为核心的神经网络优化图景。泛函与 Euler-Lagrange 方程将离散优化与连续动力学联系起来；Wasserstein 几何为概率测度空间提供了自然度量，并将热方程、Fokker-Planck 方程、多孔介质方程和 Keller-Segel 方程统一为显式能量的梯度流；平均场极限将有限宽度神经网络的训练简化为 Vlasov 型 PDE，并在合理假设下得出全局收敛性；ELBO、Adam 和 Hamilton 动力学均被纳入同一变分框架。数值实验——最速降线、一维 Fokker-Planck 模拟、两层网络的宽度扫描、非凸能量景观和小型 VAE——逐一验证了理论的关键预测。

PDE 视角仍处于早期阶段。随着数学与机器学习的持续融合，它必将为优化器设计、泛化分析和收敛性证明提供日益锐利的工具。

参考文献#

L. Chizat and F. Bach, “On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport,” NeurIPS, 2018. arXiv:1805.09545
S. Mei, A. Montanari, P.-M. Nguyen, “A Mean Field View of the Landscape of Two-Layer Neural Networks,” PNAS, 2018. arXiv:1804.06561
G. M. Rotskoff and E. Vanden-Eijnden, “Neural Networks as Interacting Particle Systems,” arXiv:1805.00915 , 2018.
A. Jacot, F. Gabriel, C. Hongler, “Neural Tangent Kernel: Convergence and Generalization in Neural Networks,” NeurIPS, 2018. arXiv:1806.07572
W. E and B. Yu, “The Deep Ritz Method,” CPAM, 2018. arXiv:1710.00211
R. T. Q. Chen, Y. Rubanova, J. Bettencourt, D. Duvenaud, “Neural Ordinary Differential Equations,” NeurIPS, 2018. arXiv:1806.07366
L. Ambrosio, N. Gigli, G. Savaré, Gradient Flows in Metric Spaces and in the Space of Probability Measures, Birkhäuser, 2008.
C. Villani, Optimal Transport: Old and New, Springer, 2009.
R. Jordan, D. Kinderlehrer, F. Otto, “The Variational Formulation of the Fokker-Planck Equation,” SIAM J. Math. Anal., 1998.
F. Otto, “The Geometry of Dissipative Evolution Equations: the Porous Medium Equation,” Comm. PDE, 2001.
Y. Lu and J. Lu, “Mean-Field Analysis of Neural SGD-Ascent ,” 2024.
A. Kazeykina and M. Fornasier, “Kernel Approximation of Fisher-Rao Gradient Flows ,” 2024.
D. Onken et al., “Deep ResNets and Conditional Optimal Transport ,” 2024.
M. Belkin, D. Hsu, S. Ma, S. Mandal, “Reconciling Modern Machine Learning Practice and the Classical Bias-Variance Trade-off,” PNAS, 2019.
S. Amari, “Natural Gradient Works Efficiently in Learning,” Neural Computation, 1998.
D. P. Kingma and J. Ba, “Adam: A Method for Stochastic Optimization,” ICLR, 2015. arXiv:1412.6980
L. Chizat, E. Oyallon, F. Bach, “On Lazy Training in Differentiable Programming,” NeurIPS, 2019. arXiv:1812.07956
G. Peyré and M. Cuturi, “Computational Optimal Transport,” Foundations and Trends in Machine Learning, 2019. arXiv:1803.00567
J. Sirignano and K. Spiliopoulos, “Mean Field Analysis of Neural Networks: A Central Limit Theorem,” Stoch. Proc. Appl., 2020. arXiv:1808.09372

可复现性说明： 本文全部七张图由站点仓库中的 scripts/figures/pde-ml/03-variational.py 生成。在项目根目录执行 python scripts/figures/pde-ml/03-variational.py 即可，PNG 文件会同时写入 EN 和 ZH 的资源目录。

偏微分方程与机器学习（三）：变分原理与优化

变分法基础：从泛函到 Euler-Lagrange 方程#

泛函与第一变分#

最速降线问题（Brachistochrone）#

Hamilton 原理与辛几何视角#

从泛函导数到梯度流#

梯度流与 Wasserstein 几何#

$\mathbb{R}^n$ 中的梯度流#

Wasserstein 距离#

JKO 格式： $\mathcal{P}_2(\mathbb{R}^d)$ 上的梯度流#

热方程作为熵的梯度流#

神经网络训练的平均场理论#

从有限宽度到无限宽度#

平均场方程的推导#

全局收敛性#

Wasserstein 梯度流形式#

深度网络的连续时间极限#

变分推断与 ELBO#

数值验证#

最速降线（图 1）#

第一变分论证（图 2）#

Wasserstein 梯度流（图 3）#

平均场极限（图 4）#

能量景观（图 5）#

另外两种几何：Fisher-Rao 与 Adam#

Fisher-Rao 与自然梯度#

Adam 作为重参数化的梯度流#

Hamilton 流 vs. 梯度流#

串联起来：从最小作用量到机器学习#

最新进展与开放问题#

展望#

总结#

参考文献#

PDE × 机器学习 8 篇

读有所得？

变分法基础：从泛函到 Euler-Lagrange 方程#

泛函与第一变分#

最速降线问题（Brachistochrone）#

Hamilton 原理与辛几何视角#

从泛函导数到梯度流#

梯度流与 Wasserstein 几何#

$\mathbb{R}^n$ 中的梯度流#

Wasserstein 距离#

JKO 格式：$\mathcal{P}_2(\mathbb{R}^d)$ 上的梯度流#

热方程作为熵的梯度流#

神经网络训练的平均场理论#

从有限宽度到无限宽度#

平均场方程的推导#

全局收敛性#

Wasserstein 梯度流形式#

深度网络的连续时间极限#

变分推断与 ELBO#

数值验证#

最速降线（图 1）#

第一变分论证（图 2）#

Wasserstein 梯度流（图 3）#

平均场极限（图 4）#

能量景观（图 5）#

另外两种几何：Fisher-Rao 与 Adam#

Fisher-Rao 与自然梯度#

Adam 作为重参数化的梯度流#

Hamilton 流 vs. 梯度流#

串联起来：从最小作用量到机器学习#

最新进展与开放问题#

展望#

总结#

参考文献#

PDE × 机器学习 8 篇

读有所得？

JKO 格式： $\mathcal{P}_2(\mathbb{R}^d)$ 上的梯度流#