比AlphaFold2快10倍！国产开源项目加速蛋白质预测，超长序列22分钟就搞定

新闻 2022-08-23 明敏

433 0

全球最快的复合物结构预测模型

22分钟推理出长达4K+的超长蛋白质序列，1分15秒搞定将近2K长的序列预测。

这是国产开源项目的最新战绩！

最近，由Colossal-AI团队(https://github.com/hpcaitech/ColossalAI)联合百图生科的蛋白质预测模型xTrimo Multimer，正式免费开源。

它在支持蛋白质单体（Monomer）和复合物（Multimer）的同时，还能大幅提升蛋白质预测速度。

面对2K到3K序列长度的预测任务，使用多卡推理，最高速度可达AlphaFold 2的11.15倍。

还能挑战4K以上的超长序列，而AlphaFold 2和OpenFold都受限于显存，无法完成这类任务。

目前，该项目已在GitHub上免费开源：

https://github.com/hpcaitech/ColossalAI/#xTrimoMultimer

怎么实现的？

本次解决方案背后优化加速技术来自Colossal-AI，这是一个旨在全面助力AI大模型工业化应用的通用深度学习系统。

xTrimo Multimer开源版模型通过引入Colossal-AI的大模型优化技术，显著降低蛋白质结构预测模型训练和推理的时间和经济成本，提高蛋白质结构预测模型领域的设计和部署效率。

这次Colossal-AI系统在AI大模型&生物医药落地应用，使其在该领域训练和使用更大的模型成为可能。

Colossal-AI是一个面向大模型时代的通用深度学习系统。

自开源以来，它在数个月内获得近五千颗Github Star，多次登上GitHub Trending榜首。

相关解决方案也已成功在生物医药、自动驾驶、云计算、零售、芯片等行业知名厂商落地应用，广受好评。

此次和百图生科联手，是结合了双方在高性能计算和生物计算领域上的前沿技术积累，提出了蛋白质单体与复合体的结构预测加速方案。

这对于精准高效探寻癌症、衰老等疾病的免疫规律，逐一攻破AI技术在靶点挖掘、创新药物设计应用上，都有所帮助，对于更多疾病的预警、控制和治愈都有深远意义。

众所周知，蛋白质结构预测是结构生物学领域最为重要的课题之一，也是人类理解基因翻译和蛋白质功能的重要手段。

蛋白质之间的相互作用，是蛋白质发挥生物学功能的重要结构基础。

但是由于蛋白质的多级结构和复杂的相互作用，使得精确预测三维结构这一关键问题十分有挑战性。

近年来，深度神经网络的成功使得人工智能在各个领域得到了广泛的应用。

自DeepMind发布的AlphaFold能够实现根据氨基酸序列精准预测蛋白质结构以来，生命科学领域掀起了使用人工智能预测蛋白结构的热潮。

AlphaFold能够实现端到端的从序列直接预测单体蛋白质三维结构，但在生物体内，多数蛋白质是以复合体的形式来发挥功能的。

因此，为了打破AlphaFold仅在预测单体蛋白方面表现优异的局限，DeepMind后来也发布了用于预测蛋白复合体结构的AlphaFold-Multimer模型。

在2022年3月，Colossal-AI 团队曾推出蛋白质结构预测模型AlphaFold的训练推理加速方案FastFold，用更低的成本将总体训练时间从11天减少到67小时，在长序列推理应用中实现超11倍的提升。

以此为起点，Colossal-AI 团队在进一步探索蛋白质结构预测领域的技术革新。

针对蛋白质单体(Monomer)与复合物(Multimer)结构预测的难题，本次Colossal-AI 团队联合百图生科提出行业内最新解决方案 xTrimo Multimer开源版模型，能够更好地理解蛋白互作关系，从而提升药物研发平台中靶点分析、蛋白质结构预测和模拟以及高精准抗体设计等方面的能力。

此外，推理阶段昂贵的经济和时间成本使得AlphaFold模型相关的研究和开发十分困难，尤其是面对长序列推理时，计算复杂度和内存消耗都面临着巨大挑战。

xTrimo Multimer开源版模型针对AlphaFold-Multimer模型中的计算特性，进行了针对性的CUDA优化，并且进行了Kernel Fusion。

相较于AlphaFold2和哥伦比亚大学的OpenFold，xTrimo Multimer开源版模型在单卡推理上性能有明显提升，推理速度分别提高1.58～2.14倍和1.14~2.23倍。

xTrimo Multimer开源版模型还支持超长序列的分布式推理。

由于使用了动态轴并行（Dynamic Axial Parallelism）技术，xTrimo Multimer可以高效地将计算和部分显存分配到不同的设备上，从而解决超长序列面临的计算和内存挑战。

在2K到3K的序列长度下，xTrimo Multimer使用多卡进行推理，相对于OpenFold和AlphaFold 2推理速度最高提升8.47倍和11.15倍，对比Uni-Fold2.0速度最高提升4.45倍。

xTrimo Multimer可以支持长达4K的序列推理，此时OpenFold和AlphaFold 2受限于显存无法完成推理，而xTrimo Multimer可以在20分钟左右完成。

百图生科首席AI科学家宋乐表示：

从对单个蛋白质结构的模拟，到对不同蛋白质之间的相互作用的识别，再到对蛋白质复合物的绘制，百图生科致力于解码、建模复杂人体免疫系统，开发突破创新药物，编程免疫系统，治愈多种免疫相关疾病。

此次xTrimo Multimer开源版模型的发布，是我们携手潞晨科技Colossal-AI 团队，借助其在高性能计算上的优势以及百图生科生物计算领域前沿的技术积淀，朝着百图生科xTrimo多模态生物计算大模型体系迈出的又一步。

正如我们一直以来所努力的方向，精准高效地探寻癌症、衰老等疾病的免疫规律，逐一攻破AI技术在靶点挖掘、创新药物设计落地应用上的难关，让更多疾病可预警、可控制、可治愈！

潞晨科技创始人尤洋教授表示：