从零开始的自动微分：PyTorch如何在物理信息神经网络中计算梯度

# 从零开始的自动微分：PyTorch 在物理信息神经网络中如何计算梯度
来源: https://arxiv.org/html/2607.13042
Abdeladhim Tahimi CECA, 阿拉戈斯联邦大学 (UFAL), 里奥拉戈, 阿拉戈斯, 57100-000, 巴西 [email protected]

###### 摘要

本文通过显式的数值，追踪了 PyTorch 的自动微分（AD）引擎在物理信息神经网络（PINN）训练中如何计算梯度——这是一个需要两级微分的场景：通过神经网络计算物理导数 \(\hat{y}'(t)=d\hat{y}/dt\)，以及计算依赖于 \(\hat{y}'(t)\) 的损失函数 \(L\) 的参数梯度 \(\nabla_\theta L\)。我们使用一个 1-3-3-1 多层感知机和初值问题 \(y'(t)+y(t)=0\)，\(y(0)=1\)（源自 [17](https://arxiv.org/html/2607.13042#bib.bib29)），在每一个节点上追踪完整的流程：前向传播过程中建立的计算图，反向模式反向遍历（一次遍历计算全部 22 个参数梯度），以及通过 `create_graph=True` 实现的“图上图”机制，该机制能够正确地对物理信息残差进行微分。每个伴随值都根据 [17](https://arxiv.org/html/2607.13042#bib.bib29) 的手工推导进行了验证，将 \(P/Q\) 敏感性框架与 PyTorch 的 autograd 引擎使用的向量-雅可比积联系起来。

**关键词**：自动微分，反向模式自动微分，反向传播，计算图，向量-雅可比积，物理信息神经网络，PyTorch

**MSC 代码**：65D25, 65-01, 68T07, 65L05

## 1 引言

一个典型的物理信息神经网络（PINN）训练循环通常只需不到二十行 PyTorch 代码。前向传播计算网络输出 \(\hat{y}(t;\theta)\)；调用 `torch.autograd.grad`（[https://arxiv.org/html/2607.13042v1/torch.autograd.grad](https://arxiv.org/html/2607.13042v1/torch.autograd.grad)）并设置 `create_graph=True` 来计算时间导数 \(\hat{y}'(t)\)；ODE 残差 \(R = \hat{y}' + \hat{y}\) 的平方形成损失函数；然后 `loss.backward()`（[https://arxiv.org/html/2607.13042v1/loss.backward()](https://arxiv.org/html/2607.13042v1/loss.backward())）在一次调用中生成所有参数梯度。整个流程依赖于自动微分——然而，这些调用实际上分别计算了什么？为什么每个标志都是必要的？回答这些问题需要打开这个黑盒。

##### 两个微分任务。

训练一个 PINN 需要两个本质不同的导数。考虑一阶 ODE \(y'(t) + y(t) = 0\)，初始条件 \(y(0) = 1\)。PINN 在配置点 \(t_c\) 处的损失为

\[
L = \underbrace{\left( \hat{y}(t_c;\theta) + \frac{\partial \hat{y}}{\partial t}\bigg|_{t_c} \right)^2}_{L_R:\; \text{ODE 残差}} + \lambda \underbrace{\bigl( \hat{y}(0;\theta) - 1 \bigr)^2}_{L_{IC}:\; \text{初始条件}}.
\tag{1}
\]

该损失包含*物理导数* \(\partial \hat{y} / \partial t\)——即当输入 \(t\) 变化而所有参数固定时网络输出的变化——而训练需要第二个导数，即*训练梯度* \(\nabla_\theta L\)，它描述了当参数 \(\theta_k\) 被调整时损失如何变化。我们用 \(\hat{y}'\) 表示 \(d\hat{y}/dt\)；由于参数 \(\theta\) 不依赖于 \(t\)，这与 \(\partial \hat{y} / \partial t\) 是一致的。

在标准的监督学习中，损失只依赖于 \(\hat{y}\)，一次反向传播就能计算 \(\nabla_\theta L\)。而此处损失同时依赖于 \(\hat{y}\) 和 \(\hat{y}'\)，因此残差损失的训练梯度为

\[
\frac{\partial L_R}{\partial \theta_k} = 2R\left( \frac{\partial \hat{y}}{\partial \theta_k} + \frac{\partial \hat{y}'}{\partial \theta_k} \right).
\tag{2}
\]

第二项 \(\partial \hat{y}' / \partial \theta_k\) 是 PINN 特有的复杂之处：它是一个二阶混合导数 \(\partial^2 \hat{y} / \partial \theta_k \partial t\)，需要对物理导数计算本身相对于参数进行微分。在 PyTorch 中，只有在物理导数计算时设置了 `create_graph=True`，这一项才会存在；省略该标志会悄无声息地丢掉它，产生错误的梯度而不报任何错误。

这个 ODE 是物理信息学习的“hello world”——足够小以允许闭式分析，同时又足够丰富以展现上述两个导数之间的耦合。配套论文 [17](https://arxiv.org/html/2607.13042#bib.bib29) 针对应用于该问题的 1-3-3-1 MLP，给出了训练梯度的分析推导；而本文则追踪 PyTorch AD 引擎内部对应的机器计算。

##### 从手工推导的梯度到框架的自动微分。

训练神经网络以满足微分方程的工作早于现代深度学习。Dissanayake 和 Phan-Thien [4](https://arxiv.org/html/2607.13042#bib.bib7) 以及 Lagaris 等人 [10](https://arxiv.org/html/2607.13042#bib.bib8) 的早期工作针对具有固定激活函数的浅层架构，分析推导了训练所需的梯度表达式；每当架构改变时，这些公式就必须重新推导。深度学习框架中通用自动微分的出现——尤其是 PyTorch [11](https://arxiv.org/html/2607.13042#bib.bib3) 所提供的反向模式 AD——消除了这一限制，它能够通用地计算物理导数（如 \(\partial \hat{y} / \partial t\)）和参数梯度 \(\nabla_\theta L\)，无论网络深度或激活函数如何选择。这一转变使得深度架构变得可行，并催生了 Raissi 等人 [12](https://arxiv.org/html/2607.13042#bib.bib11) 的现代 PINN 框架（其中边界条件或初始条件通过额外的损失项来强制执行），以及并行的神经网络 PDE 求解器，如 Deep Galerkin 方法 [15](https://arxiv.org/html/2607.13042#bib.bib10) 和 Deep Ritz 方法 [5](https://arxiv.org/html/2607.13042#bib.bib9)。自 2019 年以来，PINN 已被应用于流体动力学 [13](https://arxiv.org/html/2607.13042#bib.bib12)、反问题 [8](https://arxiv.org/html/2607.13042#bib.bib13) 以及广泛的科学计算任务 [3](https://arxiv.org/html/2607.13042#bib.bib14)，[18](https://arxiv.org/html/2607.13042#bib.bib18)]。

##### 教学上的空白。

自动微分的理论基础已经非常完善。Griewank 和 Walther [7](https://arxiv.org/html/2607.13042#bib.bib1) 给出了前向和反向模式的权威论述；Baydin 等人 [1](https://arxiv.org/html/2607.13042#bib.bib2) 面向机器学习读者综述了该领域；标准教科书 [6](https://arxiv.org/html/2607.13042#bib.bib5) 涵盖了反向传播——反向模式 AD 在神经网络训练中的应用——其源头可追溯至 Rumelhart 等人 [14](https://arxiv.org/html/2607.13042#bib.bib4) 的工作。在 PINN 方面，诸如 Katsikis 等人 [9](https://arxiv.org/html/2607.13042#bib.bib19) 的入门教程解释了损失公式和训练方法，而经过同行评议的教学性推导工作，如 Blechschmidt 和 Ernst [2](https://arxiv.org/html/2607.13042#bib.bib20)，则通过链式法则在单隐藏层网络上对残差进行符号化推导——但止步于 AD 引擎本身。目前缺少的是分析推导与机器计算之间的桥梁：一份文档，能在单个操作和显式数值的层面上，追踪 AD 引擎如何处理多层网络上公式 (1) 和 (2) 中的两个导数耦合。具体来说：

- PyTorch 在前向传播过程中构建了什么样的数据结构？
- 一次反向遍历如何同时计算所有 22 个参数梯度？
- `create_graph=True` 到底做了什么，为什么省略它会导致静默的错误梯度？
- PyTorch 传播的伴随向量与 [17](https://arxiv.org/html/2607.13042#bib.bib29) 手工推导的 \(P/Q\) 敏感性如何关联？

现有参考文献涉及了这个故事的各个部分——AD 理论没有 PINN 细节，PINN 公式没有 AD 内部机制——但没有一个能在一个可验证的实例上把它们整合成一个完整的叙述。弥合这一空白可以防止静默的实现错误，澄清物理信息训练的计算成本，并为开发改进收敛性和准确性的技术提供基础。

##### 本文涵盖的内容。

我们追踪了与 [17](https://arxiv.org/html/2607.13042#bib.bib29) 中使用的相同的 1-3-3-1 MLP 和相同的初值问题 \(y' + y = 0\)，\(y(0) = 1\) 的完整 AD 流程。这个例子足够小，以至于每个中间量——前向传播的激活值、反向传播的 22 个参数梯度、以及由 `create_graph=True` 创建的扩展图——都可以显式地展示出来；同时又足够大，使得公式 (1) 和 (2) 中两个导数的耦合能够以完全一般的形式出现：

- 第 2 节 在一个最小的两输入复合函数上介绍切向量和伴随向量。
- 第 3 节 追踪 1-3-3-1 MLP 的前向传播并列举计算图。
- 第 4 节 逐节点追踪反向传播，在一次遍历中计算全部 22 个参数梯度。
- 第 5 节 解释“图上图”：`create_graph=True` 如何实现正确的训练梯度，包括乘积法则和 \(\phi''\) 项。
- 第 6 节 讨论常见陷阱、内存和实现模式。

每个中间值都根据 [17](https://arxiv.org/html/2607.13042#bib.bib29) 的手工推导以及 PyTorch 的输出进行了验证；一个配套的 Jupyter notebook 重现了所有计算。物理导数 \(\hat{y}'\) 也可以通过前向模式 AD 计算——在前向传播过程中，一个切向量 \(\dot{a}^{(\ell)}\) 伴随每个激活值一起传递，并由每层的局部雅可比矩阵更新，这正好是 [17](https://arxiv.org/html/2607.13042#bib.bib29) 前向传播表中的对偶传播——PyTorch 通过 `torch.func.jvp` ([https://arxiv.org/html/2607.13042v1/torch.func.jvp](https://arxiv.org/html/2607.13042v1/torch.func.jvp)) 支持这一方式。我们追踪反向模式，因为它是 PINN 训练中 PyTorch 的默认方式。本文不提出新的 PINN 方法、基准测试训练策略或比较 AD 系统；其贡献在于解释说明。

## 2 切向量与伴随向量：链式法则的两面

自动微分是一系列算法，它们通过将计算分解为基本运算并系统地应用链式法则来计算导数。PyTorch 将这些运算记录在一个称为*计算图*的数据结构中。在追踪神经网络中的导数之前，我们在一个两输入的标量函数上介绍前向传播、前向模式 AD、反向模式 AD 以及“图上图”机制——这是 PINN 训练所需一切要素的最小示例。

### 2.1 最小示例

设 \(y = f(g(x_1, x_2))\)，其中：

\[
g(x_1, x_2) = 2x_1 + x_2, \qquad f(u) = u^2.
\tag{3}
\]

当 \(x_1 = 2\)，\(x_2 = 1\) 时：

\[
u = g(2, 1) = 5, \qquad y = f(5) = 25.
\tag{4}
\]

局部导数为：

\[
\frac{\partial g}{\partial x_1} = 2, \qquad \frac{\partial g}{\partial x_2} = 1, \qquad f'(u) = 2u = 10.
\tag{5}
\]

根据链式法则，\(\partial y / \partial x_1 = 10 \cdot 2 = 20\) 且 \(\partial y / \partial x_2 = 10 \cdot 1 = 10\)。我们先计算 \(y = 25\) 并追踪 PyTorch 记录了哪些内容，然后通过两种对应于自动微分两种模式的方式来计算这些导数。

### 2.2 前向传播与图 1

前向传播按顺序求值操作来计算 \(y = 25\)：\(u = 2x_1 + x_2 = 5\)，然后 \(y = u^2 = 25\)。当一个输入设置了 `requires_grad=True` 时，PyTorch 会构建一个记录此序列的计算图。每个操作成为图中的一个*节点*——它是一个包含以下内容的记录：

- 执行的操作和数值结果（原始值），
- 指向输入节点的链接以及在反向传播过程中计算导数所需的雅可比条目（梯度函数 `grad_fn`）。

图 1 显示了在前向传播过程中构建的图 1。

参见图注

**图 1：** 为 \(y = f(g(x_1, x_2))\) 在前向传播过程中构建的图 1。灰色框：叶子节点（输入）。蓝色框：操作节点，每个节点显示节点名称、操作、计算出的原始值以及关联的*梯度函数*。箭头表示依赖关系，从输出指向输入该操作的父亲节点。

图 1 是计算导数的配方：它存储了操作和在导数传播过程中需要用来计算局部导数的数值。没有输入被扰动（如有限差分那样），也没有操作符号公式——该图是基本运算的轨迹，每个运算都配有一个已知的、精确的导数规则。

### 2.3 前向模式 AD：传播一个切向量

前向模式 AD 通过从输入开始，经过计算图传播一个*切向量* \(\dot{\mathbf{x}}\) 来计算导数。通过将输入播种为一个向量 \(\mathbf{v} = (\dot{x}_1, \dot{x}_2)\)（可以视为 \(d\mathbf{x}/d\lambda\) 关于一个独立参数 \(\lambda\)），在每个节点按照以下规则计算导数：

\[
\dot{\text{output}} = J_{\text{local}} \cdot \dot{\text{input}}
\tag{6}
\]

这个操作是一个**雅可比-向量积 (JVP)**，其中 \(m \times n\) 的雅可比矩阵左乘一个 \(n \times 1\) 的切向量。这保持了数据的流动方向，将 \(n\) 个输入扰动映射到 \(m\) 个输出切向量。在从输入到节点 \(u\) 的过程中，\(1 \times 2\) 的雅可比矩阵 \(J_u = [\partial u/\partial x_1, \partial u/\partial x_2]\) 乘以 \(2 \times 1\) 的切向量 \((\dot{x}_1, \dot{x}_2)^\top\)，产生标量切向量 \(\dot{u}\)。

最终结果 \(\dot{y}\) 表示 \(y\) 在种子向量 \(\mathbf{v}\) 方向上的方向导数，并按该方向的模进行缩放。具体来说，选择标准基向量 \(e_i\) 作为种子可以恢复偏导数 \(\partial y / \partial x_i\)。

表 1 展示了获得两个偏导数所需的两遍传播过程。

**表 1：** 前向模式 AD 每个输入需要一遍传播。每遍传播通过图 1 使用相同的 JVP 序列传播一个切向种子。

图 2 显示了通过图 1 的两遍传播过程。切向值（绿色，粗体）显示在每个节点内部。第一遍播种 \(\dot{x}_1 = 1\)，\(\dot{x}_2 = 0\)