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

Posted 2022-08-26 QbitAl

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技术在靶点挖掘、创新药物设计落地应用上的难关，让更多疾病可预警、可控制、可治愈！

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

此次 Colossal-AI 团队与百图生科合作的蛋白质单体与复合物结构预测最新方案，是面向大模型时代的通用深度学习系统 Colossal-AI 在生物医药领域应用落地的重要进展。

未来我们双方将会继续在生物计算大模型上有更深入的合作，助力深度学习在创新药物研发上的应用和落地。

随着人工智能对各个领域的高度渗透，各行业的智能化产业转型升级都离不开高性能AI的加持，而潞晨的AI大模型的高效训练和推理方案，刚好可为AI市场提供解决方案。

团队介绍

关于百图生科

百图生科(BioMap)是生物计算引擎驱动的突破创新药物研发平台,由百度创始人李彦宏发起创立,致力于将先进AI等信息技术(IT)与前沿生物技术(BT)相结合,构建独特的靶点挖掘及药物设计能力,聚焦于解析免疫调控机理,开发创新的治疗性药物,造福人类健康。

目前公司拥有近万平米的实验室，50+个靶点及药物研发资产组合。

关于潞晨科技

潞晨科技致力于解放AI生产力，打造面向大模型时代的通用深度学习系统Colossal-AI ，高效促进AI大模型落地应用。

核心成员来自美国加州伯克利、斯坦福、清华、北大、新加坡国立、南洋理工大学等世界一流高校，在国际顶级学术刊物或会议共发表论文近百篇，曾在谷歌、微软、NVIDIA、IBM、英特尔等头部科技公司任职。

公司成立即获得位列“清科中国早期投资机构30强”榜首的创新工场、真格基金超千万元种子轮投资。

传送门

GitHub地址：

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

参考链接：
https://www.technologyreview.com/2021/07/22/1029973/deepmind-alphafold-protein-folding-biology-disease-drugs-proteome/

*本文系量子位获授权刊载，观点仅为作者所有。

— 完 —