谷歌 AI 最近一項研究表明，利用機器學習和硬件加速器能夠改進流體模擬，且不損害準確率或泛化性能。

流體數值模擬對于建模多種物理現象而言非常重要，如天氣、氣候、空氣動力學和等離子體物理學。流體可以用納維 - 斯托克斯方程來描述，但大規模求解這類方程仍屬難題，受限于解決最小時空特征的計算成本。這就帶來了準確率和易處理性之間的權衡。

不可壓縮流體通常由如上納維 - 斯托克斯方程來建模。 最近，來自谷歌 AI 的研究人員利用端到端深度學習改進計算流體動力學（CFD）中的近似，以建模二維渦流。對于湍流的直接數值模擬（direct numerical simulation, DNS）和大渦模擬（large eddy simulation, LES），該方法獲得的準確率與基線求解器相同，而后者在每個空間維度的分辨率是前者的 8-10 倍，因而該方法實現了 40-80 倍的計算加速。在較長模擬中，該方法仍能保持穩定，并泛化至訓練所用流以外的力函數（forcing function）和雷諾數，這與黑箱機器學習方法正相反。此外，該方法還具備通用性，可用于任意非線性偏微分方程。

論文地址：https://arxiv.org/pdf/2102.01010.pdf 該研究作者之一、谷歌研究員 Stephan Hoyer 表示：這項研究表明，機器學習 + TPU 可以使流體模擬加速多達兩個數量級，且不損害準確率或泛化性能。

至于效果如何呢？論文共同一作 Dmitrii Kochkov 展示了該研究提出的神經網絡與 Ground truth、基線的效果對比： 首先是雷諾數 Re=1000 時，在 Kolmogorov 流上的效果對比：

其次是關于衰變湍流（decaying turbulence）的效果對比：

最后是雷諾數 Re=4000 時，在更復雜流上的效果對比：

方法簡介 用非線性偏微分方程描述的復雜物理系統模擬對于工程與物理科學而言非常重要。然而，大規模求解這類方程并非易事。 谷歌 AI 這項研究提出一種方法來計算非線性偏微分方程解的準確時間演化，并且其使用的網格分辨率比傳統方法實現同等準確率要粗糙一個數量級。這種新型數值求解器不會對未解決的自由度取平均，而是使用離散方程，對未解決的網格給出逐點精確解。研究人員將受分辨率損失影響最大的傳統求解器組件替換為其學得的組件，利用機器學習發現了一些算法。 如下圖 1a 所示，對于渦流的二維直接數值模擬，該研究提出的算法可以在每個維度的分辨率粗糙 10 倍的情況下維持準確率不變，也就是說獲得了 80 倍的計算時間改進。該模型學習如何對解的局部特征進行插值，從而能夠準確泛化至不同的流條件，如不同受力條件，甚至不同的雷諾數（圖 1b）。 研究者還將該方法應用于渦流的高分辨率 LES 模擬中，獲得了類似的性能提升，在網格分辨率粗糙 8 倍的情況下在 Re = 100, 000 LES 模擬中維持逐點準確率不變，實現約 40 倍的計算加速。

圖 1：該研究提出方法與結果概覽。a）基線（direct simulation）與 ML 加速（learned interpolation）求解器的準確率與計算成本對比情況；b）訓練與驗證樣本圖示，展示出該模型強大的泛化能力；c）該研究提出「learned interpolation」模型的單時間步結構，用卷積神經網絡控制標準數值求解器對流計算中學得的近似。 研究者使用數據驅動離散化將微分算子插值到粗糙網格，且保證高準確率（圖 1c）。具體而言，將求解底層偏微分方程的標準數值方法內的求解器作為可微分編程進行訓練，在 JAX 框架中寫神經網絡和數值方法（JAX 框架支持反向模式自動微分）。這允許對整個算法執行端到端的梯度優化，與密度泛函理論、分子動力學和流體方面的之前研究類似。

研究者推導出的這些方法是特定于方程的，需要使用高分辨率真值模擬訓練粗糙分辨率的求解器。由于偏微分方程的動態是局部的，因此高分辨率模擬可以在小型域內實施。 該算法的工作流程如下：在每一個時間步中，神經網絡在每個網格位置基于速度場生成隱向量，然后求解器的子組件使用該向量處理局部解結構。該神經網絡為卷積網絡，具備平移不變性，因而允許解結構在空間中是局部的。之后，使用標準數值方法的組件執行納維 - 斯托克斯方程對應的歸納偏置，如圖 1c 灰色框所示：對流通量（convective flux）模型改進離散對流算子的近似；散度算子（divergence operator）基于有限體積法執行局部動量守恒；壓力投影（pressure projection）實現不可壓縮性，顯式時間步算子（explicit time step operator）使動態具備時間連續性，并允許額外時變力的插值。

「在更粗糙網格上的 DNS」將傳統 DNS 和 LES 建模的界限模糊化，從而得到多種數據驅動方法。 該研究主要關注兩種 ML 組件：learned interpolation 和 learned correction。此處不再贅述，詳情參見原論文。 實驗結果加速 DNS 一旦網格分辨率無法捕捉到解的最小細節，則 DNS 的準確率將快速下降。而該研究提出的 ML 方法極大地緩解了這一效應。下圖 2 展示了雷諾數 Re = 1000 的情況下在 Kolmogorov 流上訓練和評估模型的結果。

而就計算效率而言，10 倍網格粗糙度的情況下，learned interpolation 求解器取得與 DNS 同等準確率的速度也要更快。研究者在單個谷歌云 TPU v4 內核上對該求解器進行了基準測試，谷歌云 TPU 是用于機器學習模型的硬件加速器，也適用于許多科學計算用例。在足夠大的網格大小（256 × 256 甚至更大）上，該研究提出的神經網絡能夠很好地利用矩陣乘法單元，每秒浮點運算的吞吐量是基線 CFD 求解器的 12.5 倍。因此，盡管使用了 150 倍的算術運算，該 ML 求解器所用時間仍然僅有同等分辨率下傳統求解器的 1/12。三個維度（兩個空間維度和一個時間維度）中有效分辨率的 10 倍提升，帶來了 10^3/12 ≈ 80 倍的加速。 此外，研究者還考慮了三種不同的泛化測試：大型域規模；非受迫衰減渦流；較大雷諾數的 Kolmogorov 流。 首先，研究者將同樣的力泛化至較大的域規模。該 ML 模型得到了與在訓練域中同樣的性能，因為它們僅依賴流的局部特征（參見下圖 5）。 然后，研究者將在 Kolmogorov 流上訓練的模型應用于衰減渦流。下圖 3 表明，在 Kolmogorov 流 Re = 1000 上學得的離散模型的準確率可以匹配以 7 倍分辨率運行的 DNS。

最后，該模型可以泛化至更高的雷諾數嗎？也就是更復雜的流。下圖 4a 表明，該模型的準確率可以匹配以 7 倍分辨率運行的 DNS。鑒于該測試是在復雜度顯著增加的流上進行的，因此這種泛化效果很不錯。圖 4b 對速度進行了可視化，表明該模型可以處理更高的復雜度，圖 4c 的能譜進一步驗證了這一點。

與其他 ML 模型進行對比 研究者將 learned interpolation 與其他 ML 方法的性能進行了對比，包括 ResNet (RN) [50]、Encoder Processor-Decoder (EPD) [51, 52] 架構和之前介紹的 learned correction (LC) 模型。下圖 5 展示了這些方法在所有考慮配置中的結果?？傮w而言，learned interpolation (LI) 性能最佳，learned correction (LC) 緊隨其后。

對 LES 的加速效果 研究者已經描述了該方法在 DNS 納維 - 斯托克斯方程中的應用，但其實該方法是較為通用的，可用于任意非線性偏微分方程。為了證明這一點，研究者將該方法應用于 LES 加速。當 DNS 不可用時，LES 是執行大規模模擬的行業標準方法。 下圖 6 表明，將 learned interpolation 應用于 LES 也能達到 8 倍的 upscaling，相當于實現大約 40 倍的加速。

原文標題：谷歌AI利用「ML+TPU」實現流體模擬數量級加速

文章出處：【微信公眾號：深度學習實戰】歡迎添加關注！文章轉載請注明出處。

責任編輯：haq