华为边缘计算解决方案（当生物计算领域）

机器之心报道

演讲：王紫东

编辑：华卫

疾病治疗和新药研发当中，多尺度建模发挥着重要作用。而当传统科学计算方式遇到生物医药这类没有方程式可寻的问题时，亟待 AI 技术的加持。

9 月 2 日，在 2022WAIC 上海生物计算论坛上，华为中央软件院昇思MindSpore 开源项目架构师王紫东发表主题演讲《MindSpore AI 科学计算实践》。演讲中，他主要介绍了传统生物计算领域中的发展瓶颈以及 MindSpore 在这方面所做的工作。

以下为王紫东在 2022WAIC 上海生物计算论坛上的演讲内容，机器之心进行了不改变原意的编辑、整理：

大家下午好！很高兴有机会和大家一起分享 MindSpore 在 AI 生物计算方面做的实践。MindSpore 之前是传统的 AI 框架，目前正向 AI 科学计算融合的通用计算架构演进。

总体上讲，我们的实践是在高毅勤老师课题组指导下，共同合作完成的。研究中，我们主要从基于力场的分子模拟入手。

分子模拟的应用与挑战

借助分子模拟技术，我们可以模拟化学反应、蛋白质折叠以及同一物质在不同状态下的相变等；在行业之中，分子模拟也是比较重要的技术，如材料设计、药物设计、化工设计等场景都离不开分子模拟。

同时，基于力场的分子模拟也是微观世界探索中的关键问题。针对这样的问题，研究者做了很多的探索，目标是在更大的空间尺度和时间尺度上得到精确的分子模拟，以揭示微观规律。

针对微观世界，最有效的武器即为量子力学。保罗狄拉克曾说，「在有了量子力学之后，大部分物理化学现象的基本定律是完全已知的，困难在于实际应用中方程过于复杂无法求解。」

量子力学本质上是求解薛定谔方程，但薛定谔方程没办法处理较复杂的体系，只要粒子数大于 2 就没有办法求解。所以只能进行近似求解。学界为了针对更大体系进行分子模拟，发展出了很多方法，如分子动力学、粗粒化动力学和连续介质力学。

虽然能模拟的分子体系越来越大，但是精度在逐渐下降。以分子动力学为例，它利用第一性原理数据或者实验数据进行参数拟合，得到体系势能函数，也就是我们常说的分子力场。这种方法算得快，但精度很受限制，没有办法模拟精确的化学反应。

事实上，分子模拟一直处于模拟精度和效率无法兼得的状态，我把它归结为维度灾难。在这样的瓶颈下，现在已经有一些方法在解决这个问题，最主要的就是 AI。尤其是机器学习在解维度灾难问题上很有经验，也产生了一定结果，所以传统科学计算领域以后可能会成为 AI 的主战场。

「AI 分子模拟」的成就

现在，AI 和分子模拟已经开始进行融合，并且产生了有突破性的工作。这里我们看两个案例。

第一个案例 DeePMD，它主要是用 AI 神经网络去拟合分子力场，训练数据来自于第一性原理计算数据。DeePMD 兼顾精度和性能，并且获得了 2020 年的戈登贝尔奖。

案例二是 AlphaFold2，这是非常有突破性的成果，它更偏向 AI 思路，可以从海量数据中直接学习序列至结构的高维映射，效率比较高。令人惊喜的是，它的精度可以与实验精度相媲美，并被评为十大科技进步之首。

新一代分子模拟软件

AI 方法已经改变了传统计算的范式，成为一种新的可能性，软件上也要有相应的支持。我们来看一下现有软件能不能较好地支持 AI 新的范式，其实答案可能并不是很乐观。

主流分子模拟软件有这么几个特点：第一是开发时间长，用户依赖性强；第二是几乎都由西方发达国家开发，由中国人开发的模拟软件极少且用户有限；第三是模拟框架老旧，灵活性差，如需添加新算法，往往需要对程序代码进行大幅改动；第四是程序多用 C/C ，甚至 Fortran 语言编写，难以兼容目前主流以 Python 为前端语言的 AI 框架。

要想支持这种新范式，可能还是有很多路要走，现在做得并不是特别好。

在这样的背景下，MindSpore 和高毅勤老师课题组一起开发 MindSPONGE，我们把它称之为新一代分子模拟软件。

体系架构如何？

这是我们已知的第一个根植于 AI 计算框架中的分子软件，打破了传统 AI 和传统分子模拟计算的界限，并对它们的架构做了统一化的建模，取得了较好的效果。

我们发现，分子模拟和深度学习的训练虽然看着好像没什么关系，但实际上有相似的计算逻辑。比如，深度学习训练是要优化一个损失函数，分子模拟可能跟它是差不多的，只不过固定的是参数，参数指的是分子力场，再去优化以空间坐标为主的样本。

进行模块化拆分后，我们把 MindSPONGE 分为三大模块：系统建立、力场配置和迭代更新。我们有这么几个优势：1. 易用，无需推导原子受力的解析表达式；2. 高效，天然支持高通量模拟；3. 模块化，兼容 AI 势能函数。

有哪些案例实践？

这里举几个分子动力学的案例。分子动力学主要是知道分子下一步怎么动，这里面的核心就是分子力场。不管是构建分子力场，还是使用分子力场，MindSPONGE 都提供了非常方便的使用接口。

比如，克莱森重排有七元环与三元环转换的动力学过程，用我们自己训练的 AI 力场，以第一性原理精度去做动力学的迭代，精度及效率会优于传统方法。

另外，我们对新冠病毒 Delta 变异毒株也进行了模拟。模拟结果发现，Delta 变异时，其氢键数目和接触面积相对野生型毒株有一定增加，这可能就是 Delta 毒株传播力强的原因。

我们整体依托昇思MindSpore，使能 MindSpore 核心技术图算融合。MindSPONGE 的性能相较于其它框架有一定的优势。

DeepMind 开源了 AlphaFold2，这对于业界在该领域发展具有巨大的推进作用，但是其并没有开源自己的训练代码，也没有提供对应的数据集。同时，对于开源出来的网络权重，其有着非商用的 License。对于科研而言没有风险，但对于商业应用的话，存在一定的侵权的问题。

算法层面主要有两个问题，问题一：AlphaFold2 依赖于共进化信息，当共进化信息缺失严重时，蛋白结构预测精度会大幅下降。问题二：当存在多个模型推理出来的蛋白结构时，没有统一的算法评估这些结构的质量，这会导致无法选择最优的结构。

软件层面的问题主要为算法的内存消耗，算法占用的显存与序列长度是立方关系，如果基础软件内存优化做得不好，很容易导致显存爆炸。同时，推理时性能主要消耗在共进化信息检索。

针对上述不同层面的问题，我们协同高毅勤老师课题组一起进行了一些探索。

首先，我们复现了 AlphaFold2 的推理和训练，还构建了数据集，并把它开源出来，大家可以基于此数据集训练自己的模型。对于处理性能，如最耗时的前端搜索 MSA，我们都进行了优化，这项工作我们是在 4 月份完成的，当时去参加了 CAMEO 比赛，取得了较好的结果。

第二步，针对算法上依赖共进化信息的问题，我们利用了人工智能领域生成模型的思路，同时借鉴了 diffusion model 的模型构建思路，针对一些孤儿序列，构建了对应的 MSA 信息，取得了较好的效果。

最后，我们在蛋白质结构如何评估、如何挑出最好的蛋白质方面做了一些工作，表现也是良好的，比赛成绩优异。这项工作会在近期进行开源，如果大家感兴趣可以关注一下。

蛋白质结构预测和分子动力学这两个案例有意思的地方在于，可以联动。很多时候，结构并不符合物理的规律，可能很多蛋白质是可以继续优化的，很多位置是有冲突不合理的，这时需要分子动力学模拟来做进一步的动作，也就是蛋白质结构松弛。

当前 AlphaFold2, 在网络推理使用 JAX，分子动力学过程则使用 OpenMM 进行。由于 MindSPONGE 有统一建模的体系和工具，能够把蛋白质结构预测做一体化的推理，会给出更合理的结构。

还能提供什么能力？

此外，模型的内存和性能还有很多值得优化的地方，MindSPONGE 也提供了相应的解决方案，我们称之为 SOMAS 技术。该技术主要是并行拓扑内存约束优化器，可以对所有内存进行统一统筹管理，效果还是比较明显的。

性能方面，MindSPONGE 具有图算融合的特性，可以把小的细碎算子融成大算子，模型的整体性能会因此有提升。

传统优化改进到元优化

这里，对 MindSPONGE 的未来进行一些展望。

基于分子力场的模拟，分子力场肯定是最重要的。传统优化方式主要还是靠人工经验，比如先有力场，就可以做环境的模拟，模拟出的分布很可能不符合实验结果，这样就会产生误差；传统方式是通过人进行分析，是不是有些参数可以改一改，再去尝试一下。

相对来讲，这样比较盲目。MindSPONGE 采用的则是元优化方式。

现在，我们还有很多东西没有做完，后续希望把流程都打通，也希望和大家一起合作构建。这里感谢高毅勤老师课题组，也感谢所有社区贡献的小伙伴们！

最后， MindSpore 不只针对生物计算方面有相应的实践，我们在汽车、能源、气象、航空航天、EDA、材料、金融等方面都有规划。希望通过支撑更多行业应用，帮助合作伙伴商业落地，完善我们国产基础软硬件的能力。目前已经有的，主要是电磁仿真、流体仿真、分子模拟。

我们并不懂行业的技术，如果没有行业专家的参与，我们根本不知道做的这些东西对不对，所以希望和大家多多合作。如果大家感兴趣，可以随时联系我，谢谢大家！

,