When and why PINNs fail to train A neural tangent kernel perspective

When and why PINNs fail to train A neural tangent kernel perspective

PINN方法下NTK定义

Neural tangent kernel （NTK）

神经网络的训练过程中，参数更新公式如下：
$${\mathbf{\theta }}_{k+1}={\mathbf{\theta }}_k-\eta \nabla L({\mathbf{\theta }}_k)$$
其中，${\mathbf{\theta }}_k$ 表示 $k$ 时刻的神经网络参数，$\eta$ 表示学习率，$\nabla L({\mathbf{\theta }}_k)$ 则表示神将网络在参数为 ${\mathbf{\theta }}_k$ 时的损失的导数。对上式移项可得：
$$\frac{{\mathbf{\theta }}_{k+1}-{\mathbf{\theta }}_k}{\eta }=-\nabla L({\mathbf{\theta }}_k)$$
其中，左侧部分类似于一个导数的离散形式，因此，可以写成如下形式：
$$\underset{\eta \to 0}{\lim }\frac{{\mathbf{\theta }}_{k+1}-{\mathbf{\theta }}_k}{\eta }=\frac{d\mathbf{\theta }(t)}{dt}=-\nabla L(\mathbf{\theta }(t))$$
也就是，当学习率无穷小时，神经网络参数的变化情况，也称为梯度流。

对于一般的神经网络，可以假设其损失函数如下：
$$L(\mathbf{\theta })=\frac{1}{2}\Vert \mathbf{y}(\mathbf{\theta })-\bar{\mathbf{y}}{\Vert }_2^2$$
其中，$\mathbf{y}(\mathbf{\theta })$ 为网络输出，$\bar{\mathbf{y}}$ 为真解，$\frac{1}{2}$ 为避免求导产生常数项。于是可以将上式转换如下：
$$\frac{d\mathbf{\theta }(t)}{dt}=-\nabla \mathbf{y}(\mathbf{\theta })(\mathbf{y}(\mathbf{\theta })-\bar{\mathbf{y}})$$
此时，可以得到：
$$\frac{d\mathbf{y}(t)}{dt}=\nabla \mathbf{y}(\mathbf{\theta }{)}^T\frac{d\mathbf{\theta }(t)}{dt}=-\nabla \mathbf{y}(\mathbf{\theta }{)}^T\nabla \mathbf{y}(\mathbf{\theta })(\mathbf{y}(\mathbf{\theta })-\bar{\mathbf{y}})$$
观察上式红色的部分，这部分可以写成一个矩阵的形式，以 $\mathbf{H}(\mathbf{\theta })$ 来表示，称之为神经正切核（NTK）：
$$\frac{d\mathbf{y}(t)}{dt}=-\mathbf{H}(\mathbf{\theta })(\mathbf{y}(\mathbf{\theta })-\bar{\mathbf{y}})$$
可以看到，神经网络的训练过程变成了一个线性常微分方程的形式。同时，已经有人证明了，当网络足够宽时，模型输出的雅可比矩阵不会随着训练过程发生变化，也就是说：
$$\nabla \mathbf{y}(\mathbf{\theta }(t))\approx \nabla \mathbf{y}({\mathbf{\theta }}_0)⟹\mathbf{H}(\mathbf{w}(t))\approx \mathbf{H}({\mathbf{\theta }}_0)$$
对于一个足够宽的神经网络，其神经正切核将不会随着训练而改变。

PINN中的NTK

在PINN中，同样有下式：
$$\frac{d\mathbf{\theta }(t)}{dt}=-\nabla L(\mathbf{\theta }(t))$$
令 $u(t)=u({\mathbf{x}}_b,\mathbf{\theta }(t))=\{u({\mathbf{x}}_b^i,\mathbf{\theta }(t)){\}}_{i=1}^{N_b}$ ，同时 $\mathcal{L}u(t)=\mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta }(t))=\{\mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta }(t)){\}}_{i=1}^{N_r}$ ，于是可以得到如下结果：
$$\left[\begin{array}{c}\frac{du({\mathbf{x}}_b,\mathbf{\theta }(t))}{dt}\\ \frac{d\mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta }(t))}{dt}\end{array}\right]=-\left[\begin{array}{cc}{\mathbf{K}}_{u\mathbf{u}}(t)& {\mathbf{K}}_{u\mathbf{r}}(t)\\ {\mathbf{K}}_{r\mathbf{u}}(t)& {\mathbf{K}}_{r\mathbf{r}}(t)\end{array}\right]\cdot \left[\begin{array}{c}u({\mathbf{x}}_b,\mathbf{\theta }(t))-g({\mathbf{x}}_b)\\ \mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta }(t))-f({\mathbf{x}}_r)\end{array}\right]$$
其中，${\mathbf{K}}_{ru}(t)={\mathbf{K}}_{\mathbf{u}r}^T(t)$ ，${\mathbf{K}}_{uu}(t)\in {\mathbb{R}}^{N_b\times N_b},{\mathbf{K}}_{ur}(t)\in {\mathbb{R}}^{N_b\times N_r},{\mathbf{K}}_{rr}(t)\in {\mathbb{R}}^{N_r\times N_r}$ ，其第 $i,j$ 个元素可表示如下：
$$\begin{array}{rl}({\mathbf{K}}_{uu}{)}_{ij}(t)& =⟨\frac{du({\mathbf{x}}_b^i,\mathbf{\theta }(t))}{d\mathbf{\theta }},\frac{du({\mathbf{x}}_b^j,\mathbf{\theta }(t))}{d\mathbf{\theta }}⟩\\ ({\mathbf{K}}_{ur}{)}_{ij}(t)& =⟨\frac{du({\mathbf{x}}_b^i,\mathbf{\theta }(t))}{d\mathbf{\theta }},\frac{d\mathcal{L}u({\mathbf{x}}_r^j,\mathbf{\theta }(t))}{d\mathbf{\theta }}⟩\\ ({\mathbf{K}}_{rr}{)}_{ij}(t)& =⟨\frac{d\mathcal{L}({\mathbf{x}}_r^i,\mathbf{\theta }(t))}{d\mathbf{\theta }},\frac{d\mathcal{L}({\mathbf{x}}_r^j,\mathbf{\theta }(t))}{d\mathbf{\theta }}⟩\end{array}$$
其中， $⟨\cdot ,\cdot ⟩$ 表示 $\theta$ 中所有神经网络参数的内积。例如，
$$({\mathbf{K}}_{uu}{)}_{ij}(t)=\sum _{\theta \in \mathbf{\theta }}\frac{du({\mathbf{x}}_b^i,\mathbf{\theta }(t))}{d\theta }\cdot \frac{du({\mathbf{x}}_b^j,\mathbf{\theta }(t))}{d\theta }$$
下文将使用 $\mathbf{K}(t)$ 来表示矩阵 $\left[\begin{array}{cc}{\mathbf{K}}_{\mathbf{uu}}(t)& {\mathbf{K}}_{\mathbf{ur}}(t)\\ {\mathbf{K}}_{\mathbf{ru}}(t)& {\mathbf{K}}_{\mathbf{rr}}(t)\end{array}\right]$ ，也就是PINN的神经正切核（NTK）。值得注意的是，${\mathbf{K}}_{\mathbf{uu}},{\mathbf{K}}_{\mathbf{rr}}$ 以及 $\mathbf{K}(t)$ 均为半正定矩阵。事实上，令 $J_u(t)$ 和 $J_r(t)$ 分别为 $u(t)$ 和 $\mathcal{L}u(t)$ 关于 $\theta$ 的雅可比矩阵。然后，可以观察到
$${\mathbf{K}}_{uu}(t)={\mathbf{J}}_u(t){\mathbf{J}}_u^T(t),{\mathbf{K}}_{rr}(t)={\mathbf{J}}_r(t){\mathbf{J}}_r^T(t),\mathbf{K}(t)=\left[\begin{array}{c}{\mathbf{J}}_u(t)\\ {\mathbf{J}}_r(t)\end{array}\right]\left[\begin{array}{c}{\mathbf{J}}_u^T(t),{\mathbf{J}}_r^T(t)\end{array}\right].$$

PINN方法下NTK分析

对于PINN方法下的NTK，作者证明了如下结论：

假设激活函数 $\sigma$ 是平滑的并且具有有界二阶导数 $\ddot{\sigma }$。然后对于只有一个隐藏层的全连接神经网络，当 $N\to \infty$ 时
$$\begin{array}{rl}& u(x,\mathbf{\theta })\stackrel{\mathcal{D}}{}\mathcal{G}\mathcal{P}(0,{\Sigma }^{(1)}(x,x^{\prime }))\\ & u_{xx}(x,\mathbf{\theta })\stackrel{\mathcal{D}}{}\mathcal{G}\mathcal{P}(0,{\Sigma }_{xx}^{(1)}(x,x^{\prime })),\end{array}$$
其中，$\mathcal{D}$ 表示分布收敛于。同时，
$${\Sigma }_{xx}^{(1)}(x,x^{\prime })={\mathbb{E}}_{u,v\sim \mathcal{N}(0,1)}[u^4\ddot{\sigma }(ux+v)\ddot{\sigma }(u{x}^{\prime }+v)].$$
上述定理证明，在PINN方法中，足够宽的神经网络会在初始化时的函数值和PDE残差之间引入联合高斯过程，这与普通神经网络的表现一致。

同时，作者还证明了，对于初始化时只有一个隐藏层的PINN，当$N\to \infty$ 时，PINN的 NTK $\mathbf{K}(t)$ 概率收敛于确定性限制核，即
$$\mathbf{K}(0)=\left[\begin{array}{cc}{\mathbf{K}}_{u\mathbf{u}}(0)& {\mathbf{K}}_{u\mathbf{r}}(0)\\ {\mathbf{K}}_{r\mathbf{u}}(0)& {\mathbf{K}}_{r\mathbf{r}}(0)\end{array}\right]\to \left[\begin{array}{cc}{\Theta }_{\mathbf{u}\mathbf{u}}^{(1)}& {\Theta }_{\mathbf{u}\mathbf{r}}^{(1)}\\ {\Theta }_{r\mathbf{u}}^{(1)}& {\Theta }_{r\mathbf{r}}^{(1)}\end{array}\right]:={\mathbf{K}}^{\ast }$$
另外，作者也证明了：

存在常数 $C>0$，使得网络的所有参数一致有界于 $t\in T$ , 即：
$$\underset{t\in [0,T]}{\sup }\Vert \mathbf{\theta }(t){\Vert }_{\infty }\le C$$
存在一个常数 $C>0$ 使得:
∫ 0 T ∣ ∑ i = 1 N b ( u ( x b i , θ ( τ ) ) − g ( x b i ) ) ∣ d τ ≤ C ∫ 0 T ∣ ∑ i = 1 N r ( u x x ( x r i , θ ( τ ) ) − f ( x r i ) ) ∣ d τ ≤ C \begin{gathered} \int_0^T\left|\sum_{i=1}^{N_b}\left(u(x_b^i,\boldsymbol{\theta}(\tau))-g(x_b^i)\right)\right|d\tau\leq C \\ \int_0^T\left|\sum_{i=1}^{N_r}\left(u_{xx}(x_r^i,\boldsymbol{\theta}(\tau))-f(x_r^i)\right)\right|d\tau\leq C \end{gathered} ∫0T​ ​i=1∑Nb​​(u(xbi​,θ(τ))−g(xbi​)) ​dτ≤C∫0T​ ​i=1∑Nr​​(uxx​(xri​,θ(τ))−f(xri​)) ​dτ≤C​
假设激活函数 $\sigma$ 是平滑的并且 $|\sigma (k)|\le C，k=0,1,2,3,4$，其中 $\sigma (k)$ 表示 $\sigma$ 的 $k$ 阶导数，可以得到：
$$\underset{N\to \infty }{\lim }\underset{t\in [0,T]}{\sup }\Vert \mathbf{K}(t)-\mathbf{K}(0){\Vert }_2=0$$
结合上式，可以得到：
$$K(t)\approx K(0)\approx K^{\ast },\forall t$$
这意味着PINN中，当神经网络足够宽的时候，正切核也基本不会发生变化。由于正切核在训练中几乎不发生变化，因此可以将训练过程写成如下形式：
$$\left[\begin{array}{c}\frac{du({\mathbf{x}}_b,\mathbf{\theta }(t))}{dt}\\ \frac{d{\mathbf{u}}_{xx}({\mathbf{x}}_r,\mathbf{\theta }(t))}{dt}\end{array}\right]\approx -\mathbf{K}(0)\left[\begin{array}{c}u({\mathbf{x}}_b,\mathbf{\theta }(t))-g({\mathbf{x}}_b)\\ u_{xx}({\mathbf{x}}_r,\mathbf{\theta }(t))-f({\mathbf{x}}_r)\end{array}\right]$$
或者：
$$\left[\begin{array}{c}\frac{du({\mathbf{x}}_b,\mathbf{\theta }(t))}{dt}\\ \frac{d{\mathbf{u}}_{xx}({\mathbf{x}}_r,\mathbf{\theta }(t))}{dt}\end{array}\right]\approx \left(I-e^{-\mathbf{K}(0)t}\right)\cdot \left[\begin{array}{c}g({\mathbf{x}}_b)\\ f({\mathbf{x}}_r)\end{array}\right]$$
同时，由于正切核是半正定矩阵，因此可以将其分解为 $\mathbf{K}(0)=Q^T\Lambda Q$ 的形式。于是上式可以改写为：
$$\begin{array}{rl}\left[\begin{array}{c}\frac{du({\mathbf{x}}_b,\mathbf{\theta }(t))}{dt}\\ \frac{d{u}_{x\mathbf{x}}({\mathbf{x}}_r,\mathbf{\theta }(t))}{dt}\end{array}\right]-\left[\begin{array}{c}g({\mathbf{x}}_b)\\ f({\mathbf{x}}_r)\end{array}\right]& \approx \left(I-e^{-\mathbf{K}(0)t}\right)\cdot \left[\begin{array}{c}g({\mathbf{x}}_b)\\ f({\mathbf{x}}_r)\end{array}\right]-\left[\begin{array}{c}g({\mathbf{x}}_b)\\ f({\mathbf{x}}_r)\end{array}\right]\\ & \approx -{\mathbf{Q}}^T{e}^{-\mathbf{\Lambda }t}\mathbf{Q}\cdot \left[\begin{array}{c}g({\mathbf{x}}_b)\\ f({\mathbf{x}}_r)\end{array}\right]\end{array}$$
上式等价于：
$$Q\left(\left[\begin{array}{c}\frac{du({\mathbf{x}}_b,\mathbf{\theta }(t))}{dt}\\ \frac{d{\mathbf{u}}_{xx}({\mathbf{x}}_r,\mathbf{\theta }(t))}{dt}\end{array}\right]-\left[\begin{array}{c}g({\mathbf{x}}_b)\\ f({\mathbf{x}}_r)\end{array}\right]\right)\approx -e^{-\mathbf{\Lambda }t}Q\cdot \left[\begin{array}{c}g({\mathbf{x}}_b)\\ f({\mathbf{x}}_r)\end{array}\right]$$
可以看出，上式中左侧的第 $i$ 个分量将以大约 $e^{-{\lambda }_{i}t}$ 的速率衰减。换句话说，内核的特征值表征了绝对训练误差下降的速度。特别是，与具有较大特征值的核特征向量相对应的目标函数的分量将被更快地学习。对于全连接网络，NTK 矩阵较高特征值对应的特征向量通常表现出较低的频率。

更一般地，梯度下降 $t$ 步后 PINN 的 NTK 由下式给出：
$$\mathbf{K}(t)=\left[\begin{array}{cc}{\mathbf{K}}_{uu}(t)& {\mathbf{K}}_{ur}(t)\\ {\mathbf{K}}_{ru}(t)& {\mathbf{K}}_{rr}(t)\end{array}\right]=\left[\begin{array}{c}{\mathbf{J}}_u(t)\\ {\mathbf{J}}_r(t)\end{array}\right]\left[\begin{array}{c}{\mathbf{J}}_u^T(t),{\mathbf{J}}_r^T(t)\end{array}\right]=\mathbf{J}(t){\mathbf{J}}^T(t)$$

$$\begin{array}{rl}\sum _{i=1}^{N_b+N_r}{\lambda }_i(t)& =Tr\left(\mathbf{K}(t)\right)=Tr\left(\mathbf{J}(t){\mathbf{J}}^T(t)\right)=Tr\left({\mathbf{J}}^T(t)\mathbf{J}(t)\right)\\ & =Tr\left({\mathbf{J}}_u^T(t){\mathbf{J}}_u(t)+{\mathbf{J}}_r^T(t){\mathbf{J}}_r(t)\right)=Tr\left({\mathbf{J}}_u(t){\mathbf{J}}_u^T(t)\right)+Tr\left({\mathbf{J}}_r(t){\mathbf{J}}_r^T(t)\right)\\ & =\sum _{i=1}^{N_b}{\lambda }_i^{\mathbf{u}\mathbf{u}}(t)+\sum _{i=1}^{N_r}{\lambda }_i^{r\mathbf{r}}(t)\end{array}$$

其中，${\lambda }_i(t),{\lambda }_i^{\mathbf{u}\mathbf{u}}(t),{\lambda }_i^{\mathbf{r}\mathbf{r}}(t)$ 分别表示 $\mathbf{K}(t),{\mathbf{K}}_{uu}(t),{\mathbf{K}}_{rr}(t)$ 的特征值。这表明总训练误差的整体收敛速度由 ${\mathbf{K}}_{uu}$ 和 ${\mathbf{K}}_{rr}$ 的特征值共同表征。同时，$u(x_b,\mathbf{\theta })$ 和 $$u_{xx}(x_r,\mathbf{\theta })$$ 的单独训练误差分别由 ${\mathbf{K}}_{uu}$ 和 ${\mathbf{K}}_{rr}$ 的特征值决定。上述观察促使作者给出以下定义：

对于正半定核矩阵 $K\in {\mathbb{R}}^{n\times n}$，平均收敛速度 $c$ 定义为其所有特征值 ${\lambda }_i$ 的平均值，即
$$c=\frac{{\sum }_{i=1}^n{\lambda }_i}{n}=\frac{Tr(K)}{n}$$
特别地，对于任意两个分别具有平均收敛速度 $c_1$ 和 $c_2$ 的核矩阵 ${\mathbf{K}}_1,{\mathbf{K}}_2$ ，如果 $c_1\gg c_2$ ，则可以说 ${\mathbf{K}}_1$ 支配 ${\mathbf{K}}_2$ 。

上图为作者在一维泊松问题上用具有100个神经元的单隐层神经网络进行的试验，从图中可以观察到 PINN 的 NTK 特征值快速衰减。这导致目标函数高频分量的收敛速度极其缓慢。因此作者认为，PINN 也受到光谱偏差的影响。此外，对于不同频率的函数（不同的 $a$），$\mathbf{K}$ 的特征值分布看起来相似，这可以启发性地解释 PINN 倾向于几乎同时学习所有频率，如 Lu 等人所观察到的

这里的另一个关键观察结果是， ${\mathbf{K}}_{rr}$ 的特征值远大于 ${\mathbf{K}}_{uu}$ ，即， ${\mathbf{K}}_{rr}$ 支配 ${\mathbf{K}}_{uu}$ 。因此，PDE 残差的收敛速度比拟合 PDE 边界条件的速度快得多，这可能会阻止网络逼近正确的解。根据作者的经验，高频函数通常会导致 ${\mathbf{K}}_{rr}$ 的特征值较高，但在某些情况下 ${\mathbf{K}}_{uu}$ 可以主导 ${\mathbf{K}}_{rr}$ 。作者认为 ${\mathbf{K}}_{uu}$ 和 ${\mathbf{K}}_{rr}$ 之间的这种差异是 PINN 经常无法训练和产生准确预测的关键根本原因之一。

损失加权方法

根据上文分析， ${\mathbf{K}}_{rr}$ 的特征值远大于 ${\mathbf{K}}_{uu}$ ，这可能会阻止网络逼近正确的解。因此，可以根据NTK理论来为损失函数动态加权，损失函数可以表示如下：
$$\begin{array}{rl}\mathcal{L}(\mathbf{\theta })& ={\mathcal{L}}_b(\mathbf{\theta })+{\mathcal{L}}_r(\mathbf{\theta })\\ & =\frac{{\lambda }_b}{2N_b}\sum _{i=1}^{N_b}|u(x_b^i,\mathbf{\theta })-g(x_b^i){|}^2+\frac{{\lambda }_r}{2N_r}\sum _{i=1}^{N_r}|r(x_r^i,\mathbf{\theta }){|}^2\end{array}$$
其中，${\lambda }_b$和 ${\lambda }_r$ 是一些超参数，可以在训练期间利用反向传播梯度统计数据手动或自动调整。这里，训练数据 $\{x_b^i,g(x_b^i){\}}_{i=1}^{N_b}$ 和 $\{x_r^i,\bar{f}(x_b^i)\}$ 可以对应于在梯度下降的每次迭代中随机采样的完整数据批次或小批次。可以推导出对应于上述损失函数的输出 $u(x,\theta )$ 和 $u(x,\theta )$的变化过程：
$$\begin{array}{c}\left[\begin{array}{c}\frac{du({\mathbf{x}}_b,\mathbf{\theta }(t))}{dt}\\ \frac{d\mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta }(t))}{dt}\end{array}\right]=-\left[\begin{array}{cc}\frac{{\lambda }_b}{N_b}{\mathbf{K}}_{uu}(t)& \frac{{\lambda }_r}{N_r}{\mathbf{K}}_{ur}(t)\\ \frac{{\lambda }_b}{N_b}{\mathbf{K}}_{ru}(t)& \frac{{\lambda }_r}{N_r}{\mathbf{K}}_{rr}(t)\end{array}\right]\cdot \left[\begin{array}{c}u({\mathbf{x}}_b,\mathbf{\theta }(t))-g({\mathbf{x}}_b)\\ \mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta }(t))-f({\mathbf{x}}_r)\end{array}\right]\\ :=\mathbf{K}(t)\cdot \left[\begin{array}{c}u({\mathbf{x}}_b,\mathbf{\theta }(t))-g({\mathbf{x}}_b)\\ \mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta }(t))-f({\mathbf{x}}_r)\end{array}\right]\end{array}$$
观察上式可以发现，${\lambda }_b$和 ${\lambda }_r$ 的大小以及小批量的大小将对 $\mathbf{K}(t)$ 的奇异值产生至关重要的影响，训练误差 $u({\mathbf{x}}_b,\mathbf{\theta })$ 和 $\mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta })$ 的收敛速度也同样如此。基于这些观察，作者认为可以通过校准权重或批量大小来克服之前讨论的 ${\mathbf{K}}_{rr}$ 和 ${\mathbf{K}}_{rr}$ 之间的差异，使得 $u({\mathbf{x}}_b,\mathbf{\theta })$ 和 $\mathcal{L}u({\mathbf{x}}_r,\mathbf{\theta })$ 的每个分量在幅度上具有相似的收敛速度。由于操纵批量大小可能会涉及额外的计算成本（例如，可能会导致批量过大），因此这里作者选择固定批量大小并仅考虑根据调整权重${\lambda }_b$和 ${\lambda }_r$ 。

首先，方程 6.5 和 6.6 中的更新可以在梯度下降循环的每次迭代时进行，也可以以用户指定的频率进行（例如，每 10 个梯度下降步骤）。为了计算特征值之和，只需计算相应NTK矩阵的迹即可，这样可以节省一些计算资源。此外，作者指出 NTK $\mathbf{K}(t)$ 的计算与每次迭代时输入网络的训练数据点相关，这意味着每次迭代时内核的值不一定相同。然而，如果假设所有训练数据点都是从相同的分布中采样的，并且每次迭代时 NTK 的变化可以忽略不计，那么计算出的内核应该近似等于置换矩阵。因此，每次迭代时 $\mathbf{K}(t)$ 特征值的变化也可以忽略不计，因此上述算法的训练过程应该是稳定的。

实验结果

PINN NTK 收敛性实验

考虑具有如下强迫项的泊松方程：
$$\begin{array}{l}f(x)=-{\pi }^2\sin (\pi x),x\in [0,1]\\ g(x)=0,x=0,1.\end{array}$$
其精确解为 $u(x)=\sin (\pi x)$ ，其损失函数如下：
$$\begin{array}{rl}\mathcal{L}(\mathbf{\theta })& ={\mathcal{L}}_b(\mathbf{\theta })+{\mathcal{L}}_r(\mathbf{\theta })\\ & =\frac{1}{2}\sum _{i=1}^{N_b}|u({\mathbf{x}}_b^i,\mathbf{\theta })-g({\mathbf{x}}_b^i){|}^2+\frac{1}{2}\sum _{i=1}^{N_r}|u_{xx}({\mathbf{x}}_r^i,\mathbf{\theta })-f({\mathbf{x}}_r^i){|}^2\end{array}$$

上图中，左中两幅图为具有不同宽度的单隐层神经网络参数与NTK随着训练的变化情况，可以看到，随着宽度的提升，网络的参数和NTK变化幅度变小。右侧的图为宽度为500的神经网络在第0步和第10000步时NTK的特征值，从大到小排列的可视化，可以看出，其特征值几乎没有发生改变。

上图为基于三层神经网络做的实验，可以看出，结果与单层神经网络类似，因此推测先前基于单隐层神经网络的推导结果同样可以应用于多隐层神经网络上来。

PINN 的自适应训练

考虑具有如下强迫项的泊松方程：
$$\begin{array}{l}f(x)=-16{\pi }^2\sin (4\pi x),x\in [0,1]\\ g(x)=0,x=0,1.\end{array}$$
其精确解为 $u(x)=\sin (4\pi x)$ .

由于 NTK 在训练期间几乎没有变化。这意味着权重 ${\lambda }_b$和 ${\lambda }_r$ 是由NTK在初始化时确定的，因此它们可以被视为训练期间的固定权重。此外，由于先前实验已经知道在此示例中 ${\mathbf{K}}_{rr}$ 支配 ${\mathbf{K}}_{uu}$ 。因此，梯度下降第 $t$ 步的超参数 ${\lambda }_b$和 ${\lambda }_r$ 的更新规则可以简化为:
$$\begin{array}{rl}{\lambda }_b& =\frac{{\sum }_{i=1}^{N_b+N_r}{\lambda }_i(t)}{{\sum }_{i=1}^{N_b}{\lambda }_i^{uu}(t)}\approx \frac{{\sum }_{i=1}^{N_r}{\lambda }_i^{rr}(t)}{{\sum }_{i=1}^{N_b}{\lambda }_i^{uu}(t)}\approx \frac{Tr(K_{rr}(0))}{Tr(K_{uu}(0))}\\ {\lambda }_r& =\frac{{\sum }_{i=1}^{N_b+N_r}{\lambda }_i(t)}{{\sum }_{i=1}^{N_r}{\lambda }_i^{rr}(t)}\approx 1.\end{array}$$

上图为不使用动态损失权重和使用动态损失权重后的结果对比，明显可以看出，使用动态损失权重取得了更好的效果。（正如上图中高亮部分所示，由于NTK几乎不变，因此由NTK导出的损失权重在训练过程中其实也是固定的，这里的100是由于之前 ${\mathbf{K}}_{rr}$ 差不多是 ${\mathbf{K}}_{uu}$ 特征值均值的100倍。所谓的动态应该是指针对不同问题，其权重会发生变化）

上图为作者手动调整权重得到的不同权重下训练得到的 $L^2$ 误差。可以看到，当权重在100左右时，模型效果就比较好了，继续增大权重，后续方差变大，会增加训练过程的不稳定性。

一维波动方程

考虑如下一维波动方程：
$$\begin{array}{rl}& u_{tt}(x,t)-4u_{xx}(x,t)=0,(x,t)\in (0,1)\times (0,1)\\ & u(0,t)=u(1,t)=0,t\in [0,1]\\ & u(x,0)=\sin (\pi x)+\frac{1}{2}\sin (4\pi x),x\in [0,1]\\ & u_t(x,0)=0,x\in [0,1].\end{array}$$
其真解为：$u(x,t)=\sin (\pi x)\cos (2\pi t)+\frac{1}{2}\sin (4\pi x)\cos (8\pi t).$

由于其形式与先前推倒时不同，因此，重写如下：
[ d u ( x u , θ ( t ) ) d t d u t ( x u t , θ ( t ) ) d t ] : = K ~ ( t ) ⋅ [ u ( x b , θ ( t ) ) − g ( x b ) u t ( x u t , θ ( t ) ) L u ( x r , θ ( t ) ) ] \begin{bmatrix}\frac{du(\boldsymbol{x}_u,\boldsymbol{\theta}(t))}{dt}\\\frac{d\boldsymbol{u}_t(\boldsymbol{x}_{u_t},\boldsymbol{\theta}(t))}{dt}\\\end{bmatrix}:=\widetilde{\boldsymbol{K}}(t)\cdot\begin{bmatrix}u(\boldsymbol{x}_b,\boldsymbol{\theta}(t))-g(\boldsymbol{x}_b)\\u_t(\boldsymbol{x}_{u_t},\boldsymbol{\theta}(t))\\\mathcal{L}u(\boldsymbol{x}_r,\boldsymbol{\theta}(t))\end{bmatrix} [dtdu(xu​,θ(t))​dtdut​(xut​​,θ(t))​​]:=K (t)⋅ ​u(xb​,θ(t))−g(xb​)ut​(xut​​,θ(t))Lu(xr​,θ(t))​ ​
其中：
K ~ ( t ) = [ λ u N u J u ( t ) λ u t N u t J u t ( t ) λ r N r J r ( t ) ] ⋅ [ J u T ( t ) , J u t T ( t ) , J r T ( t ) ] , [ K u ( t ) ] i j = [ J u ( t ) J u T ( t ) ] i j = ⟨ d u ( x u i , θ ( t ) ) d θ , d u ( x u j , θ ( t ) ) d θ ⟩ [ K u t ( t ) ] i j = [ J u t ( t ) J u t T ( t ) ] i j = ⟨ d u t ( x u t i , θ ( t ) ) d θ , d u ( x u t j , θ ( t ) ) d θ ⟩ [ K r ( t ) ] i j = [ J r ( t ) J r T ( t ) ] i j = ⟨ d L u ( x r i , θ ( t ) ) d θ , d L u ( x r j , θ ( t ) ) d θ ⟩ . \begin{aligned} &\widetilde{\boldsymbol{K}}(t)=\begin{bmatrix}\frac{\lambda_u}{N_u}\boldsymbol{J}_u(t)\\\frac{\lambda_{u_t}}{N_{u_t}}\boldsymbol{J}_{u_t}(t)\\\frac{\lambda_r}{N_r}\boldsymbol{J}_r(t)\end{bmatrix}\cdot\begin{bmatrix}\boldsymbol{J}_u^T(t),\boldsymbol{J}_{u_t}^T(t),\boldsymbol{J}_r^T(t)\end{bmatrix}, \\ &\left[\boldsymbol{K}_u(t)\right]_{ij}=\left[\boldsymbol{J}_u(t)\boldsymbol{J}_u^T(t)\right]_{ij}=\left\langle\frac{du(\boldsymbol{x}_u^i,\boldsymbol{\theta}(t))}{d\boldsymbol{\theta}},\frac{du(\boldsymbol{x}_u^j,\boldsymbol{\theta}(t))}{d\boldsymbol{\theta}}\right\rangle \\ &[\boldsymbol{K}_{u_t}(t)]_{ij}=\left[\boldsymbol{J}_{u_t}(t)\boldsymbol{J}_{u_t}^T(t)\right]_{ij}=\left\langle\frac{du_t(\boldsymbol{x}_{\boldsymbol{u}_t}^i,\boldsymbol{\theta}(t))}{d\boldsymbol{\theta}},\frac{du(\boldsymbol{x}_{\boldsymbol{u}_t}^j,\boldsymbol{\theta}(t))}{d\boldsymbol{\theta}}\right\rangle \\ &\left[\boldsymbol{K}_r(t)\right]_{ij}=\left[\boldsymbol{J}_r(t)\boldsymbol{J}_r^T(t)\right]_{ij}=\left\langle\frac{d\mathcal{L}u(\boldsymbol{x}_r^i,\boldsymbol{\theta}(t))}{d\boldsymbol{\theta}},\frac{d\mathcal{L}u(\boldsymbol{x}_r^j,\boldsymbol{\theta}(t))}{d\boldsymbol{\theta}}\right\rangle. \end{aligned} ​K (t)= ​Nu​λu​​Ju​(t)Nut​​λut​​​Jut​​(t)Nr​λr​​Jr​(t)​ ​⋅[JuT​(t),Jut​T​(t),JrT​(t)​],[Ku​(t)]ij​=[Ju​(t)JuT​(t)]ij​=⟨dθdu(xui​,θ(t))​,dθdu(xuj​,θ(t))​⟩[Kut​​(t)]ij​=[Jut​​(t)Jut​T​(t)]ij​=⟨dθdut​(xut​i​,θ(t))​,dθdu(xut​j​,θ(t))​⟩[Kr​(t)]ij​=[Jr​(t)JrT​(t)]ij​=⟨dθdLu(xri​,θ(t))​,dθdLu(xrj​,θ(t))​⟩.​
于是，其损失权重可表示如下：
$$\begin{array}{c}{\lambda }_u=\frac{Tr({\mathbf{K}}_u)+Tr({\mathbf{K}}_{{\mathbf{u}}_t})+Tr({\mathbf{K}}_r)}{Tr({\mathbf{K}}_u)}\\ {\lambda }_{u_t}=\frac{Tr({\mathbf{K}}_u)+Tr({\mathbf{K}}_{u_t})+Tr({\mathbf{K}}_r)}{Tr({\mathbf{K}}_{u_t})}\\ {\lambda }_r=\frac{Tr({\mathbf{K}}_u)+Tr({\mathbf{K}}_{u_t})+Tr({\mathbf{K}}_r)}{Tr({\mathbf{K}}_r)}\end{array}$$

上图为结果对比，可以看见使用损失加权后，PINN学习到了真实解。

上图为三个NTK的特征值可视化，可以发现，随着训练的进行，特征值趋于稳定

上图为对权重的可视化，可以发现，随着训练的进行，损失权重很快就收敛了，在此后的训练过程中就几乎没有发生变化了。

总结

这篇文章从NTK的角度出发，首先证明了PINN中NTK的性质没有发生变化，随后根据分析过程中发现的不同损失项之间的NTK特征值具有数量级上的差异这一特性，提出应该为不同损失项添加权重。整个过程非常自然，并且包含了大量的数学推导。虽然是基于泊松方程以及单隐层神经网络做的推导，但作者很贴心地附上了多隐层神经网络的实验结果，利用实验表明上述结论大概率适用于多隐层神经网络。

但文中提到的 PINN也受频谱偏差影响 和 对于不同频率的函数，NTK的特征值分布看起来相似，这可以启发性地解释 PINN 倾向于几乎同时学习所有频率，如 Lu 等人所观察到的 。这里感觉有点矛盾？主要我不是很清楚一般的神经网络NTK特征值分布的情况，同时PINN中NNTK 矩阵较高特征值对应的特征向量是否也表现出较低的频率。

单就我的实验来看，PINN表现出的传播性会更强一些，对于不同的频率，没有特别明显的倾向性。

上图为一般的神经网络对 $\sin (\pi x)+\cos (x)$ 的拟合，可以明显看出神经网络先拟合了频率较小的部分，然后再去拟合高频的部分

上图为利用PINN对泊松方程的近似，其中，泊松方程的精确解同样为 $\sin (\pi x)+\cos (x)$. 从上图中可以看出，PINN似乎是一个“波浪”一个“波浪”地在近似，表现出一定的传播性，而对频率没有表现出明显的偏好。

相关链接：

• 原文：When and why PINNs fail to train: A neural tangent kernel perspective - ScienceDirect
• 代码：PredictiveIntelligenceLab/PINNsNTK (github.com)