Efficient physics-informed neural networks using hash encoding

简介

在 PINN 中，一个偏微分方程问题的求解就需要一次对神经网络的优化才能获得，这使得求解 PDE 问题的成本可能很高。传统的 PINN 训练需要消耗 10000 个 epoch，其成本通常远高于传统的数值方法，这使得除了求解奇异方程之外，将 PINN 用于实际应用的吸引力较小。 PINN 成本高的主要原因是收敛所需的 epoch 数量，特别是对于复杂的函数解决方案。其中一些限制可归因于神经网络的低频偏差。随机初始化的多层感知器需要一段时间的训练才能找到通往潜在局部最小值（有时是数千个epoch）的路径。换句话说，考虑到神经网络空间通常具有高维性质，在找到局部最小值之前，它首先会在随机搜索中迷失。部分问题是由于解的坐标值（输入）受其维度（通常是三个空间维度再加上时间维度）的限制，对神经网络施加的影响很小，这迫使网络最初在参数空间中自由漫游，几乎没有受到任何指引。通过编码增加输入的影响，这一限制已得到部分解决，编码允许输入空间具有更高维度的表示，其中标量输入被向量替换。频率或位置编码给出的示例将相邻样本坐标的微小差异替换为位置向量表示中较大差异所表示的网络中更深刻的变化。

提高 PINN 训练的效率的方法大致可以分为以下三个组成部分：神经网络 (NN) 架构设计、损失函数设计和训练过程。本文重点关注神经网络的设计，特别是表示输入（使用嵌入层）。在某种程度上，PINN 可以被视为神经场表示任务（例如NeRF），其输入是空间坐标，输出是体素或物理场，但以偏微分方程损失作为训练信号。受嵌入方法成功的启发，特别是对神经场表示的输入坐标的编码，以及多分辨率哈希编码在 NeRF 上最近取得的巨大进展，引发了一个逻辑问题：可以利用哈希编码来加速 PINN 训练吗？多年来用作输入坐标加密形式的哈希编码可以提供更本地化的坐标值表示。这一多分辨率特征为图像神经网络功能表示（NeRF）的收敛提供了相当大的加速。 1000 个 epoch 的收敛率降低到 100 以下。这是一项重大改进，使 NeRF 得以实际使用。然而，此类表示中涉及的监督训练不需要计算导数，而 PINN 中的 PDE 损失函数则需要计算导数。

本文将哈希编码引入 PINN，并研究了它通过减少收敛所需的 epoch 数来解决 PINN 的成本限制。本文还研究了通过哈希编码自动微分的适用性。具体来说，本文使用有限差分方法来确保具有哈希编码的 NN 函数导数的稳定计算，并将其用于 PINNs 训练。在数值示例中，在几个偏微分方程上测试了本文的方法，以证明其在显着降低 PINN 成本方面的潜力，并分享了可能遇到的限制。

方法

PINN

考虑 $n$ 维 $\Omega \subseteq {\mathbb{R}}^n$ 和边界 $\partial \Omega$ 的连通域，一般的时间相关 PDE 可以定义为：
$$u_t(\mathbf{x})+S(u(\mathbf{x}),a(\mathbf{x}))=0,\mathbf{x}\in \Omega ,t\in [0,T]$$
其中 $t$ 和 $x$ 分别是时间和空间坐标，$S(u,a)$ 是非线性微分算子，$a\in A$ 表示偏微分方程的参数，例如系数和初始或边界条件，$u$ 表示我们想要解决的物理场。在普通 PINN 中，由可训练参数 $\theta$ 参数化的神经网络 $\Phi (\theta ,x,t)$ 经过训练，将输入坐标（包括与时间相关的方程的时间）映射到输出，输出代表物理场（例如， 速度、压力、涡度等）在输入坐标位置，满足以下方程：
$$\frac{\partial \Phi (\theta ,\mathbf{x})}{\partial t}+S(\Phi (\theta ,\mathbf{x}),a(\mathbf{x}))=0.$$
因此，我们可以在损失函数中使用该偏微分方程的均方残差以及任何初始或边界条件，
$$\mathrm{L}=\underset{\text{Interior PDE loss}}{\underbrace{\frac{1}{N_i}\sum _{j=1}^{N_i}{\left(\frac{\partial \Phi (\theta ,{\mathbf{x}}_j)}{\partial t}+S(\Phi (\theta ,{\mathbf{x}}_j),a({\mathbf{x}}_j))\right)}^2}}+\underset{\text{Supcrivesd loss on boundary}}{\underbrace{\frac{1}{N_b}\sum _{i=1}^{N_b}{\left(\mathcal{B}(\Phi (\theta ,{\mathbf{x}}_i))-u_b\left({\mathbf{x}}_i\right)\right)}^2,(3)}}$$
优化神经网络的参数 $\theta$，其中 $N_i$ 是域内收集点的数量，$N_b$ 是边界上的收集点的数量，$u_b$ 是边界值，$\mathcal{B}$ 是边界算子，表示域的导数或值。

哈希编码

对于函数学习，本文的目标是提高神经网络对偏微分方程的逼近质量，以及给定神经网络的训练速度。请注意，速度是本文的主要目标。一种聪明的方法是将输入查询（例如空间坐标和时间）编码到高维空间。在这里，本文使用 Müller 等人提出的哈希编码。哈希编码的总体思想是将多分辨率分解与哈希表相结合，以紧凑地表示 3D 形状。然后，复杂的 3D 现实世界可以通过小型神经网络和可训练的哈希编码参数来表示。

上图描述了根据 Müller 等人使用的图表建模的哈希编码机制。 模拟域 $x_i$ 中的每个样本都可以使用从低到高的 $S$ 级分辨率来描述。对于每个分辨率级别，如上图中的粉色或蓝色点，计算嵌入向量。具体来说，首先找到它的体素顶点（2D 情况下有 4 个顶点，3D 情况下有 8 个顶点），然后使用可训练的哈希表，其中包括每个分辨率级别的长度为 $L$ 的固定特征和哈希表大小 $T$，评估每个顶点相应的嵌入向量。然后，我们使用顶点向量的线性插值来获得 $x_i$ 在不同级别的嵌入向量。最后，$x_i$ 的哈希编码是来自不同级别的这些嵌入向量的串联。

具体地，给定编码级别S以及最细和最粗分辨率 $R_f$ 和 $R_c$，每个级别的分辨率 $R_s$ 由几何级数确定，如下：
$$\begin{array}{rl}{R}_s& :=\lfloor R_c\cdot b^s\rfloor ,\\ b& :=\exp \left(\frac{\ln R_f-\ln R_c}{S-1}\right).\end{array}$$
那么对于 $x_i$，它的顶点是 $\lfloor\mathbf{x}_{i,s}\rfloor:=\lfloor\mathbf{x}i\cdot R_s\rfloor $ 和 $\lceil\mathbf{x}{i,s}\rceil:=\lceil\mathbf{x}_i\cdot R_s\rceil $。对于粗分辨率，当顶点数 $(R^s+1{)}^d$ 小于哈希表 $T$ 时，哈希表可以提供一对一的查询。然而，为了更高分辨率，顶点和哈希表之间的映射是通过哈希函数来实现的
$$h({\mathbf{x}}_i)=\left(\underset{j=1}{\overset{d}{⨁}}{x}_{i,j}{\pi }_j\right)\mathrm{mod}T$$
其中 $⨁$ 是按位“异或”XOR，${\pi }_j$ 是唯一且大的素数。这种编码不仅提供了输入维度的紧凑表示，而且由于几乎即时的哈希表查找机制，计算复杂度为 $O(1)$，非常高效。

具有哈希编码的 PINN

本文的目标是利用哈希编码来加速 PINN 的收敛。然而，与 NeRF 应用不同，PINN 的损失函数需要输出场相对于输入坐标的导数。由于所提出的哈希编码包括线性插值，因此这些导数可能是不连续的，这导致评估不准确，尤其是在分辨率网格的边界附近，并且这些潜在的不连续导数对于哈希编码的高分辨率级别更为频繁。以一个简单函数 $f(x)=sin(x)$ 为例，其各种阶导数很容易获得，可以在一个简单的输入为 $x$ 的网络上测试自动微分（PINN 中经常使用）的性能，该网络函数经过训练以输出 $f(x)$ 的值。然而，这个简单的网络将包含一个哈希编码层。

如上图所示，可以观察到基于哈希编码的神经网络的导数值并不准确，其准确性很大程度上取决于哈希编码的超参数。具体来说，需要仔细选择最粗和最细的分辨率、编码级别以及哈希表大小，以减轻不连续性的影响，这也取决于搭配点。本文的观点是，与高分辨率哈希中的高频细节非常适合的哈希编码的强度将被自动微分计算中的这一弱点所抵消。例如，使用很多级别并包括最精细的分辨率将使整个神经网络（带编码的神经网络）更好近似训练样本点的值，但生成的函数将缺乏平滑性，导致不稳定的导数。这是用于哈希向量的线性插值的直接结果。因此，使用当前哈希编码实现通过自动微分评估的神经网络的导数是不稳定的。

为了寻求有效的实现，本文转而使用有限差分（FD）方法来计算导数，因为自动微分对于高阶导数来说也很昂贵。由于有限差分（顾名思义）计算有限长度上的导数，因此它相对不受点引起的导数不连续性的影响。然而，在处理突然变化的函数时，准确性可能会略有下降，这是 PINN 的普遍弱点。 FD 方法建立在泰勒级数展开的基础上。给定一个网格点 $x_i$，可以通过限制其泰勒级数展开的长度来近似其物理场$u(x_i)$，如下：
$$u\left({\mathbf{x}}_i+\Delta \mathbf{x}\right)=u\left({\mathbf{x}}_i\right)+{\Delta \mathbf{x}\frac{\partial u}{\partial \mathbf{x}}|}_{{\mathbf{x}}_i}+{\frac{\Delta {\mathbf{x}}^2}{2}\frac{{\partial }^{2}u}{\partial {\mathbf{x}}^2}|}_{{\mathbf{x}}_i}+\cdots$$
停在二阶精度，有限差分一阶和二阶导数由下式给出：
$$\begin{array}{rl}{\frac{\partial u}{\partial \mathbf{x}}|}_{{\mathbf{x}}_i}& \approx \frac{u\left({\mathbf{x}}_i+\Delta \mathbf{x}\right)-u\left({\mathbf{x}}_i-\Delta \mathbf{x}\right)}{2\Delta \mathbf{x}},\\ {\frac{{\partial }^{2}u}{\partial {\mathbf{x}}^2}|}_{{\mathbf{x}}_i}& \approx \frac{u\left({\mathbf{x}}_i+\Delta \mathbf{x}\right)-2u\left({\mathbf{x}}_i\right)+u\left({\mathbf{x}}_i-\Delta \mathbf{x}\right)}{\Delta {\mathbf{x}}^2}.\end{array}$$
在训练过程中，导数计算所需的网格点应输入神经网络以获得相应的场值。

如上图所示，具有哈希编码的 NN 的导数通常比自动微分的导数更准确，但同样需要仔细选择编码超参数。在这里，上图© 中 FD 方法出现了失败的情况，该案例是由于使用小型哈希表而导致的，这迫使神经网络学习区分不同位置的样本，从而导致导数计算的准确性下降。与自动微分方法相比，(b)中的二阶导数虽然不平滑，但其趋势与解析解一致。

实验

本文通过其在三个众所周知的偏微分方程（PDE）上的应用来展示该方法的有效性。在所有情况下，都使用 MLP 架构作为主干，使用 Tanh 作为激活函数。为了公平的效率比较，使用 FD 方法来获得普通 PINN 和带有哈希编码的 PINN 的导数。同时稍微增加了普通 PINN 的宽度，使可训练参数的数量几乎等于具有哈希编码的 PINN。本文在所有实验中使用 Adam 优化器和衰减学习率来训练这些神经网络。

对于每个测试，通过设置解决方案精度的阈值（通过神经网络预测的解决方案与参考解决方案之间的绝对误差）以停止训练并根据轮数和每个轮次的成本评估方法。由于 FD 方法在导数计算方面的灵活性，本文使用 tiny-cuda-nn 框架实现了神经网络以进一步加速训练过程。

Burgers 方程

首先，考虑一个具有代表粘性流体一维流动的狄利克雷边界条件的一维时间相关方程，称为 Burgers 方程，它是 PINN 中广泛使用的基准。控制方程如下：
$$\begin{array}{rl}& \begin{array}{rlrl}\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}& =\nu \frac{{\partial }^{2}u}{\partial x^2},& t\in [0,1],& x\in [-1,1],\end{array}\\ & u(t,-1)=u(t,1)=0,\\ & u(0,x)=-sin(\pi x),\end{array}$$
其中 $\nu$ 是粘度参数，此处设置为 $0.01\pi$。对于普通 PINN，使用具有三个隐藏层 { 96 , 96 , 96 } \{96,96,96\} {96,96,96} 的 MLP，而对于具有哈希编码的 PINN，使用大小为 { 64 , 64 , 64 } \{64,64,64\} {64,64,64} 的 MLP。学习率为 $1\text{e}-3$，在 3000、5000、7000 和 10000 epoch 时减少 0.8 倍。使用数值解作为评估预测准确性的参考。

上图 (a) 显示了所提出的方法和使用 12800 个配置点的普通 PINN 的收敛速度。如前所述，通过测量达到预定义精度阈值所需的 epoch 数来关注收敛速度，可以发现采用哈希编码的 PINN 可以在不到 2500 epoch 内达到此目标精度，而普通 PINN 需要近20000个epoch。这里考虑的阈值精度为与参考解等效的解上图 (b)

Helmholtz 方程

接下来，在二阶导数问题上测试偏微分方程，该问题由著名的亥姆霍兹方程给出，该方程描述了波动现象并在地震和电磁领域有很多应用。这里考虑亥姆霍兹方程的简单形式：
$$\begin{array}{rl}& \begin{array}{r}\frac{{\partial }^{2}u}{{\partial }^{2}x}+\frac{{\partial }^{2}u}{{\partial }^{2}y}+\lambda u-f(x,y)=0,\end{array}\\ & \begin{array}{rlrl}u(x,2)=0,& u(-2,y)=0,& u(x,-2)=0,& u(y,2)=0,\end{array}\\ & f=-{\left(a_1\pi \right)}^2\sin \left(a_1\pi x\right)\sin \left(a_2\pi y\right)\\ & -{\left(a_2\pi \right)}^2\sin \left(a_1\pi x\right)\sin \left(a_2\pi y\right)\\ & +\lambda \sin \left(a_1\pi x\right)\sin \left(a_2\pi y\right),\end{array}$$
其中 $f$ 是源函数，$u$ 是波场，$\lambda$ 是波数的平方，$a_1$ 和 $a_2$ 是控制源项正弦特性的参数。该方程存在解析解：
$$u(x,y)=sin(a_1\pi x)sin(a_2\pi y)$$
在本例中，本文对普通 PINN 使用具有三个隐藏层 { 144 , 144 , 144 } \{144,144,144\} {144,144,144} 的 MLP，而对于具有哈希编码的 PINN，使用 { 128 , 128 , 128 } \{128,128,128\} {128,128,128} 的 MLP。学习率为 $1.5\text{e}-3$，在 3000、5000 和 7000 epoch 时降低了 0.8 倍。统一采样 10000 个收集点来训练神经网络。

亥姆霍兹方程训练的收敛速度如上图 (a) 所示。可以观察到，使用哈希编码的 PINN 收敛速度更快，并且使用哈希编码的 PINN 的 PDE 损失和测试数据错误仍然可以减少。尽管如此，预测解和参考解（如上图 (b) 所示）看起来是相同的。

N-S 方程

最后，在动态流体中的一个著名方程——纳维-斯托克斯方程上测试了所提出的方法。具体来说，考虑不可压缩流体的情况，产生基于质量和动量守恒的两个控制方程，如下所示：
$$\begin{array}{rlr}{\partial }_t\vec{u}(x,y,t)+\vec{u}(x,y,t)\cdot \nabla \vec{u}(x,y,t)+\nabla p& =\frac{1}{Re}\Delta \vec{u}(x,y,t)+f(x,y),& x\in (0,1{)}^2,t\in (0,T]\\ \nabla \cdot \vec{u}(x,t)& =0& x\in (0,1{)}^2,t\in [0,T]\\ \vec{u}(x,0)& ={\vec{u}}_0(x)& x\in (0,1{)}^2\end{array}$$
其中 $Re$ 是雷诺数，在实验中设置为100，$\nabla$ 是散度算子，$\Delta$ 是拉普拉斯算子， $\vec{u}$ 是速度场，${\vec{u}}_0$ 是初始速度场， $p$ 是压力，$f$ 是外力，这里设置为零。普通 PINN 具有三个隐藏层 { 112 , 112 , 112 } \{112,112,112\} {112,112,112} ，相比之下，使用 { 64 , 64 , 64 } \{64,64,64\} {64,64,64} 作为具有哈希编码的 PINN。学习率为 $1.2\text{e}-3$，在 3000、5000 和 7000 epoch 时降低了 0.8 倍。统一采样 10000 个收集点来训练神经网络。

结果如上图所示，其中参考解是从数值求解器获得的。与之前的实验一样，所提出的方法具有快速收敛性，并且只需 2270 个 epoch 即可达到目标精度（1.5e-3）。然而，即使有 50000 个 epoch，vanilla PINN 也无法满足目标精度。这表明所提出的方法加速了训练并提高了准确性。

训练效率对比

上述实验表明，使用哈希编码的 PINN 可以在比普通 PINN 少得多的 epoch 内进行训练，从而实现良好的目标准确度。

在上表中，对此处使用的三个示例的两种方法进行了定量比较。可以发现采用哈希编码的 PINN 可以使用单个 NVIDIA A100 GPU 卡在 30 秒内求解这三个著名方程。

总结

本文将哈希编码引入 PINN 中，并使用有限差分方法克服了哈希编码带来的自动微分的不准确性，加速了 PINN 的训练过程。

之前一直觉得 NeRF 和求解 PINN 之间有共同之处，也在2022中科大计算机图形学暑期课程里看过哈希编码加速NeRF的那篇文章，但没有意识到在这点上是共通的，看来在这方面敏感度还是不太够。当然，将输入映射到高维这一方法也是十分有趣的话题，而本人对这方面其实不太感兴趣，或许我该多看点这方面的方法。

相关链接：

• 原文：[2302.13397] Efficient physics-informed neural networks using hash encoding (arxiv.org)