华为 MindSpore 架构师王紫东:当生物计算领域「遇灾」,MindSpore 做了何解?

再生纤维素纤维2022-09-27  8

华为 MindSpore 架构师王紫东:当生物计算领域「遇灾」,MindSpore 做了何解?

机器心脏报告

演讲:王子东

编辑:华为

多尺度建模在疾病治疗和新药开发中发挥着重要作用。而当传统科学计算方法遇到生物医学等无方程可寻的问题时,AI技术急需加持。

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

以下是王子东在2022WAIC上海生物计算论坛上的演讲内容。《机器之心》在不改变初衷的情况下进行了编辑整理:

大家下午好!很高兴有机会和大家分享MindSpore在AI+生物计算方面的实践。MindSpore曾经是一个传统的AI框架,现在正在进化为一个融合AI和科学计算的通用计算架构。

总的来说,我们的实践是在高教授和课题组的指导下共同合作完成的。在研究中,我们主要从基于力场的分子模拟入手。

分子模拟的应用和挑战

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

同时,基于力场的分子模拟也是微观世界探索中的一个关键问题。针对这一问题,研究人员做了大量探索,目标是在更大空尺度和时间尺度上获得精确的分子模拟,从而揭示微观规律。

微观世界最有效的武器是量子力学。保罗·狄拉克曾说,“量子力学问世后,大多数物理和化学现象的基本规律都完全为人所知。难点在于方程过于复杂,在实际应用中难以求解。」

量子力学的本质是求解薛定谔方程,但是薛定谔方程无法处理更复杂的系统,只要粒子数大于2,就无法求解。所以只能近似解决。为了模拟更大的系统,发展了许多方法,如分子动力学、粗化动力学和连续介质力学。

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

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

“人工智能+分子模拟”成果[S2/]

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

第一种情况,DeePMD,主要利用AI神经网络拟合分子力场,训练数据来自第一性原理计算数据。MD兼顾精度和性能,2020年获得戈登·贝尔奖。

第二个案例是AlphaFold2,这是一个非常突破性的成就。它更倾向于AI思维,可以直接从海量数据中学习序列到结构的高维映射,效率高。令人惊讶的是,它的精确度与实验相当,被评为十大科技进步之首。

新一代分子模拟软件

人工智能方法改变了传统计算的范式,成为一种新的可能性,这种可能性应该得到软件的支持。让我们看看现有的软件是否能更好地支持AI的新范式。其实答案可能不太乐观。

主流分子模拟软件有几个特点:一是开发时间长,用户依赖性强;第二,几乎都是西方发达国家开发的,中国人开发的仿真软件很少,用户有限。第三,仿真框架老旧,灵活性差。如果需要添加新的算法,往往需要对程序代码进行实质性的修改。第四,程序多采用C/C++甚至Fortran语言编写,难以兼容目前主流的以Python为前端语言的AI框架。

要支撑这个新范式,可能还有很长的路要走,现在也不是特别好。

在此背景下,MindSpore与高课题组共同开发了MindSPONGE,我们称之为新一代分子模拟软件。

架构呢?

这是我们知道的第一个植根于人工智能计算框架的分子软件。它打破了传统人工智能和传统分子模拟计算的界限,对两者的体系结构进行了统一建模,取得了良好的效果。

我们发现,虽然分子模拟和深度学习的训练看似毫无关系,但实际上它们有着相似的计算逻辑。比如深度学习训练就是优化一个损失函数,分子模拟可能与之类似,只是参数是固定的,参数参考分子力场,然后优化坐标在空之间的样本。

模块化拆分后,我们将MindSPONGE分为系统建立、力场配置、迭代更新三个模块。我们有几个优势:1。使用方便,无需推导原子力的解析表达式;2.高效率和对高通量模拟的自然支持;3.模块化,兼容AI势能函数。

有哪些案例做法?

这里有一些分子动力学的例子。分子动力学主要是了解分子下一步如何运动,而这其中的核心就是分子力场。无论是构建还是使用分子力场,MindSPONGE都提供了非常方便的用户界面。

比如克莱森重排有一个七元环和三元环转换的动态过程。用我们自己训练的AI力场做第一性原理精度的动态迭代,精度和效率会比传统方法更好。

此外,我们还模拟了新冠肺炎的Delta突变株。模拟结果表明,与野生型菌株相比,Delta菌株的氢键数量和接触面积增加,这可能是Delta菌株传递能力强的原因。

我们整体依赖MindSpore,使MindSpore的核心技术得以整合。与其他框架相比,MindSPONGE的性能具有一定的优势。

DeepMind开放了AlphaFold2,极大地推动了该领域行业的发展,但并没有开放自己的训练代码,提供相应的数据集。同时,开源网络权重拥有非商业许可。科研没有风险,商业应用有一定侵权。

算法层面主要有两个问题。第一个问题是AlphaFold2依赖于协同进化信息,当协同进化信息严重缺失时,蛋白质结构的预测精度会显著下降。问题二:当有多个模型推导出的蛋白质结构时,没有统一的算法来评价这些结构的好坏,会导致无法选择最优的结构。

软件层面的问题主要是算法的内存消耗。算法占用的内存与序列长度成三次关系。如果基础软件的内存优化没有做好,很容易导致内存爆炸。同时,推理的性能主要消耗在协同进化信息检索中。

针对上述不同层面的问题,我们与高课题组共同进行了一些探索。

首先我们重现AlphaFold2的推理和训练,并建立数据集,并作为开源开放。你可以根据这个数据集训练你自己的模型。对于处理性能,比如最耗时的前端搜索MSA,我们进行了优化。我们在4月份完成了这项工作,当时参加了CAMEO比赛,取得了不错的成绩。

其次,针对算法依赖于协同进化信息的问题,我们利用人工智能领域的模型生成思想,借鉴扩散模型的模型构建思想,对部分孤儿序列构建相应的MSA信息,取得了较好的效果。

最后,我们在如何评价蛋白质的结构和如何选择最佳蛋白质方面做了一些工作,我们的表现也不错,比赛成绩优异。这项工作将在不久的将来开放源代码。有兴趣的可以关注一下。

关于蛋白质的结构预测和分子动力学,有趣的是它们可以联系在一起。在很多情况下,结构不符合物理定律。也许很多蛋白质可以继续优化,很多立场是冲突的,不合理的。这时候就需要分子动力学模拟来做进一步的动作,也就是蛋白质结构松弛。

目前AlphaFold2在网络推理上使用JAX,在分子动力学上使用OpenMM。MindSPONGE有统一的建模系统和工具,可以对蛋白质结构预测进行集成推理,给出更合理的结构。

还能提供哪些功能?

另外,模型的内存和性能还有待优化,MindSPONGE也提供了相应的解决方案,我们称之为SOMAS技术。该技术主要是一种并行拓扑内存约束优化器,可以统一管理所有内存,效果明显。

在性能方面,MindSPONGE具有图形和计算融合的特性,可以将细分的小算子融合成大算子,从而提高模型的整体性能。

传统优化改进为元优化

在这里,我期待MindSPONGE的未来。

基于分子力场的模拟,分子力场肯定是最重要的。传统的优化方法主要依靠人工经验。比如环境的模拟可以先施加一个强场来做,模拟的分布很可能不符合实验结果,会导致误差。传统的方式是由人来分析,是否可以改变一些参数,再试一次。

相对来说,这是盲目的。MindSPONGE采用元优化方法。

现在我们还有很多未完成的事情,希望以后能把所有的流程都打通,也希望能和大家合作建设。在此,我要感谢高先生的课题组,以及社会各界做出贡献的朋友们!

最后,MindSpore不仅对生物计算有相应的实践,对汽车、能源、气象、航空空航天、EDA、材料、金融等都有相应的实践。我们希望通过支持更多的行业应用,帮助合作伙伴开展业务,提升我们国内的基础软硬件能力。目前主要有电磁模拟、流体模拟和分子模拟。

我们不了解行业的技术。没有行业专家的参与,我们根本不知道这些事情做得对不对,所以希望和大家多合作。如果你有兴趣,请随时联系我。谢谢大家!

剧终

授权请联系本微信官方账号。

投稿或寻求报道:content@jiqizhixin.com

转载请注明原文地址:https://juke.outofmemory.cn/read/1808935.html

最新回复(0)