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://avoid.overfit.cn/post/4b21ca89cc714512bff16ffd1af69538
参考:
[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