今天介绍一篇文章，使用神经网络将分子动力学第一原理计算推进到新的尺度和精度。分子动力学模拟一般被用来揭示化学反应的微观物理过程，用于光合作用，新药寻找，纳米材料性质预测等等。
刘慈欣告诉我们，经典力学下的三体运动已经非常复杂，而量子力学统治下的多体相互作用，计算复杂度更是高到天际。分子动力学模拟要根据每一个时刻千万粒子的位置计算势能面（势能的负梯度等于粒子受力），再根据受力和运动方程，更新所有粒子的座标与速度。此过程即便使用超级计算机，也只能计算很小的系统，远离物理现实。
文中对比了各种大规模分子动力学模拟的原子数 （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 来计算纳米晶体铜颗粒的拉力变形。