PINN训练新思路：把初始条件和边界约束嵌入网络架构，解决多目标优化难题

PINNs出了名的难训练。主要原因之一就是这个多目标优化问题。优化器很容易找到投机取巧的路径——比如拼命降低微分方程残差，但完全不管初始条件和边界约束。只要给初始条件和边界损失配的权重够低，它们增加的那点损失完全能被残差损失的大幅下降抵消掉。调整权重也许能暂时缓解这个问题，但谁也不能保证最优权重在整个训练过程中一直有效。

标准的PINN用复合损失函数，把三项加权求和：

初始条件损失

边界损失

微分方程残差损失

要让解有用，必须让所有损失项同时降下来。用复合损失训练PINN的时候，优化器面对多个目标，有什么办法让优化器的工作简单点呢？

硬约束

要简化训练，最好能把复合损失换成单一项。所以如果能设计一种神经网络架构，让它自动满足初始和边界条件那么事情就会简单很多了。初始条件和边界条件直接"钉"在网络结构里，所以叫"硬约束"。相比之下，原始PINN把初始和边界条件放在损失函数里，那些约束是"软"的——不保证精确满足。

下面看看怎么把这些约束具体嵌入网络架构。

强制初始条件

初始条件一般是空间域上的一系列离散测量点。比如1D + t（一个空间维度加时间）的情况，初始条件就是沿x轴的一串值：(0 m, 100°C), (0.01 m, 105°C), (0.02 m, 106°C), … (0.3 m, 31°C)，诸如此类。

假设能用函数u₀(x)在整个空间域上插值。对1D + t的例子，如果x在样本点上，u₀(x)返回观测值；否则返回平滑插值，x∈ [0 m, 0.3 m]。三次样条很适合干这活。现在把初始条件的插值函数u₀(x)整合进PINN架构的输出：

方程(1)给出了一种强制u(x, 0) = u₀(x)的方式。β是常数或者可训练的时间衰减参数，决定了从初始状态到热扩散状态的过渡期长短。任意函数v(x, t)是神经网络的输出，不过还有点额外的处理。

强制边界条件

如果方程(1)里的函数v(x, t)在边界上取零值，再加上u₀(x)满足Dirichlet边界条件的假设，那么u(xᵇ, t) = u₀(xᵇ)，这里xᵇ是边界上的点。

Ω(x)是个平滑函数，边界处输出零，空间域内部输出非零有限值。可以想象成平面上的肥皂泡。

文献[1]把Ω(x)叫做approximate distance function（近似距离函数），因为它的行为类似于输出到最近边界距离的函数。我觉得叫boundary bubble function（边界泡泡函数）更形象点。

对1D + t情况，Ω(x)可以取抛物线形式，在空间边界x_L和x_R处为零：

z(x, t)是任意函数的输出。这里用神经网络。

输出u(x, t)不是直接从可训练网络z(x, t)出来的，还能用梯度下降优化网络参数吗？

能。方程(3)里涉及的其他函数（u₀(x)、时间衰减函数、边界泡泡函数Ω(x)）对x和t都可导。深度学习框架反向传播损失张量的时候会把它们的梯度算进去。所以初始条件插值、时间衰减、边界泡泡函数的具体形式其实不太关键，神经网络会自己适应。

用方程(3)作为PINN输出的过滤器，就可以开始求解偏微分方程(PDE)了。

实验

考虑一个薄金属棒，初始温度分布已知，来自一系列测量值。

初始条件，一系列温度测量点定义。图片由作者提供。

描述温度随时间演化的PDE是1D热扩散方程：

这个问题的解析解是知道的，所以能拿PINN的结果跟解析解对比。

PINN架构核心是个三层ResNet，每层宽度32。

初始条件用三次样条插值，边界泡泡函数用抛物线：

class HardConstrained1dResNet(torch.nn.Module):

  def __init__(self, number_of_blocks, block_width, number_of_outputs,

               initial_profile, time_bubble_beta=2.0):

      super().__init__()

      self.bubble = Bubble()

      self.z_predictor = pinn1d.ResidualNet(

          number_of_inputs=2,

          number_of_blocks=number_of_blocks,

          block_width=block_width,

          number_of_outputs=number_of_outputs

      )

      self.initial_profile = initial_profile

      delta_x = 1.0/(len(self.initial_profile) - 1)

      xs = np.arange(0, 1 + delta_x/2, delta_x)

      xy_list = []

      for x_ndx in range(len(xs)):

          x = xs[x_ndx]

          xy_list.append((x, self.initial_profile[x_ndx]))

      self.initial_profile_interpolator = interpolation1d.CubicSpline(

          xy_list, boundary_condition='2nd_derivative_0'

      )

      self.time_bubble_beta = torch.nn.Parameter(torch.tensor([time_bubble_beta]))

  def forward(self, x_t):  # x_t.shape = (B, 2)

      x_tsr = x_t[:, 0].unsqueeze(1)  # (B, 1)

      t_tsr = x_t[:, 1].unsqueeze(1)  # (B, 1)

      bubble_t_tsr = self.time_bubble(t_tsr)  # (B, 1)

      bubble_x_tsr = self.bubble(x_tsr)  # (B, 1)

      z_tsr = self.z_predictor(x_t)  # (B, 1)

      initial_interpolation_tsr = self.initial_profile_interpolator.batch_evaluate(x_tsr)

      return initial_interpolation_tsr + bubble_t_tsr * bubble_x_tsr * z_tsr

  def time_bubble(self, t_tsr):  # t_tsr.shape = (B, 1)

      return 1 - torch.exp(-self.time_bubble_beta * t_tsr)  # (B, 1)

训练程序train.py里，损失函数只有一项——微分方程残差损失。初始条件和边界损失架构设计就保证了，不用放到损失函数里：

(...)

# Differential equation residual loss

diff_eqn_residual_x_t_tsr.detach_()

diff_eqn_residual_x_t_tsr.requires_grad = True

du_dx__du_dt = first_derivative(neural_net, diff_eqn_residual_x_t_tsr)

du_dt = du_dx__du_dt[:, 1]  # (N_res)

d2u_dx2__d2u_dxdt = second_derivative(neural_net, diff_eqn_residual_x_t_tsr, 0)  # (N_res, 2)

d2u_dx2 = d2u_dx2__d2u_dxdt[:, 0]  # (N_res)

diff_eqn_residual = 1.0/duration * du_dt - alpha/length**2 * d2u_dx2  # (N_res)

diff_eqn_residual_loss = criterion(diff_eqn_residual, torch.zeros_like(diff_eqn_residual))

loss = diff_eqn_residual_loss

is_champion = False

(...)

下图是典型训练过程的损失变化：

动画1对比了10秒模拟期间的近似解和解析解：

PINN预测与解析解的比较。图片由作者提供。

符合预期，t=0时的预测就是初始条件的插值，x=0 m和x=0.3 m处的温度值精确等于设定的边界温度（Dirichlet边界条件）。

总结

多目标优化给PINN训练带来的困难是实实在在的。设计上做点简单修改，让网络输出在t=0时自动匹配初始条件插值，在边界上自动满足边界条件，就能把问题简化不少。

训练PINN解1D+t热扩散问题的结果还不错，从可视化能清楚看到PINN学会了满足PDE，同时被强制满足初始和边界条件。

这套方法对那些物理约束不能妥协的领域可能挺有价值——气候建模、生物医学仿真之类的场景。

本文的代码在这里：

https://github.com/sebastiengilbert73/tutorial_pinn_hard_constraints

参考：

[1] Exact imposition of boundary conditions with distance functions in physics-informed deep neural networks, N. Sukumar, Ankit Srivastava, https://arxiv.org/abs/2104.08426

作者：Sébastien Gilbert

点个在看你最好看！