我们需要什么样的机器学习仿真算法？

前几天在 Taichi GAMES 201 微信群里看到大家在吐槽某微信公众号报道的今年 ICML 的一篇文章 arxiv:2002.09405[1], 用图神经网络学习了SPH，MPM等图形学仿真里面几个常见的体系，既说不上比物理模拟 (physically-based simulation) 精确，又不如物理模拟高效，再加上公众号文章夸大其词称“特效师可以下班了”，把Science新闻报道说成发表在Science正刊，甚至群里的张心欣老师直接发了篇文章开怼。正好我PhD科研是做机器学习处理物理模拟体系和数据的，自己这几天也一直在想这个问题，作为AI和模拟交叉的研究者我们应该怎样应用和发展机器学习方法，于是在这做一个粗略的讨论和简单的总结。

我自己的背景是做分子模拟的，说来也巧，分子模拟两大方法MD（分子动力学）和MC（Monte Carlo），都和图形学有着很近的关系。MD即是对粒子的运动方程进行数值积分，对应于图形学里的众多仿真方法（“模拟”和“仿真”对应于同一个词simulation，大概由于领域不同的关系中文叫法不一样），MC用来计算采样intractable的物理体系积分（热力学里的配分函数），对应图形学渲染里的path tracing等一众Monte Carlo方法。近半年来开始逐渐接触计算机图形学，发现很多东西和研究生所学的其实是相通的，所以在这里从另一个比纯CS更偏自然科学角度来聊聊AI+模拟这个问题，算是抛砖引玉吧。我本科和PhD读的都是化学（因为我学了垃圾专业所以我所有观点都是垃圾的现在可以愉快地关掉这篇文章了~），CS半路出家还只算萌新，有错误或者不严谨的地方也欢迎指正。

我们要用什么机器学习模型解决物理模拟问题？

近20年来机器学习，尤其是深度学习在处理需要人类感知才能玩成的任务上取得了空前的成功，而自然科学正是人类智慧登峰造极的成果之一，于是使用机器学习(ML)的方法解决自然科学问题也成了近几年AI界的热点。ML里有一个很重要的概念，叫inductive bias，指的是我们对模型所学习的数据有着怎样的假设，已知这些数据有着怎样的特性，比如CNN卷积层里的平移不变性就是对图像数据的inductive bias。而对于物理体系来说，我们的inductive bias正是几百年来无数代人付出艰巨努力发现的物理规律。

更具体点地说，如果物理体系中存在一系列守恒量和对称性（这二者更深层次上反映的也是同一种规律），那么ML模型就也需要尽可能满足这些性质。物理定律之所以能保持上百年的生命力，是因为它们能够进行extrapolation，对已知范围之外的现象做出预测。在机器学习中这叫做泛化(generalization)，而泛化在物理体系中是至关重要的。这时候各种物理规律就一定程度上“regularize”了ML模型，使得遵守物理规律的模型有更强的泛化能力。

举个最简单的例子，能量守恒。这里提一下那个图网络模拟SPH的前期工作Lagrangian Neural Networks[2]，整个工作简单来说就是神经网络(NN)输入广义动量和坐标，输出体系的Lagrangian，通过欧拉—拉格朗日方程来进行训练。这样得到的体系运动轨迹，能量是守恒的。这就比输入动量和坐标，输出加速度的暴力拟合强得多，因为如果机器学习预测的轨迹甚至不满足能量守恒，那么它预测的结果对需要能量守恒物理模拟就没有任何价值。

Ref. 2 对双摆动力学模拟的结果。LNN能满足能量守恒，暴力拟合能量不守恒。

再说回那篇ICML，它的模型首先根据体系粒子的状态encode成图结构，在图里进行message passing，之后再decode回粒子状态，然而对于message passing来说这个图的结构不是已知的么？如果知道体系是长程力，那么它就一定是全连通图，又如果是SPH一类使用kernel的体系，那么它就是一个epsilon-threshold graph，于是我们就费不着用神经网络再推导一遍图的连接结构了。

Ref.1 使用encoder和decoder推导出哪些粒子之间有相互作用，然而如果模拟体系的相互作用形式已知，这一步多半不需要神经网络。

所以针对物理模拟问题的ML方法，需要对这些物理规律有“敬畏之心”，而不是片面地数据驱动，暴力破解，再造一遍轮子出来，否则得到的结果要么没有卵用，要么事倍功半。

机器学习的模拟需要解决什么问题？

我的PhD老板不止一次地说过，机器学习解决的物理和化学问题一定需要是难啃的，传统方法效果不好的，这样的工作才能有影响力。随着计算力的指数增长和科学计算方法的不断发展，科学计算和模拟产生了大量的数据，这的确给了基于数据的算法一个很大的舞台。那么什么样的物理模拟问题，适合或者有潜力用机器学习方法解决呢？这里随便提几个想法。

1. 填补不同尺度模拟体系之间的巨大差异

聪明的人类根据物理体系的性质，需要的精度，感兴趣的结果发明了不同尺度的模拟算法，小到基本粒子，大到宇宙天体，中间差了几十个数量级。然而不同尺度体系中可能存在着巨大的鸿沟，使得人们往往需要在精度和速度中做出鱼和熊掌的选择。这个鸿沟可以是计算复杂度的差异，也可以是空间尺度的差异。

比如，对化学反应的模拟一直是分子模拟研究中的重点。化学反应，即是分子的分解和重组。经典力学算法里一般把分子作为整体，如果要模拟化学反应的话，就需要人工预设哪些反应能够发生，而这个预设往往是不够完善的。要想更精确地模拟化学反应，就需要从分子再往下一级，考虑到原子内部原子核、电子的相互作用，用量子化学的手段从第一性原理出发计算。[3]然而量子化学算法知名的问题在于，它的时间复杂度随体系大小增加极其恐怖。最“简单”的Hartree-Fock算法的复杂度是O(N⁴),[4] 如果进一步考虑电子和电子之间的相互作用，算法的复杂度将会更高 (可以通过FFT等方法优化，但复杂度仍然远高于O(N²))。这个方向上的一个开创性工作是2007年发在PRL上的Behler-Parrinello Neural Network[5]，甚至比AlexNet还早。中心思想是模拟体系总能量等于各原子能量的加和，使用量子化学的计算结果训练来预测体系的势能面（Potential energy surface，即总能量和每个原子坐标之间的关系）。只要知道了势能，我们就可以求梯度得到力，从而进行动力学模拟。有了机器学习近似的加持，我们就可以对大量粒子参与的化学反应，进行长时间的模拟，[6]而这对只依赖第一性原理的方法来说是难以想象的。

我们实验室之前的工作 [Ref. 3], 用第一性原理的Monte Carlo模拟氮氧化合物的化学反应。我自己以前试着跑过类似的体系，每个step需要的时间非常非常久，如果把机器学习的方法利用起来，计算速度有很大的提升空间。

反过来说，正因为量子化学的计算成本太大，机器学习方法才有应用的价值。至于SPH，MPM这类算法，因为相互作用通过kernel (或者eulerian grid)来计算，在体系密度变化不大的情况下复杂度已经是O(N)了。一个O(N)自由度的系统，机器学习难道能把O(N)的算法加速到哪儿去？

再举一个空间尺度的例子。预测多相流体(mult-phase flow)和湍流在工程中有很重要的意义，比如船只螺旋桨卷起的微小气泡会对螺旋桨叶片造成巨大的损伤。而传统CFD（计算流体力学）方法一般要么无法预测多相流体的性质，要么需要大量的经验参数。而解决这个问题的最直接方法之一，便是精确到微观尺度，对气泡体系中每一个水分子进行分子动力学模拟。利用分子动力学模拟，我们可以在含有气泡的水(bubbly water)体系中发现负压强，而这在基于equation of state的CFD里是不允许的。[7]但是这个做法面临的关键问题是，CFD模拟和分子动力学模拟空间尺度的差别太大了。我们实验室现在的搞法是强行堆算力，在超级计算机上我们能够模拟0.5微米大小的bubbly water体系（未发表工作，65536个核要跑两三天），而想要进一步把体系加大到1微米，则需要动用国家实验室整台超算的全部资源。不说计算量，就光存储上百亿个分子的位置和速度，内存的开销就是几百TB。这基本上已经是全世界最大的分子动力学模拟了，然而可能还摸不到CFD的边（我不是专业做流体力学的，欢迎更正），这就意味着我们的结果只能“为CFD研究提供参考”，而不能真正地和更高尺度的模拟吻合。这时候，如果机器学习方法能够填补二者之间的鸿沟，对于我们对这一类物理体系的理解就格外有意义。机器学习模型能不能从分子动力学数据出发，抽象出可以应用于宏观体系的动力学规律或者状态方程？机器学习能不能帮助我们进行多尺度模拟，识别出体系中微观相互作用重要的部分，让模拟既接近微观的精度，又达到宏观的尺度？这些可能都是值得研究的问题。

研究螺旋桨上的气泡这样的特殊物理系统，需要填补宏观模拟和微观模拟之间的鸿沟（图出自Ref. 7）

2. 加速复杂相互作用体系的优化或采样

这个可能和图形学的关系远一点，有时候我们对物理体系进行模拟，不是关注它模拟的过程，而是想得到一个答案。最典型的例子就是蛋白质折叠。蛋白质从化学看就是氨基酸缩合成长链，从机器学习看就是序列数据，然而它能形成复杂的二级、三级、四级结构，发挥生物活性。所以分子动力学中，一大个方向就是研究蛋白质的折叠，找到这样高度折叠的结构，以及弄明白它们是如何形成的。本质上这个问题是一个优化问题，即如何能让这个蛋白质分子的自由能最小。说到优化CS方向的大家应该就很熟悉了，优化算法最怕的就是目标函数的性质不好，到处都是局部极小值。不巧的是，分子动力学体系如果使用经典力学模型计算分子的能量，偏偏就是这样一种情况，这让单纯使用物理方法模拟蛋白质折叠成为了一个困难的问题，也导致了一大批利用生成式模型预测蛋白质结构工作的出现。其中我个人印象比较深刻的是ICLR 2019的一篇Oral，既然蛋白质折叠的困难在于经典力学描述的能量性质不好，那用神经网络从已知结构学一个能量不就好了？[8]然后从这个“能量”出发，使用动力学方程模拟蛋白质分子的运动，就更容易地得到折叠好的结构。

Ref. 8 的机器学习模型，从氨基酸序列和分子结构得到一个机器学习的“能量”。

除了蛋白质折叠动力学这个例子以外，另外一个例子是对微观体系状态的分布进行采样。分子模拟中大家一般关心体系的物理性质，根据统计力学，这些物理性质可以对通过微观体系的分布求期望得到。说到采样问题自然就是Markov Chain Monte Carlo (MCMC)了，然而即使是MCMC算法，有时候也因为系统内复杂的相互作用，难以找到可能性最大的状态，这使得人们精心设计了各种各样的MC增强采样算法（我们实验室本行就是做这个的）。那么我们能不能借助机器学习，基于大量模拟数据做到一步到位，避免Markov chain被困在局部极小值里呢？去年Science上发表了一个开创性的工作，叫做Boltzmann generator（玻尔兹曼生成器），就在尝试解决这个问题。[9]平衡态分子体系在微观上服从玻尔兹曼分布，那么既然玻尔兹曼分布难以采样，需要各种MCMC的特殊技巧，那就干脆用一个可逆生成模型，把玻尔兹曼分布给变换成正态分布，如果模型训练的好，只要在这个变换后的正态分布里采样，就可以得到符合玻尔兹曼分布的微观状态了。

Boltzmann Generator (Ref. 9) 使用可逆神经网络模型把统计力学的Boltzmann分布转换成正态分布，从而避免MC采样的困难。

他们在今年ICML上又发了一个follow-up，讲怎么把物理系统里的等变性/对称性加入到Boltzmann generator的架构中，[10]不过似乎因为实现方法比较复杂，体系分子数的scaling不太好（文章里搞了个13个Lennerd-Jones分子的系统，专业做分子模拟的多半会开嘲讽），还需要期待之后的改进。除了模拟分子体系的Monte Carlo方法，在渲染中也可以找到类似的工作，通过可逆生成模型来加速path tracing中的采样，发表于去年的SIGGRAPH。[11]

3. 端到端地从直接观测学习模拟体系

目前为止，机器学习最擅长的就是模仿人类的感知，比如图像，声音和语言。虽然根据物理模型和物理定律模拟体系的变化不能也不应该靠感知（“人脑模拟物理运动”实在是无稽之谈），但是从现实观测中抽象出物理模型，却的确是一个一般只有人类能做到的感知任务。所以，机器学习的手段很有可能能够帮助我们从“第一手”数据中构建物理模型，进而进行模拟计算。比如去年NeurIPS的一个工作Hamiltonian Neural Networks[12], 其中一个重要的结果是利用哈密顿力学的运动方程训练神经网络，可以直接从单摆运动轨迹的视频中学习出体系的运动规律，预测结果显著好于暴力拟合的baseline。

除了动力学模拟之外，类似从“观测数据”建立物理模型的任务在渲染中也会遇到。渲染所使用的物体材质用参数化的着色器(shader)描述，而从渲染结果或者现实世界的照片中求出shader的组成和参数，就可以借助机器学习的方法。著名计算机科学YouTuber Two Minute Papers PhD期间的主要工作便是这样的一个材质合成(Materials synthesis)的问题。[13]首先把一张参考渲染图进行后期处理（俗称PS），然后通过一个神经网络，就能得出使用什么样的shader就可以得到看上去和处理后的图一样的材质。

这个工作的看上去有点meta-learning的感觉，假如把每个shader到渲染结果的转换看作一个task，那么这样的材质合成问题就成了meta-learning了，而它的网络结构也和DeepMind之前的一篇文章有些像[14]。不过他从shader到渲染结果的转换还是使用的神经网络，如果把他变得更物理一点，用一个differentiable renderer作为optimizer，再配合MAML之类的meta-learning方法，会不会更好呢？

另外一个视觉层面上例子是对模拟体系的风格迁移(style transfer)。今年的SIGGRAPH报道了一篇在拉格朗日视角的流体模拟中进行风格迁移的工作。[15]风格迁移本身是一个基于图像的深度学习技术，应用于物理模拟中，显然是使用规则网格的欧拉视角更符合直觉。然而如果把欧拉视角的每一帧当作图片处理，就会出现风格化之后帧与帧之间不一致的问题，为了确保这个时间一致性，欧拉视角的风格迁移就需要用一个滑动窗口把好几帧的数据逐个“对齐”。而这篇文章从混合欧拉/拉格朗日视角的模拟方法受到启发，巧妙地解决了这个问题：从每个流体粒子所带的物理量出发，做一遍粒子-网格转换(P2G)，再送入CNN计算风格化loss，这样风格迁移得到的是每个粒子各个时间的物理量变化，由于粒子的运动是连续的，我们只要把风格迁移的结果稍加regularization和平滑话，就可以得到帧与帧之间一致的风格化效果了。

用《神奈川冲浪里》和火焰的图片对流体模拟风格迁移，可以达到“水火交融”的艺术效果。

这一类的机器学习方法，主要受益的很可能是处在图形学“下游”的视觉艺术工作者。一个游戏美术可能不熟悉shader，但是她觉得游戏程序（职位）给的渲染图效果不好，拿PS修成了满意的样子还给程序，那么通过上面提到的机器学习模型，我们就能知道shader应该怎么改才能让效果更好看。又比如特效师设计了效果绚丽的魔法原型图，却没想好魔法释放的动态过程，这时候我们又可能可以借助上面的风格迁移方法，来把”科学侧”和“魔法侧”生动地结合。（换句话说也就意味着我这种3D美术视觉设计图形程序都是半吊子的离失业更近了一步）

讲了这么多机器学习的进展，终究还是要来泼一泼冷水的，唯物辩证法嘛。对于生产级别的物理模拟来说，比如HPC（高性能计算）的主流应用计算物理、计算化学、工业模拟（有限元，CFD等），以及电影的视觉特效，业界成熟的模拟软件都是为它们的应用经过数十年的开发并且高度优化的，比如scale到几万个gpu或者几十万个cpu核，而现有的机器学习框架直接运行物理模拟效率却不高。（这方面taichi开了一个好头）所以基于物理的模拟现在和近期的未来仍然会是科学计算和视觉特效中的主流。机器学习想要广泛地运用到物理模拟中，还需要经过时间的考验。

最后强烈安利一下GAMES201这个公开课~~~~~胡老师（他虽然和我同届，但是讲课已经强过我大学绝大多数老师了）讲的很清楚人也特别nice，自从上了这个课我才发现原来我研究生所学的和图形学有很近的关系，而不是在科学的象牙塔之外就一无是处。

1. ^Learning to Simulate Complex Physics with Graph Networks, arxiv:2002.09405 https://arxiv.org/pdf/2002.09405.pdf
2. ^Lagrangian Neural Networks, arxiv: 2003.04630 https://arxiv.org/pdf/2003.04630.pdf
3. ^First-Principles Monte Carlo Simulations of Reaction Equilibria in Compressed Vapors, ACS Cent. Sci. 2016, 2, 6, 409–415 https://pubs.acs.org/doi/abs/10.1021/acscentsci.6b00095
4. ^T. Helgaker, P. Jorgensen, J. Olsen, Molecular Electronic Structure Theory, Wiley, 2004
5. ^Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces，Phys. Rev. Lett. 2007. 98, 146401 https://journals.aps.org/prl/pdf/10.1103/PhysRevLett.98.146401
6. ^Neural network potential-energy surfaces in chemistry: a tool for large-scale simulations，Phys. Chem. Chem. Phys., 2011, 13, 17930–17955 https://pubs.rsc.org/en/content/articlepdf/2011/cp/c1cp21668f
7. ^Molecular Simulations Probing the Thermophysical Properties of Homogeneously Stretched and Bubbly Water Systems，J. Chem. Eng. Data 2019, 64, 3755−3771 https://pubs.acs.org/doi/pdfplus/10.1021/acs.jced.9b00284
8. ^Learning Protein Structure with a Differentiable Simulator, ICLR 2019 https://openreview.net/pdf?id=Byg3y3C9Km
9. ^Boltzmann generators: Sampling equilibrium statesofmany-body systems with deep learning, Science 2019, 365, 1001 https://science.sciencemag.org/content/365/6457/eaaw1147
10. ^Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities, ICML 2020 https://proceedings.icml.cc/static/paper_files/icml/2020/6711-Paper.pdf
11. ^Neural Importance Sampling, SIGGRAPH 2019 https://arxiv.org/pdf/1808.03856.pdf
12. ^Hamiltonian Neural Networks, NeurIPS 2019 https://arxiv.org/pdf/1906.01563.pdf
13. ^Photorealistic Material Editing Through Direct Image Manipulation，Eurographics Symposium on Rendering 2020 https://users.cg.tuwien.ac.at/~zsolnai/wp/wp-content/uploads/2019/09/pme.pdf
14. ^Meta-Learning with Latent Embedding Optimization, ICLR 2019 https://arxiv.org/pdf/1807.05960.pdf
15. ^Lagrangian Neural Style Transfer for Fluids， SIGGRAPH 2020 https://arxiv.org/pdf/2005.00803v1.pdf