迈入世界领先梯队!天壤自研蛋白质结构预测平台 TRFold,单张 GPU 16 秒实现精准预测

2021-12-20 14:08:33 浏览数 (1)

近日,天壤 XLab 发布蛋白质结构预测平台 TRFold,其最新版本的预测精度接近 AlphaFold2,达到世界领先水平

此外,TRFold 突破了 AlphaFold2 需要超大算力的瓶颈,采取权重共享的方式节约算力,预测绝大多数蛋白质链所需时间不超过 16 秒,以更小算力、更高效率推动 AI 生命科学的应用发展。

TRFold 基于 CASP14 蛋白质测试集的评估数据

绿色:真实结构,蓝色:预测结构

AI 攻克生物界五十年难题

蛋白质参与行使几乎所有的细胞功能,而为了行使这些特定功能,蛋白质必须折叠成对应的结构。因此,蛋白质的三维结构直接决定了他的功能,一旦三维结构被破坏,蛋白质功能随之丧失或改变,许多疾病就是由体内重要的蛋白质结构异常所导致。

“蛋白质折叠问题”一直是生物学中的重大挑战,已经困扰生物学家 50 多年,始终面临着难度大、成本高的局面,直到 AI 加入。

2020 年,DeepMind 开发的 AlphaFold2 在第十四届国际蛋白质结构预测竞赛(CASP14)上不仅击败了全球一百多个科研团队,其准确度更是达到原子级别,可与冷冻电子显微镜等实验技术相媲美。至此,单个蛋白质折叠预测问题得到基本解决。

TRFold:媲美 AlphaFold2

国产自研蛋白质结构预测平台

TRFold 不仅拥有媲美 AlphaFold2 的蛋白质结构预测精度,还在训练数据处理、内存优化、预测效率和模型表达等方面拥有明显优势。

预测精度接近 AlphaFold2

天壤 XLab 基于第十四届国际蛋白质结构预测竞赛(CASP14)的蛋白质测试集进行企业内测,TRFold 取得 82.7 分(TM-Score,评估蛋白质结构拓扑相似性的指标),已超过 Baker Lab 在 GitHub 公布的 RoseTTAFold 预测结果的 81.3 分,仅次于 AlphaFold2 的 91.1 分。

算力消耗约 AlphaFold2 的 1/32

TRFold 采取权重共享的方式以节约算力,在训练时,AlphaFold2 使用了 128 个 TPUv3 核心(约等于 256 张 GPU),TRFold 仅使用了 8 张 Nvidia RTX 3090 GPU,以极小算力取得接近 AlphaFold2 的成绩。

16秒精准预测

TRFold 采用 5 千万的参数循环多轨注意力网络,同时支持氨基酸残基间的距离预测和蛋白质全链结构预测,400 个氨基酸的蛋白链使用单张 Nvidia RTX 3090 GPU 的预测时间只需 16 秒。

TRFold:不是对 AlphaFold2 的简单复刻

更关注下游延展

除了对 AlphaFold2 的性能改进,天壤 XLab 还从自身经验出发,取得多项具有前瞻性的突破。

天壤 XLab 团队负责人苗洪江博士表示,“我们在研发 TRFold 的过程中充分考虑了后续的应用方向。比如 TRFold 有不同的模型版本,面对不同情况,端到端版本(end2end)用于快速生成结构,分段式版本用于大规模测算蛋白质氨基酸之间的距离和角度。

在后续构建蛋白质相互作用网络的过程中,涉及一个蛋白和另一个蛋白的相互作用,以及多个蛋白之间的相互作用,计算量呈指数级增长。因此,TRFold 以较小算力需求快速预测蛋白质结构对后续开展深入研究具有深远意义。”

充分挖掘数据价值

相比 AlphaFold2,TRFold 的算法模型仅采用少量的真实数据训练。即从多个宏蛋白质组序列库中寻找出蕴含更精准共进化信息的多序列排列,使模型在训练过程中能够获得对真实共进化信息更好的识别能力,从而取得对氨基酸残基距离和坐标更准确的预测结果,有望在近期加入蒸馏数据增强工作,进一步提高模型的预测和泛化能力。

多版本实现

在模型表达方面,不同于 AlphaFold2 仅采用端到端模型(end2end),TRFold 采用多版本实现。除了端到端模型之外,还有分段式结构模型即先预测氨基酸残基间的距离再生成三维结构,不同的模型实现可以分别应用在不同的场景中。

意味着在需要快速生成蛋白质结构时,可以快速使用端到端模型,在需要大规模测算蛋白质与蛋白质的相互关系时,分段式版本会有更优异的表现。

蛋白质与蛋白质的相互作用

除了单蛋白结构预测,TRFold 的另一重要功能是对蛋白质相互作用的预测。RoseTTAFold 和 AlphaFold-multimer 都是对已知相互作用的复合体进行结构模拟,无法判断两个或多个输入的蛋白质是否存在作用,而 TRFold 复合体预测模型能够首先对输入蛋白质链是否相互作用进行快速判别,再对其蛋白质复合体结构进行准确预测。

天壤认为,在蛋白质结构预测方面取得优异成绩只是第一步,更重要的是将在开发算法过程中积累的经验和技术,进一步应用于破解更具挑战的难题。

值得一提的是,曾在 CASP14 比赛期间,CASP 官方发起了一次针对新型冠状病毒肺炎(COVID-19)的蛋白质结构预测,天壤 XLab 基于 TRFold 提交的模型(nsp6-D2)预测结果被 CASP 选为六个“最具可信度模型”之一。

天壤 XLab 基于 TRFold 提交的 nsp6-D2 预测结果

被 CASP 选为六个“最具可信度模型”之一

面对团队取得的进展,天壤创始人兼 CEO 薛贵荣博士表示:“国内以后会涌现出一批在蛋白质结构预测领域取得优异成绩的公司,大家发展路径会各不相同,有的侧重构建平台,建设社区、生态,也有专注下游应用与落地等。

天壤通过两年半的时间掌握了蛋白质结构预测的核心技术,将以此为基础继续推进 AI 在蛋白质相互作用网络构建、蛋白质/抗体设计优化等下游领域的应用。希望能参与到国际最前沿的生命科学领域的应用中,开辟出一条 AI 生命科学的应用之路。”

天壤三部曲

AI 围棋、AI 交通、AI 生命科学

自创立以来,天壤始终专注研究复杂系统的决策难题,凭借对复杂问题的理解以及在深度学习领域的沉淀,先后破解了 AI 围棋、AI 交通和 AI 生命科学领域三大难题。

这些重大课题的突破,其背后都是由天壤通用人工智能平台提供技术支撑。天壤自主研发的迭代升级神经网络技术,可以大幅降低超大规模网络模型的算力消耗,能够以最小的计算资源,通过渐进式迭代,获得最佳算法模型,通过智能化方式最大限度地降本增效。

2018 年,天壤自主研发的深度强化学习AI围棋 TRGo,以 AlphaGo 1% 的计算资源击败了世界围棋冠军朴廷桓,朴廷桓曾经获得三次世界冠军,当时的个人围棋等级分为 3660 分,比柯洁领先 49 分。

此后团队开始积极探索落地场景,并率先在城市级交通调度领域取得了重大的进展。2018 年 9 月,天壤发布了 TRTraffic,是全球首个 7×24 小时在线的AI信号灯调控系统,接管了杭州全城高架路,同时使通行效率提高了 20%。

2020 年,天壤 TRTraffic 以数据为中心,从城市全局视角,为江西省南昌市提供了系统级解决方案,助力南昌成为全国首例从交通“限行”到“不限行”的城市。

直到最近,天壤 XLab 研发的 TRFold 以更具优势的方式解决了蛋白质结构预测难题。对此,天壤创始人兼 CEO 薛贵荣博士表示,“我们长期以来面临的都是此类复杂系统的决策问题,在技术层面已经有了深厚的积累,蛋白质折叠预测虽然是生命科学课题,但也属于这个领域,并且体现了 AI 在基础科学中的巨大潜能,这也是每个立志于基础研发的 AI 从业者的价值取向,我们不愿意错过这道风景。”

百舸争流

AI 将成为生命科学的“基础设施”

2016 年,AlphaGo 击败世界围棋冠军李世石,AI 进入公众视野,大量资金和人才涌入。此后数年,AI 进入千行百业,引发广泛关注的同时也深受质疑,普遍认为 AI 公司的商业化路径仍然模糊不清。

2020 年,AlphaFold2 破题蛋白质结构预测,再次引起高度关注。相比 AlphaGo 通过围棋突显了AI 的技术水平,AlphaFold 2 则加强了从业者对 AI 在生命科学大规模商用化的信心,且被认为是有史以来 AI 在基础科学研究上最大的突破。同时,AlphaFold2 对算法领域也有较大贡献,其中的模型间相互增强、数据蒸馏以及空间结构建模等算法极具应用价值。

在新冠肺炎疫情大流行的背景下,全球生命科学领域正面临转型,AI 生命科学元年正在开启。相信在未来几年内,将有大量的机构和公司加入到技术创新与生命科学研究的热潮中。

未来,天壤 XLab 将专注研究多蛋白的交互网络,挖掘蛋白质与蛋白质相互作用的精准链路,并以此构建全局视图研究生命科学,夯实技术创新能力。同时,天壤非常期待与各界同仁共同探索新的课题,携手推动 AI 生命科学高速发展。

0 人点赞