追赶AlphaFold2！天壤自研深度学习蛋白质折叠预测平台TRFold， 单张GPU实现秒级预测

近日，人工智能企业天壤对外宣布，其自研的深度学习蛋白质折叠预测平台TRFold在基于CASP14（2020年第14届国际蛋白质结构预测竞赛）蛋白质测试集的企业内测中，获得82.7/100的成绩（TM-Score），已经超过来自华盛顿大学的生物学家David Baker团队研发的RoseTTAFold 81.3/100的成绩，仅次于AlphaFold2的91.1/100的成绩。在400个氨基酸的蛋白链预测时，TRFold仅耗时16秒。

这是目前国内所有公开蛋白质结构预测模型中取得的最好成绩，它标志着我国计算生物学领域的表现已经处于世界第一梯队。

天壤也参加了CASP14竞赛，成绩处于中上，但是对比AlphaFold2的128个TPUv3核心（约等于256张GPU）算力配比，TRFold仅使用了8张Nvidia RTX 3090 GPU，以极小算力取得接近AlphaFold2的成绩。

在新型冠状病毒疫情爆发的初期，CASP组织者还发起了一次针对难以实验解析的COVID-19蛋白质结构预测，天壤提交的nsp6-D2预测结果被CASP官方选为六个“最具可信度模型”之一。

CASP14之后，始终同台竞争的David Baker团队则开发出了RoseTTAFold，一种“三轨（three-track）”神经网络，能够在十几分钟之内解析给定序列的三维结构，虽然无法超越AlphaFold2的优异表现，但也获得了不俗的成绩。

天壤成立于2016年，其XLab团队在两年半前开始进入计算生物学领域，当时Google旗下DeepMind公司研发的AlphaFold一代在CASP13中初出茅庐，首次展现了深度学习的方法在蛋白质结构预测领域的巨大想象力。在意识到它可能带来划时代的意义后，天壤立即投身其中。

蛋白质是构成生命的基本元件，破解它们的功能是揭开各种生命现象的金钥匙。然而再复杂的蛋白质也仅由20多种氨基酸以不同长度排列组合组成，在几微秒至几毫秒内迅速地折叠成一个特殊的三维结构，不同的结构决定了蛋白质之间迥异的功能。

在生命科学领域，蛋白质结构是个令人着迷的话题，引发了众多科学家的侧目攻坚，却也一直面临着难度大、成本高、进展缓慢的局面，直到迎来了人工智能的加入。

历时两年半，天壤TRFold性能已跃全球第二

今年7月，DeepMind公开了 AlphaFold2 的源代码，并在顶级科技杂志Nature上发表论文阐述了AlphaFold2的技术细节。同一天David Baker也公开了RoseTTAFold 的算法，并将研究成果经Science刊出。

这次开源在生物学界掀起了巨大的波澜，意味着生物学家有机会摆脱先进设备的掣肘，而这些设备往往都非常昂贵，只有经费充足的大学或研究机构才有条件配置。此后，那些小型团队或者个人研究者也能参与到蛋白质的研究中来。

对于像天壤这样一直投入在该领域的团队来说，开源带来便利的同时，也带来了挑战。大众的普遍认知里，每个人只要能登录github，就能下载AlphaFold2的源代码，并根据氨基酸序列预测蛋白质的三维结构。它已经是目前世界上效果最优的模型，同等类型的研发团队即便再努力，也很难超越这个结果。

“开源对我们自己要实现的技术路径带来了挑战和机遇，因为假如没有深入探索过程，就无法对不同的方法优势进行量化，也不能将最有价值的创新进行更大范围的应用。此外这次开源并未公布训练代码，意味着即便下载了它的源代码，也只能对单蛋白结构进行预测而无法根据项目需求进行进一步的优化改进。比起单个蛋白的预测，天壤更关注的是这项技术的落地，所以我们一定需要从头搭建一套自研算法，才能展开后续的工作。”天壤XLab团队负责人苗洪江博士表示。

4个月后，天壤蛋白质预测模型的迭代版在CASP14的竞赛数据集中取得了82.7/100的成绩，已经超过RoseTTAFold的81.3/100的成绩（该分数基于RoseTTAFold在github公布的数据集计算得出），是目前AlphaFold2之外所有公开蛋白质结构预测结果中最好成绩。

对比AlphaFold2，它耗费更小的计算资源，仅基于8张GPU，目前算法还在迭代优化，团队有信心最终能够以比AlphaFold2更小的计算资源取得90以上的成绩。

全方位创新，开启AI生命科学应用之路

天壤XLab团队表示，像蛋白质折叠预测这种全方位的创新项目非常珍贵，它覆盖了交叉学科的创新、行业的创新、基础科学的创新以及AI算法和工程能力的创新。为了获得更佳的性能，TRFold也做了许多创新和优化。

1) 在内存优化方面，TRFold模型参数量接近5千万，约为 AlphaFold2 的一半。TRFold在注意力权重分配时，在多个模块中采取权重共享的方式以节约计算资源。对于绝大多数蛋白质链（氨基酸残基数<400)，使用单张Nvidia RTX3090 GPU花费的预测时间不超过16秒。

2) 在训练数据方面，TRFold目前的算法模型用了比AlphaFold2更少的数据，也能取得相对优秀的表现，仅基于少量的真实数据训练而成，这得益于更充分的数据准备工作。

团队从多个宏蛋白质组序列库中寻找出蕴含更精准共进化信息的多序列排列，使模型在训练过程中能够获得对真实共进化信息更好的识别能力，从而取得对氨基酸残基距离和坐标更准确的预测结果。

目前的训练仅使用了来自蛋白质结构数据库PDB中的已知结构，团队期待在不久的将来加入了蒸馏数据增强工作之后，能够继续提高模型的预测和泛化能力。

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

当需要快速生成蛋白质结构时，可以快速使用end2end版本；当需要大规模测算蛋白质与蛋白质的相互关系时，分段式版本会有更优异的表现。而且分段式版本对后续研究基因突变的工作有指导意义。在这种情况下，共进化信息往往失去作用，无法按照原来的算法模型预测蛋白质结构。

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

面对团队取得的阶段性进展，天壤创始人薛贵荣博士表示：“我们通过两年半的时间掌握了蛋白质折叠核心技术，并以此参与到国际最前沿的生物医学领域的应用中，希望能够开辟出一条AI+生命科学的应用之路。”

正如不久前，DeepMind创始人Demis Hassabis宣布成立的新公司Isomorphic Labs（同构实验室），将聚焦AI对生命的基本机制的建模和理解，重构药物发现的过程。AI在生物计算领域的出色表现或将带来一个前所未有的AI生物研究新时代。

天壤XLab团队认为，目前的单个蛋白质折叠预测只是一个起点，更加精准的侧链优化、蛋白质的动态分析、蛋白质与其配体（如小分子、DNA、RNA、多肽、蛋白质等）的相互作用等一些列的问题还没有解决，接下来的工作重点将是利用目前的全蛋白质组协同进化分析，建立起蛋白质与蛋白质之间的相互作用的精准链路。

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

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