4. 自动微分原理

自动微分（Automatic Differentiation, AD）是 PINN（物理信息神经网络）及所有现代深度学习框架的数学基石。与符号微分（产生复杂表达式）和数值微分（存在截断误差）不同，自动微分通过将复合函数分解为一系列基本运算，并精确应用链式法则，在机器精度下计算导数。理解其底层机制，对于调试 PINN 训练中的梯度消失/爆炸、设计自定义损失函数以及优化计算性能至关重要。本章将从链式法则出发，剖析前向与反向两种模式的计算逻辑，并深入 PyTorch 与 JAX 的底层实现，揭示其设计哲学与工程权衡。

上图展示了自动微分在计算图上的两种传播路径：蓝色实线为前向计算，红色虚线为反向梯度回传。在 PINN 中，损失函数通常为标量（如 PDE 残差的均方误差），因此反向模式（即反向传播）是默认选择。下面我们将逐步拆解其数学原理与工程实现。

4.1 链式法则与计算图：自动微分的数学骨架

核心原理：自动微分的本质是将复合函数 y = f(g(h(x))) 分解为一系列原子操作（如加减乘除、指数、三角函数等），并构建一个有向无环图（DAG）。每个节点代表一个中间变量，每条边代表一个基本运算。链式法则告诉我们：

dy/dx = (dy/dvₙ) × (dvₙ/dvₙ₋₁) × ... × (dv₁/dx)

其中 vᵢ 为中间变量。计算图的价值在于将这一连串乘法结构化，使得我们可以选择从输入侧（前向）或输出侧（反向）开始累积这些局部导数。

判断方法：在 PINN 中，判断使用哪种模式的关键在于维度分析。设输入维度为 n（如网络权重数量），输出维度为 m（通常为标量损失）。前向模式每次计算 Jacobian 的一列（即所有输出对单个输入的偏导），需要 n 次前向传播；反向模式每次计算 Jacobian 的一行（即单个输出对所有输入的偏导），仅需 m 次反向传播。由于 PINN 的损失函数几乎总是标量（m=1），反向模式的计算复杂度为 O(1) 次原函数代价，而前向模式为 O(n) 次，因此反向模式占据主导。

4.2 前向模式：从输入侧累积导数

操作流程：前向模式的核心是引入一个"种子"向量 ẋ，表示输入方向的导数。对于每个中间变量 vᵢ，同时计算其值 vᵢ 和导数 v̇ᵢ = ∂vᵢ/∂x · ẋ。以函数 y = sin(x₁) + x₁·x₂ 为例，若设置种子 ẋ = (1, 0)（即只对 x₁ 求导）：

v₁ = sin(x₁)      → v̇₁ = cos(x₁)·1
v₂ = x₁·x₂        → v̇₂ = 1·x₂ + x₁·0 = x₂
y  = v₁ + v₂      → ẏ = v̇₁ + v̇₂ = cos(x₁) + x₂

整个过程与原始计算同步进行，无需存储中间结果。在 PINN 中，前向模式可用于高效计算单个输入点上的高阶导数（通过嵌套），但很少用于权重梯度。

案例拆解：考虑一个简单的 PINN 残差计算：f(x) = u_net(x) - u_exact(x)，其中 u_net 是神经网络。若我们需要 ∂f/∂x（x 为空间坐标），且网络输入维度为 2（x, y），则前向模式只需 2 次传播即可获得完整的 Jacobian 行。而反向模式需要 1 次反向传播获得标量梯度，但若输出维度为 1，两者效率相近。实际中，JAX 的 jax.jacfwd 在输入维度小于输出维度时自动选择前向模式。

4.3 反向模式：从输出侧回传梯度

核心原理：反向模式（即反向传播）分为两个阶段：前向阶段计算所有中间变量值并存储；反向阶段从输出开始，利用链式法则逐层回传梯度。定义 v̄ᵢ = ∂y/∂vᵢ，则对于节点 vₖ = vᵢ + vⱼ，有 v̄ᵢ = v̄ₖ · 1；对于 vₖ = vᵢ · vⱼ，有 v̄ᵢ = v̄ₖ · vⱼ。这一过程本质上是计算图上的动态规划。

操作流程：以 y = e^(sin(x)·x²) 为例，前向阶段记录：v₁=x², v₂=sin(x), v₃=v₁·v₂, v₄=e^v₃。反向阶段从 v̄₄=1 开始：

v̄₃ = v̄₄ · e^v₃ = e^v₃
v̄₂ = v̄₃ · v₁ = e^v₃ · x²
v̄₁ = v̄₃ · v₂ = e^v₃ · sin(x)
x̄  = v̄₁·2x + v̄₂·cos(x) = e^v₃·(2x·sin(x) + x²·cos(x))

关键点在于：反向阶段需要用到前向阶段存储的 v₁, v₂, v₃ 值，因此内存开销与网络深度成正比。这也是 PyTorch 中 backward() 会释放中间变量的原因——默认在反向传播后释放计算图以节省内存。

案例拆解（PINN 损失梯度）：在 PINN 中，损失函数通常为 L = λ₁·L_PDE + λ₂·L_BC，其中 L_PDE = ||∂u/∂t + u·∂u/∂x - ν·∂²u/∂x²||²。反向模式需要计算 ∂L/∂θ（θ 为网络权重）。前向阶段计算 u 及其各阶导数（通过嵌套自动微分），反向阶段将损失梯度逐层回传至每个权重。PyTorch 的 autograd 会自动处理这一过程，但需要注意：∂²u/∂x² 的计算需要两次前向模式的自动微分（或一次反向+一次前向），这会产生额外的计算图节点。

4.4 PyTorch 与 JAX 的自动微分实现机制对比

PyTorch：动态计算图（Define-by-Run）

PyTorch 在每次前向传播时动态构建计算图。每个张量包含 data、grad、grad_fn 三个关键属性。当执行 loss.backward() 时，PyTorch 从 loss 节点出发，沿 grad_fn 的引用链反向遍历，调用每个节点的 backward() 方法。其核心优势是灵活性——可以在循环中动态改变网络结构，非常适合 PINN 中需要根据迭代次数调整损失权重的场景。

import torch
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = torch.sin(x[0]) + x[0]*x[1]
y.backward()
print(x.grad)  # tensor([2.0, 1.0])

JAX：静态计算图（Define-and-Run）

JAX 采用函数式编程范式，通过 jax.grad、jax.jit、jax.vmap 等变换器操作纯函数。其核心是 jaxpr（JAX 表达式），在首次调用时通过追踪（tracing）将 Python 函数编译为静态计算图，然后由 XLA 编译器优化执行。JAX 的优势在于：jit 编译后执行速度极快，且 grad 可以嵌套使用（如 jax.grad(jax.grad(f)) 计算二阶导数），非常适合 PINN 中需要高阶导数的场景。

import jax.numpy as jnp
from jax import grad, jit

def loss_fn(params, x, t):
    u = predict(params, x, t)
    return jnp.mean((u - u_exact)**2)

grad_fn = jit(grad(loss_fn))  # 编译并返回梯度函数
grads = grad_fn(params, x, t)

常见误区与总结复盘：

• 误区一：认为自动微分与数值微分等价。实际上，自动微分是精确的（机器精度），而数值微分（如有限差分）存在截断误差和舍入误差，在 PINN 中会导致梯度噪声，影响收敛。
• 误区二：在 PyTorch 中重复调用 backward() 而不清零梯度。默认情况下梯度会累积，需在每次迭代前调用 optimizer.zero_grad()。
• 误区三：在 JAX 中对带有 Python 控制流的函数使用 jit 时，控制流会被静态展开，可能导致计算图过大。应使用 jax.lax.cond 或 jax.lax.scan 保持动态性。
• 误区四：忽略内存峰值。在 PINN 中计算二阶导数时，PyTorch 会创建新的计算图节点，导致内存翻倍。建议使用 torch.autograd.grad 的 create_graph=False 参数来避免不必要的图保留。