今天介紹一篇文章，使用神經網絡將分子動力學第一原理計算推進到新的尺度和精度。分子動力學模擬一般被用來揭示化學反應的微觀物理過程，用於光合作用，新藥尋找，納米材料性質預測等等。
劉慈欣告訴我們，經典力學下的三體運動已經非常複雜，而量子力學統治下的多體相互作用，計算複雜度更是高到天際。分子動力學模擬要根據每一個時刻千萬粒子的位置計算勢能面（勢能的負梯度等於粒子受力），再根據受力和運動方程，更新所有粒子的座標與速度。此過程即便使用超級計算機，也只能計算很小的系統，遠離物理現實。
文中對比了各種大規模分子動力學模擬的原子數 （atoms), CPU計算核心數，GPU 計算核心數，以及單個原子單步模擬花費時間 TtS (秒/步/原子，或 s/step/atom）。
對比來看，2019年之前模擬的系統最大包含 100 萬個硅原子（速度= [公式] s/step/atom)，模擬速度最快達到 [公式] s/step/atom (9000個水原子系統）。
今天介紹的這篇文章，模擬了4.03億個水分子H2O，1.13億個銅原子Cu。速度也比之前最快的模擬提高了5000倍，爲 [公式] s/step/atom， (下表的最後兩行）。
計算方法：
作者們詳細介紹瞭如何對開源代碼 DeePMD-kit 二次開發，各種優化，加速計算。這些優化包括超級計算機集羣可擴展性優化， 面向 GPU 的數據結構優化，TensorFlow 算子優化與單雙精度混合計算加速優化。這篇博客也集中於介紹這些優化方法（相比其他文章，Deep Potential 的網絡架構極其簡單）。
從文章的配圖 2 也能看出，子圖（a）和（b）是優化之前的 Deep Potential 模型，而子圖 （c) - (g) 都在介紹如何面向GPU，更改數據結構，計算單雙精度以及 Tensorflow 算子，優化和加速計算。
子圖（a): 給定系統中每個粒子的位置，使用物理模型計算系統勢能, 積攢訓練數據。神經網絡爲三隱藏層全聯接網絡。輸入粒子云座標，預測系統勢能 U®。
子圖（b): 定義了一些類似於數據增廣的對稱性算子，即在平移，旋轉，重排等對稱性變換下使得系統勢能不變。這提供了一個從數據中學習物理對稱性的好方法。
子圖（c，d): 面向 GPU 的第一個優化。GPU 計算的底層很有特點，以並行編程語言OpenCL和Cuda 的混合術語舉例，一塊 GPU 一般有幾千個微處理器（Processing Element), 它們被分成幾十個小組，每個小組稱作一個Computing Unit (CU), CU 裏面的32或64個線程運算過程同步，被稱爲一個 Warp。計算的時候，如果同一個Warp裏不同線程走了不同的 If Else 通道，那麼此 Warp 內的所有線程都會既走 If 通道，又走 Else 通道。Warp 就像儀仗隊的一行士兵，隊列內所有士兵步調一致，不搞特殊。
第一個優化是減少分支，讓同一個Warp內的所有線程進行同樣的操作。做法是將一個粒子的鄰居們（也是粒子）按種類重排，這對應 子圖 c 和 d，重排後同一種類的鄰居粒子被同一Warp內的線程操作，步調完全一致，摒除分支。
如果同一種類的粒子個數小於 Warp 內線程數，則用空白補齊。如果多於 Warp 內線程數，則丟掉一些距離較遠的粒子（電磁相互作用，距離越遠，貢獻越小）。這也是爲何子圖 © 中除了按種類別 type 排序外，還要將同一種類的粒子按照距離 dist 從近到遠排序。
子圖 (e): 相比 Array of Structures (AoS), GPU 計算更喜歡 Structure of Arrays (SoA)數據結構，這樣GPU從顯存（Global Memory） 讀數據時速度更快。子圖 (e) 中對那些數據 Reshape 來 Reshape 去就是爲了將 AoS 轉化爲 SoA，讀到每個線程的 Private Memory 之後再轉化回來。
文中介紹了他們使用的另一個技巧，數據壓縮。將粒子種類 [公式] ，粒子距離 [公式] 和粒子 index = j 壓縮後保存在64位整數中。
基本上，這篇文章進行了非常精細的 GPU 並行加速優化，將所有計算耗時或者傳輸數據耗時的運算都移植到了 GPU 上進行。從 CPU 內存往 GPU 顯存移動數據非常緩慢，直接在 GPU顯存上開闢數據空間，在GPU上計算，減少內存-顯存之間的數據流動，是重要的GPU並行計算優化步驟。
當然，從 GPU 的 Global Memory（顯存）往線程的 Private Memory （寄存器）傳輸數據也有延遲，數據的 SoA 存儲格式能夠使用緩存加速讀取，對多線程共享的數據使用每個 Computing Unit 共享的內存 (Shared Memory)也能減小讀取延遲。目前未見文章在這一塊的優化。
子圖 (f): GPU上單精度(float)計算速度快於雙精度(double)計算速度。子圖 (f) 通過單雙精度混合計算，既保證精度，又加快運算速度。在一般的計算程序中，可多次嘗試，觀察哪些參數（數據）從雙精度改爲單精度浮點數後不影響計算結果。此研究發現對原子的位置和能量使用雙精度，對神經網絡的參數使用單精度，可以實現混合精度優化。混合精度相比雙精度加速1.5倍 （此處存疑，個人經驗GPU上單精度計算要遠快於雙精度），顯存節約 50%。
子圖 (g): 這一部分涉及對 Tensorflow 計算圖和算子的優化。此優化只針對當前物理問題。在勢能面的計算中，需要對大量粒子的貢獻求和，涉及很高很瘦的矩陣乘積 MATMUL 與求和 SUM。在此過程中，SUM 操作的計算需求劇增，Tensorflow 未對此做特殊優化。本文使用了 CUBLAS 提供的 General Matrix Multiply(GEMM) 來加速 x*W + b，Concat 和 Sum 的計算。另外， [公式] 的微分等於 [公式] , 作者們定製了兩個 kernel 來計算 TANH 和其梯度 TANHGrad。感覺回到了自己手寫一大堆 OpenCL kernels 解流體力學方程的日子。
子圖之外，文章還對不同計算節點的 MPI 通信進行了優化。MPI 節點間通信主要有兩個目的：（1）將系統分塊後，每個節點計算一個區塊內粒子云的運動，實現粗粒度並行。每個區塊周邊要有一定的輔助區域作爲邊界條件，輔助區域的粒子需要使用MPI通信從周邊格子抓取信息來更新。 （2）輸出時需要對所有區塊的能量求和，MPI_Allreduce 約化操作會因不同格子的阻塞同步帶來延遲，文章使用非同步 Reduce，並將求和頻率更改爲 20 步一次。
文章意義：
基於物理模型的計算 + 機器學習 + 高性能 GPU 並行集羣的共同操作，將超大系統的分子動力學模擬帶到了一個新時代。文章在最後簡單廣告了他們的下一篇物理文章，使用本文發展的大尺度 DeePMD-kit 來計算納米晶體銅顆粒的拉力變形。