系列导读
这是 Physics-Informed Vibe Coding 系列的第四期。经过前三期对 NTK 自适应权重、多尺度 Fourier 特征嵌入以及坐标尺度变换的系统探索，人与 AI 智能体之间的协作已经形成了一套可复现的工作范式。在这一范式中，研究者承担算法甄选与实验方案的顶层设计，AI 智能体则作为高效的执行引擎，将设计意图转化为可运行的代码与可验证的数值结果——整个过程中没有一行代码出自人类之手。我们将继续沿用 Vibe Coding & Vibe Researching 的方法论，在每一期中挑战一个具有代表性的 PINN 改进算法，检验这种"零手写代码"的科研协作模式究竟能走多远。
所有源代码与教程文档均开源于 GitHub: Physics-informed-vibe-coding，全部数值实现采用 JAX 框架编写。
本期主题：梯度加权 PINN —— 基于梯度流统计的自适应损失加权策略，以非线性 Klein-Gordon 方程为算例，从梯度病态诊断的视角出发，详解如何通过梯度统计量动态平衡多任务损失分量，将求解精度提升近两个数量级。
物理信息神经网络（PINN）的训练本质上是一个多任务优化问题：网络需要同时最小化 PDE 残差损失、初始条件损失和边界条件损失。当各损失分量的梯度信号尺度严重失衡时，优化器将被梯度主导项"绑架"，导致其余约束条件得不到充分满足——这一现象被称为梯度流病态（gradient flow pathology）。
Wang 等人 [1] 提出了一种基于梯度统计量的学习率退火算法（learning rate annealing），其核心思想是：在每一步训练中，计算 PDE 残差损失相对于网络参数的梯度最大值，与初始/边界条件损失梯度均值之比，作为自适应权重
本教程以非线性 Klein-Gordon 方程为算例，基于 JAX 框架完整复现了梯度加权 PINN 算法，设置两组对比实验：
| 模型 | 权重策略 | 最优 |
|---|---|---|
| M1（标准 PINN） |
|
|
| M2（梯度加权 PINN） | 基于梯度统计的自适应 |
M2 相较于 M1 实现了 95.5% 的误差降低，充分验证了梯度加权策略在缓解梯度流病态方面的有效性。
PINN [2] 通过将 PDE 的物理约束编码到神经网络的损失函数中来求解偏微分方程。对于一般的 PDE 初边值问题，PINN 的损失函数可以分解为：
其中 $\mathcal{L}r$ 为 PDE 残差损失，$\mathcal{L}{ic}$ 和 $\mathcal{L}{bc}$ 分别为初始条件和边界条件的损失，$\lambda{ic}$、$\lambda_{bc}$ 为对应的权重系数。
标准 PINN 采用
然而，上述朴素设定在实际训练中往往导致严重的收敛困难。其根本原因在于，不同损失分量对网络参数的梯度信号在量级和方向上可能存在巨大差异：
-
PDE 残差
$\mathcal{L}_r$ 涉及网络输出的高阶微分（如$u_{tt}$ 、$u_{xx}$），其梯度信号经过多次求导后呈指数放大或衰减，主导了总梯度的方向； - 初始/边界条件 $\mathcal{L}{ic}$、$\mathcal{L}{bc}$ 仅约束网络的零阶或一阶输出，梯度信号相对较弱。
这种不平衡会导致优化器几乎完全沿着 $\nabla_\theta \mathcal{L}r$ 的方向更新参数，而初始条件和边界条件的梯度贡献被"淹没"，最终表现为 $\mathcal{L}{ic}$ 和
Wang 等人 [1] 观察到，如果能够使各损失分量的反向传播梯度在量级上保持大致平衡，则网络可以在每一步更新中同时改善所有约束条件。基于此，他们提出了一种简单而有效的策略：利用各损失分量梯度的统计量（最大值与均值之比）来动态调整权重
与 NTK 自适应权重方法 [3] 相比，梯度加权策略无需计算庞大的核矩阵（$O(N^2 P)$ 复杂度），仅需一次额外的反向传播（$O(P)$ 复杂度），在实现上更为轻量。
本教程采用的基准算例为非线性 Klein-Gordon 方程：
取参数
精确解为：
源项
| 条件类型 | 数学表达式 | 定义域 |
|---|---|---|
| 初始条件（Dirichlet） | ||
| 初始条件（Neumann） | ||
| 边界条件（$x=0$） | ||
| 边界条件（$x=1$） |
该问题的难点在于：解包含
对于上述 Klein-Gordon 问题，PINN 的损失函数具体分解为：
PDE 残差损失：
初始条件损失：
边界条件损失：
总损失为：
为理解梯度流病态，考察各损失分量对某一网络参数层
而对于边界条件损失，仅涉及网络的零阶输出：
在网络初始化阶段，高阶微分项可能产生极大的梯度值（尤其当激活函数的导数在某些区间较大时），而零阶项的梯度相对温和。这意味着 $|\nabla_\theta \mathcal{L}r| \gg |\nabla\theta \mathcal{L}{bc}|$，在 $\lambda{bc} = 1$ 的设定下，边界条件的梯度贡献在总梯度中几乎可以忽略不计。
Wang 等人 [1] 提出的梯度加权策略可以形式化描述如下。
步骤一：计算各损失分量的梯度统计量。 对于网络的每一层
步骤二：跨层聚合。 对所有层取全局统计量：
步骤三：计算瞬时自适应权重。
其中
直觉上，当
直接使用瞬时估计
EMA 系数
为改善网络的训练稳定性，对输入坐标施加 z-score 归一化：
其中
由于网络实际接收的是归一化坐标
这一修正确保了在归一化空间中计算的导数能正确映射回物理空间的导数值。
本实验采用的全连接网络结构为
下表展示了一个 batch size 为
| 阶段 | 操作 | 输入维度 | 输出维度 | 数学表达式 |
|---|---|---|---|---|
| 输入归一化 | z-score | |||
| 输入层 → 隐藏层 1 | 线性 + tanh | |||
| 隐藏层 1 → 隐藏层 2 | 线性 + tanh | |||
| 隐藏层 2 → 隐藏层 3 | 线性 + tanh | |||
| 隐藏层 3 → 隐藏层 4 | 线性 + tanh | |||
| 隐藏层 4 → 隐藏层 5 | 线性 + tanh | |||
| 隐藏层 5 → 输出层 | 线性（无激活） |
权重矩阵
偏置
| 层 | 权重矩阵尺寸 | 权重参数量 | 偏置参数量 | 小计 |
|---|---|---|---|---|
| 第 1 层 | 100 | 50 | 150 | |
| 第 2 层 | 2500 | 50 | 2550 | |
| 第 3 层 | 2500 | 50 | 2550 | |
| 第 4 层 | 2500 | 50 | 2550 | |
| 第 5 层 | 2500 | 50 | 2550 | |
| 第 6 层 | 50 | 1 | 51 | |
| 总计 | — | 10150 | 251 | 10401 |
PDE 残差的计算是 PINN 中最核心的部分。以下代码展示了如何利用 JAX 的自动微分能力计算二阶导数，并通过链式法则修正归一化带来的尺度因子：
def residual_single(params, t_n, x_n, f_val, sigma_t, sigma_x):
u = net_u_scalar(params, t_n, x_n)
u_t = grad(net_u_scalar, 1)(params, t_n, x_n) / sigma_t
u_tt = grad(lambda p, tn, xn:
grad(net_u_scalar, 1)(p, tn, xn) / sigma_t,
argnums=1)(params, t_n, x_n) / sigma_t
u_xx = grad(lambda p, tn, xn:
grad(net_u_scalar, 2)(p, tn, xn) / sigma_x,
argnums=2)(params, t_n, x_n) / sigma_x
pde = u_tt + ALPHA * u_xx + BETA_PDE * u + GAMMA * u ** K_EXP
return pde - f_val数学原理解读：
-
net_u_scalar(params, t_n, x_n)计算网络在归一化坐标$(\tilde{t}, \tilde{x})$ 处的标量输出$u$ ； -
grad(net_u_scalar, 1)对第 1 个参数（t_n）求偏导，得到$\partial u / \partial \tilde{t}$ ，除以$\sigma_t$ 后得到物理导数$\partial u / \partial t$ ； - 二阶导数
$u_{tt}$ 通过嵌套grad实现，外层再除以$\sigma_t$ ，总效果为$u_{tt} = (1/\sigma_t^2) \cdot \partial^2 u / \partial \tilde{t}^2$ ； -
vmap将标量函数向量化，使其能够并行处理一个 batch 的数据。
梯度加权策略的核心在于 compute_adaptive_lambdas 函数：
def compute_adaptive_lambdas(params, t_r, x_r, f_r, t_ic, x_ic, u_ic,
t_bc, x_bc, u_bc, sigma_t, sigma_x,
lam_ics_cur, lam_bcs_cur):
grads_res = grad(loss_res_fn)(params, t_r, x_r, f_r, sigma_t, sigma_x)
def weighted_ics(params, t_ic, x_ic, u_ic, sigma_t):
return lam_ics_cur * loss_ics_fn(params, t_ic, x_ic, u_ic, sigma_t)
grads_ics = grad(weighted_ics)(params, t_ic, x_ic, u_ic, sigma_t)
def weighted_bcs(params, t_bc, x_bc, u_bc):
return lam_bcs_cur * loss_bcs_fn(params, t_bc, x_bc, u_bc)
grads_bcs = grad(weighted_bcs)(params, t_bc, x_bc, u_bc)
max_res_list, mean_ics_list, mean_bcs_list = [], [], []
for g_r, g_i, g_b in zip(grads_res, grads_ics, grads_bcs):
max_res_list.append(jnp.max(jnp.abs(g_r['w'])))
mean_ics_list.append(jnp.mean(jnp.abs(g_i['w'])))
mean_bcs_list.append(jnp.mean(jnp.abs(g_b['w'])))
max_grad_res = jnp.max(jnp.array(max_res_list))
mean_grad_ics = jnp.mean(jnp.array(mean_ics_list))
mean_grad_bcs = jnp.mean(jnp.array(mean_bcs_list))
lam_ics_hat = max_grad_res / (mean_grad_ics + 1e-10)
lam_bcs_hat = max_grad_res / (mean_grad_bcs + 1e-10)
return lam_ics_hat, lam_bcs_hat逐步解析：
- 三次独立反向传播：分别对 $\mathcal{L}r$、$\lambda{ic}\mathcal{L}{ic}$、$\lambda{bc}\mathcal{L}{bc}$ 求关于参数 $\theta$ 的梯度。注意 IC/BC 损失在求梯度时已经乘上了当前权重 $\lambda{ic}^{(k)}$、$\lambda_{bc}^{(k)}$，这与论文的定义保持一致；
-
逐层统计：对每一层的权重梯度矩阵，残差损失取
$\max|\cdot|$ （捕捉最大梯度信号），IC/BC 损失取$\text{mean}|\cdot|$ （反映平均梯度水平）； - 跨层聚合：残差的逐层最大值取全局最大，IC/BC 的逐层均值取全局均值；
- 权重计算：$\hat{\lambda} = G_r^{\max} / G_{ic/bc}^{\text{mean}}$。
两种模式在训练步骤上的差异集中在权重更新机制：
# M1: 固定权重
@jit
def train_step_m1(params, opt_state, t_r, x_r, f_r,
t_ic, x_ic, u_ic, t_bc, x_bc, u_bc):
lam_i = jnp.float32(1.0)
lam_b = jnp.float32(1.0)
(loss, (l_res, l_ics, l_bcs)), grads = jax.value_and_grad(
loss_total_fn, has_aux=True)(
params, t_r, x_r, f_r, t_ic, x_ic, u_ic,
t_bc, x_bc, u_bc, sigma_t, sigma_x, lam_i, lam_b)
updates, new_opt = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)
return new_params, new_opt, loss, l_res, l_ics, l_bcs
# M2: 自适应权重
@jit
def train_step_m2(params, opt_state, t_r, x_r, f_r,
t_ic, x_ic, u_ic, t_bc, x_bc, u_bc,
lam_ics, lam_bcs):
(loss, (l_res, l_ics, l_bcs)), grads = jax.value_and_grad(
loss_total_fn, has_aux=True)(
params, t_r, x_r, f_r, t_ic, x_ic, u_ic,
t_bc, x_bc, u_bc, sigma_t, sigma_x, lam_ics, lam_bcs)
updates, new_opt = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)
return new_params, new_opt, loss, l_res, l_ics, l_bcs| 特性 | M1（标准 PINN） | M2（梯度加权 PINN） |
|---|---|---|
| 权重 |
固定为 1 | 每 10 步自适应更新 |
| 额外梯度计算 | 无 | 3 次额外反向传播（每 10 步） |
| EMA 平滑 | 不适用 | |
| 计算开销 | 基准 | 约增加 5-10%（因仅每 10 步更新） |
训练流程的关键步骤如下：
-
随机采样：每步从计算域内随机采样
$N = 128$ 个配点（domain points）、IC 点和 BC 点； -
输入归一化：将物理坐标
$(t, x)$ 转换为归一化坐标$(\tilde{t}, \tilde{x})$ ； - 前向传播与损失计算：通过网络和自动微分计算 PDE 残差、IC 损失和 BC 损失；
- 权重更新（仅 M2）：每 10 步计算一次梯度统计量，更新自适应权重；
-
参数更新：使用 Adam 优化器 [5] 结合指数学习率衰减（初始
$\eta_0 = 10^{-3}$ ，每 1000 步衰减$\times 0.9$ ）更新网络参数。
| 配置项 | 设定值 |
|---|---|
| PDE 方程 | Klein-Gordon：$u_{tt} - u_{xx} + u^3 = f$ |
| 计算域 | |
| 精确解 | |
| 网络结构 |
|
| 参数总量 | 10401 |
| 初始化 | Xavier（Glorot） |
| 优化器 | Adam（$\eta_0 = 10^{-3}$） |
| 学习率调度 | 指数衰减（rate=0.9，每 1000 步） |
| Batch size | 128（每类配点） |
| 总迭代次数 | 40001 |
| EMA 系数 |
0.9 |
|
|
每 10 步 |
| 随机种子 | 1234 |
下图展示了 M1 与 M2 在整个计算域
图 1 上排为 M1（标准 PINN），下排为 M2（梯度加权 PINN）。
分析：
- M1 的预测解在
$t > 0.3$ 区域出现严重偏差，逐点误差分布显示高值区域几乎覆盖了整个计算域的大部分区域； - M2 的预测解在视觉上与精确解高度吻合，逐点误差集中在极少数局部区域且量级远低于 M1。
以下是 M1 和 M2 各自的解场详细视图：
图 2 M1（标准 PINN）解场：精确解、预测解、绝对误差。
图 3 M2（梯度加权 PINN）解场：精确解、预测解、绝对误差。
图 4 M1 与 M2 的三组损失分量（$\mathcal{L}r$、$\mathcal{L}{ic}$、$\mathcal{L}_{bc}$）收敛曲线对比。
分析：
| 损失分量 | M1 终值 | M2 终值 | 降低比例 |
|---|---|---|---|
|
|
80.9% | ||
|
|
94.0% | ||
|
|
99.9% |
从上表可以看出：
- M2 的边界条件损失
$\mathcal{L}_{bc}$ 相比 M1 降低了近 3 个数量级（从$10^{-2}$ 到$10^{-5}$ ），说明自适应权重极大增强了边界条件的约束力度； - M2 的初始条件损失
$\mathcal{L}_{ic}$ 也降低了约 1 个数量级； - M1 的残差损失在后期未能有效收敛，而 M2 的残差损失降至更低水平。
图 5 M1 与 M2 的
分析：
- M1 的
$L^2$ 误差在约 5000 步后即停滞在$\sim 0.17$ 附近，此后几乎没有改善。这表明标准 PINN 在该问题上陷入了梯度流病态导致的局部困境； - M2 的
$L^2$ 误差在前 5000 步内即快速下降至$\sim 0.05$ ，此后持续缓慢收敛，最终达到$7.75 \times 10^{-3}$ 。
下表展示了不同迭代阶段的
| 迭代次数 | M1 |
M2 |
M2 优势倍数 |
|---|---|---|---|
| 0 | |||
| 1000 | |||
| 5000 | |||
| 10000 | |||
| 20000 | |||
| 30000 | |||
| 40000 |
图 6 M2 中 $\hat{\lambda}{ic}$ 和 $\hat{\lambda}{bc}$ 随迭代次数的演化。
分析：
- $\hat{\lambda}{bc}$ 的值远大于 $\hat{\lambda}{ic}$，这表明边界条件的梯度信号远弱于残差梯度，需要更大的权重来补偿；
-
$\hat{\lambda}_{bc}$ 在训练过程中整体呈现先升后稳的趋势，最终稳定在$200 \sim 300$ 的量级——这意味着边界条件损失被放大了约 200-300 倍，才能与残差损失的梯度贡献相匹配； -
$\hat{\lambda}_{ic}$ 相对稳定，在$30 \sim 80$ 之间波动。
| 指标 | M1（标准 PINN） | M2（梯度加权 PINN） | 改善比例 |
|---|---|---|---|
| 参数量 | 10401 | 10401 | — |
| 迭代次数 | 40001 | 40001 | — |
| 最优 |
95.5% | ||
| 终态 |
93.9% | ||
| 终态 |
80.9% | ||
| 终态 |
94.0% | ||
| 终态 |
99.9% |
下表从损失结构的角度对比了两种方法的训练行为，有助于理解梯度加权策略的作用机理：
| 损失分量 | M1 行为 | M2 行为 |
|---|---|---|
|
|
在前 5000 步快速下降后趋于震荡，终态值仍偏高 | 更加平稳地持续下降，终态值降至 M1 的约 1/5 |
|
|
始终在 |
通过自适应权重增强约束，快速降至 |
|
|
停滞在 |
自适应权重 |
|
|
在约 |
前期快速收敛，后期仍有缓慢改善 |
本教程基于 JAX 框架完整复现了 Wang 等人 [1] 提出的梯度加权 PINN 算法，并以非线性 Klein-Gordon 方程为算例进行了系统的对比实验。主要结论如下：
-
梯度流病态是标准 PINN 训练失败的关键原因。 M1（$\lambda = 1$）的
$L^2$ 误差停滞在$1.73 \times 10^{-1}$ ，边界条件损失始终无法有效收敛； -
梯度加权策略以极低的计算开销实现了显著的精度提升。 M2 仅需每 10 步进行 3 次额外的反向传播，即可将
$L^2$ 误差降低至$7.75 \times 10^{-3}$ ，改善幅度达 95.5%； -
自适应权重的幅值揭示了梯度失衡的严重程度。
$\hat{\lambda}_{bc}$ 稳定在$200 \sim 300$ ，表明在标准设定下边界条件的梯度贡献被低估了两个数量级以上。
-
依赖梯度统计量的代表性：$\max|\cdot|$ 和
$\text{mean}|\cdot|$ 是粗粒度的统计量，可能无法完全捕捉梯度分布的复杂结构； -
EMA 系数
$\beta$ 的敏感性：$\beta$ 过小导致权重振荡，过大则响应迟钝，最优值可能因问题而异； - 权重更新频率的选择：本实验中每 10 步更新一次，该频率对不同问题可能需要调整。
- 将梯度加权策略与多尺度 Fourier 特征嵌入相结合，以同时解决梯度失衡和谱偏差问题；
- 探索更鲁棒的梯度统计量（如分位数或对数均值）替代 max/mean；
- 将该方法推广到更高维度和更复杂的 PDE 系统（如 Navier-Stokes 方程）。
Q1：为什么 M1 的
这是典型的梯度流病态现象。当
Q2：EMA 系数
-
$\beta \in [0.8, 0.95]$ ：适用于大多数问题； -
$\beta < 0.8$ ：权重响应更快但可能导致振荡； -
$\beta > 0.95$ ：权重更平滑但适应速度变慢。
Q3：梯度加权与 NTK 自适应权重有何区别？
| 特性 | 梯度加权（本方法） | NTK 自适应权重 [3] |
|---|---|---|
| 数学基础 | 梯度统计量 | 神经正切核特征值 |
| 计算复杂度 |
|
|
| 实现难度 | 低 | 中等 |
| 物理可解释性 | 直接反映梯度失衡程度 | 反映各约束的有效学习速率 |
| 更新频率 | 可以每步或间隔更新 | 通常间隔较大步数更新 |
两种方法解决的核心问题相同（多任务损失的梯度失衡），但切入角度不同。梯度加权方法更轻量，适合作为即插即用的改进策略。
Q4：为什么对 IC/BC 损失的梯度取
论文 [1] 的设计意图是：PDE 残差的梯度信号分布通常更分散（因为残差表达式更复杂），取
Q5：输入归一化是否必要？
从本项目的消融实验（case3 和 case4）可以看出，单独的归一化对精度改善有限（仅约 1.3%），但归一化与梯度加权结合后的效果远超两者之和。这表明归一化主要起到改善数值条件的辅助作用，而非核心改进手段。
[1] S. Wang, Y. Teng, and P. Perdikaris, "Understanding and mitigating gradient flow pathologies in physics-informed neural networks," SIAM Journal on Scientific Computing, vol. 43, no. 5, pp. A3055–A3081, 2021.
[2] M. Raissi, P. Perdikaris, and G. E. Karniadakis, "Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations," Journal of Computational Physics, vol. 378, pp. 686–707, 2019.
[3] S. Wang, X. Yu, and P. Perdikaris, "When and why PINNs fail to train: A neural tangent kernel perspective," Journal of Computational Physics, vol. 449, p. 110768, 2022.
[4] X. Glorot and Y. Bengio, "Understanding the difficulty of training deep feedforward neural networks," in Proceedings of the 13th International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 249–256, 2010.
[5] D. P. Kingma and J. Ba, "Adam: A method for stochastic optimization," in Proceedings of the 3rd International Conference on Learning Representations (ICLR), 2015.