Respecting causality is all you need for training physics-informed neural networks
论文阅读：Respecting causality is all you need for training physics-informed neural networks
- Respecting causality is all you need for training physics-informed neural networks
-
- 问题分析
-
- 以Allen-Cahn方程为例
- 损失函数的重新表述
- 解决方法
-
- 想法
- 以Allen-Cahn方程为例
- 停止条件
- 完整算法
- 实验结果
-
- Lorentz system
- Kuramoto–Sivashinsky equation
- Navier-Stokes equation
- 总结
Respecting causality is all you need for training physics-informed neural networks
问题分析
以Allen-Cahn方程为例
这个例子很难用 Raissi 等人原始的连续时间公式直接求解。模型似乎卡在了某个中间状态，无法进一步细化以提供对方程的准确近似。
损失函数的重新表述
于是，作者就利用前向欧拉对损失函数进行了重新表述，原始的残差函数为:
L r ( θ ) = 1 N r ∑ i = 1 N r ∣ ∂ u θ ∂ t ( t r i , x r i ) + N [ u θ ] ( t r i , x r i ) ∣ 2 \mathcal{L} _r(\boldsymbol{\theta })=\frac{1}{N_r}\sum_{i=1}^{N_r}{\left| \frac{\partial \boldsymbol{u}_{\boldsymbol{\theta }}}{\partial t}(t_{r}^{i},\boldsymbol{x}_{r}^{i})+\mathcal{N} [\boldsymbol{u}_{\boldsymbol{\theta }}](t_{r}^{i},\boldsymbol{x}_{r}^{i}) \right|^2} Lr(θ)=Nr1i=1∑Nr ∂t∂uθ(tri,xri)+N[uθ](tri,xri) 2
可以按时间片差分为如下形式：
L r ( θ ) = 1 N t ∑ i = 1 N t L r ( t i , θ ) = 1 N t N x ∑ i = 1 N t ∑ j = 1 N x ∣ ∂ u θ ∂ t ( t i , x j ) + N [ u θ ] ( t i , x j ) ∣ 2 \begin{aligned} \mathcal{L} _r(\boldsymbol{\theta })&=\frac{1}{N_t}\sum_{i=1}^{N_t}{\mathcal{L} _r}(t_i,\theta )\\ &=\frac{1}{N_tN_x}\sum_{i=1}^{N_t}{\sum_{j=1}^{N_x}{|}}\frac{\partial \boldsymbol{u}_{\boldsymbol{\theta }}}{\partial t}(t_i,\boldsymbol{x}_{\boldsymbol{j}})+\mathcal{N} [\boldsymbol{u}_{\boldsymbol{\theta }}](t_i,\boldsymbol{x}_{\boldsymbol{j}})|^2\\ \end{aligned} Lr(θ)=Nt1i=1∑NtLr(ti,θ)=NtNx1i=1∑Ntj=1∑Nx∣∂t∂uθ(ti,xj)+N[uθ](ti,xj)∣2
进一步，使用前向欧拉法，将其进行分解，得到：
L r ( t i , θ ) ≈ 1 N x ∑ j = 1 N x ∣ u θ ( t i , x j ) − u θ ( t i − 1 , x j ) Δ t + N [ u θ ] ( t i , x j ) ∣ 2 ≈ ∣ Ω ∣ Δ t 2 ∫ Ω ∣ u θ ( t i , x ) − u θ ( t i − 1 , x ) + Δ t N [ u θ ] ( t i , x ) ∣ 2 d x \begin{aligned} \mathcal{L} _r(t_i,\boldsymbol{\theta })&\approx \frac{1}{N_x}\sum_{j=1}^{N_x}{\left| \frac{u_{\theta}(t_i,x_j)-u_{\theta}(t_{i-1},x_j)}{\Delta t}+\mathcal{N} [u_{\theta}](t_i,x_j) \right|^2}\\ &\approx \frac{|\Omega |}{\Delta t^2}\int_{\Omega}{|}\boldsymbol{u}_{\boldsymbol{\theta }}(t_i,x)-\boldsymbol{u}_{\boldsymbol{\theta }}(t_{i-1},x)+\Delta t\mathcal{N} [\boldsymbol{u}_{\boldsymbol{\theta }}](t_i,x)|^2dx\\ \end{aligned} Lr(ti,θ)≈Nx1j=1∑Nx Δtuθ(ti,xj)−uθ(ti−1,xj)+N[uθ](ti,xj) 2≈Δt2∣Ω∣∫Ω∣uθ(ti,x)−uθ(ti−1,x)+ΔtN[uθ](ti,x)∣2dx
从上式中可以看出， t i t_i ti时刻残差的最小化应该基于 t i t_i ti 和 t i − 1 t_{i-1} ti−1 时刻 u θ u_\theta uθ 的正确预测，而原始公式倾向于同时最小化所有时刻的残差。因此，原始公式即使 t i t_i ti 和之前时间的预测不准确，残差损失也会最小化。这种行为不可避免地违反了时间因果关系，使 PINN 模型容易学习到错误的解决方案。
上图左侧是PINN训练 2 × 1 0 5 2 \times 10^5 2×105 轮数中误差的变化情况；中间是训练完后，PINN在不同时间片上的误差；右侧是PINN在不同迭代次数时在不同时间片上的收敛速度（根据神经网络的神经正切核得到，感兴趣可以去看原文）。
从中间的图中可以看到，训练后的PINN在 t 0 t_0 t0 时刻仍有较大误差，但在 t 0.5 t_{0.5} t0.5 时刻误差反而较小，结合上述公式，这说明PINN可能在 t 0.5 t_{0.5} t0.5 时刻学习到了错误的解决方案。同时从右侧的图中可以看到，PINN在不同训练轮数时，普遍倾向于首先学习更后面的时刻。
解决方法
想法
为了让PINN能够按时间顺序来学习PDE，自然地作者就想到如下方法：
**对损失函数加权：**为了平衡PINN固有的偏见，可以对收敛速度不同的时间点的残差损失加上不同的权重，即：
L r ( θ ) = 1 N t ∑ i = 1 N t w i L r ( t i , θ ) . \mathcal{L} _r(\boldsymbol{\theta })=\frac{1}{N_t}\sum_{i=1}^{N_t}{w_i}\mathcal{L} _r(t_i,\boldsymbol{\theta }). Lr(θ)=Nt1i=1∑NtwiLr(ti,θ).
其中，
w i = exp ( − ϵ ∑ k = 1 i − 1 L r ( t k , θ ) ) , f o r i = 2 , 3 , . . . , N t w_i=\exp \left( -\epsilon \sum_{k=1}^{i-1}{\mathcal{L} _r}(t_k,\theta ) \right) ,\mathrm{for}i=2,3,...,N_t wi=exp(−ϵk=1∑i−1Lr(tk,θ)),fori=2,3,...,Nt
其中 ϵ \epsilon ϵ 被称为控制权重，是用来控制权重的陡峭程度的超参数。
由于权重与先前时间步的累积残差损失的大小成反比。因此，只有 t i t_i ti 时刻前的所有残差减小到某个小值，使得 t i t_i ti 时刻的权重足够大时，网络才会最小化 t i t_i ti 时刻的残差损失。
以Allen-Cahn方程为例
采用上述修改，取 $ \epsilon = 100$ 来最小化加权后的损失来训练相同的网络，在超参数完全相同的设置下进行 300000 次梯度下降迭代。
可以看到，改进后的PINN确实学习到了正确的解。
停止条件
实验结果也表明，通过观察残差权重 w i w_i wi 的大小可以为评估训练期间 PINN 模型的收敛性提供有效的停止标准。具体来说，对于某个选定的阈值参数 δ ∈ ( 0 , 1 ) \delta \in (0, 1) δ∈(0,1) ，可以选择当 m i n ( w i ) > δ min(w_i) > \delta min(wi)>δ 时停止训练。
从上图中可以看出，改进后PINN的loss相比原始PINN下降了3个数量级，并且不同时间片上的loss也是按时间顺序下降的。
完整算法
实验结果
作者又在Lorentz system、Kuramoto–Sivashinsky equation以及Navier-Stokes equation上进行了试验。以下是一些结果展示。
Lorentz system
d x d t = σ ( y − x ) , d y d t = x ( ρ − z ) − y , d z d t = x y − β z . \begin{aligned} &\frac{\mathrm{d}x}{\mathrm{d}t} =\sigma(y-x), \\ &\frac{\mathrm{d}y}{\mathrm{d}t} =x(\rho-z)-y, \\ &\frac{\mathrm{d}z}{\mathrm{d}t} =xy-\beta z. \end{aligned} dtdx=σ(y−x),dtdy=x(ρ−z)−y,dtdz=xy−βz.
Kuramoto–Sivashinsky equation
u t + α u u x + β u x x + γ u x x x = 0 u ( 0 , x ) = u 0 ( x ) u_t+\alpha u u_x+\beta u_{xx}+\gamma u_{xxx}=0\\ u(0,x)=u_0(x) ut+αuux+βuxx+γuxxx=0u(0,x)=u0(x)
regular:
chaotic:
Navier-Stokes equation
w t + u ⋅ ∇ w = 1 Re Δ w , in [ 0 , T ] × Ω , ∇ ⋅ u = 0 , in [ 0 , T ] × Ω , w ( 0 , x , y ) = w 0 ( x , y ) , in Ω , \begin{aligned} {w}_{t}+\boldsymbol{u}\cdot\nabla{w}& =\frac{1}{\text{Re}}\Delta w,\quad\text{in}[0,T]\times\Omega, \\ \nabla\cdot\boldsymbol{u}& =0,\quad\text{in}[0,T]\times\Omega, \\ w(0,x,y)& =w_0(x,y),\quad\text{in}\Omega, \end{aligned} wt+u⋅∇w∇⋅uw(0,x,y)=Re1Δw,in[0,T]×Ω,=0,in[0,T]×Ω,=w0(x,y),inΩ,
CausalPINN NS
总结
本文发现PINN的训练过程并不遵守实际物理过程，内点残差并不能真实反映网络的状况。于是作者通过加权的方式，使得PINN可以根据不同时间片的残差来自动调整训练的时间区域。目前也有一些基于时间半离散或者分片递推的方式，相比这些方式，本文的方法可以直接得到PDE在整个时空域上的解，更加方便。
相关链接：
-
原文：2203.07404] Respecting causality is all you need for training physics-informed neural networks (arxiv.org)
-
原文代码：PredictiveIntelligenceLab/CausalPINNs (github.com)