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 ?{\mathbb{R}}^n$ 和邊界 $?\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{?\Phi (\theta ,\mathbf{x})}{?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{?\Phi (\theta ,{\mathbf{x}}_j)}{?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& :=?R_c?b^s?,\\ 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{?u}{?\mathbf{x}}|}_{{\mathbf{x}}_i}+{\frac{\Delta {\mathbf{x}}^2}{2}\frac{{?}^{2}u}{?{\mathbf{x}}^2}|}_{{\mathbf{x}}_i}+?$$
停在二階精度，有限差分一階和二階導數由下式給出：
$$\begin{array}{rl}{\frac{?u}{?\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{{?}^{2}u}{?{\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{?u}{?t}+u\frac{?u}{?x}& =\nu \frac{{?}^{2}u}{?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{{?}^{2}u}{{?}^{2}x}+\frac{{?}^{2}u}{{?}^{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}{?}_t\vec{u}(x,y,t)+\vec{u}(x,y,t)??\vec{u}(x,y,t)+?p& =\frac{1}{Re}\Delta \vec{u}(x,y,t)+f(x,y),& x\in (0,1{)}^2,t\in (0,T]\\ ??\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，$?$ 是散度算子，$\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)