AI工程的极致:AlphaFold

图片

蛋白质是细胞的生命基础引擎,蛋白质的功能取决于它们的三级结构(protein tertiary structure)。当一条新生的氨基酸链从核糖体脱落,几秒内就会折叠成特定的三级结构(立体结构)。这种结构是由相邻的氨基酸中原子的相互作用决定的,直到蛋白质达到一种稳定的状态,从物理化学的角度看,这个结构的熵值应该是接近最低的。掌握蛋白质结构( protein folding problem)至关重要,例如,对于新药研发,从蛋白质的形状入手就能有效研发出药物,使药物影响特定蛋白质的功能来达到治疗的目的。

为了破解蛋白质的三级结构,生物学家一般使用X射线衍射、核磁共振光谱和冷冻电子显微镜。但这些技术非常复杂并且需要非常长的时间,例如,为了确定血红蛋白的三级结构,Max Ferdinand Perutz和他的团队花费了23年时间,确定血红蛋白结构的成就大到让Perutz拿到了诺贝尔化学奖。一些科学家毕生的工作就是为了确定20000个人类蛋白中的一个的结构。

在1961年,生物化学家Christian Anfinsen认为有更简单的办法,Anfinsen发现蛋白质的折叠并不是毫无章法,原理上任何一段线性氨基酸链的三级结构应该是可预测的。Anfinsen在他1972年的诺贝尔化学奖获奖感言中,提出了这样一个愿景:有一天,仅仅根据其氨基酸序列就可以预测任何蛋白质的3D结构。

到了1970年,科学家发明了核酸测序技术,在人类基因组计划的刺激下,自动快速测序技术被研发出来,并且揭开了完整的人类、动物、植物的基因组(记录蛋白质)的编码基因。从1970年代开始,化学家开始开发软件,在计算机上模拟蛋白质的结构,这项多尺度复杂化学系统建模工作在2013年获得了诺贝尔化学奖。

从1994开始,生物界开始举办代号CASP(Critical Assessment of Protein Structure Prediction)的蛋白质结构预测大赛。参赛者会被分到大约100个未知的蛋白的氨基酸序列,这些蛋白质的三级结构已经被确定但从未公布过。参赛团队将有几个月的时间去研发数学模型以预测这些未知的结构。CASP的计分体系为0到100分;超过90分意味着结构预测接近完美,将赢得比赛。

图片

在前面的12次比赛中,预测蛋白质的模型有些许改善,但除了最小和最简单的蛋白质,模型的精确度非常低,得分都没有超过40分。前面12次比赛的软件都是基于物理化学规则实现的,20多年来参赛者的模型都一样地差。

在2018年,DeepMind参加了CASP13比赛,软件取名为AlphaFold。这款软件融合了化学家、物理学家和生物学家掌握的蛋白质的物理化学知识,并且开始使用深度学习。CASP13中,AlphaFold进展明显,获得了60多分,但仍然不够好,其实现的精度仍然没有应用的价值。

到2020年,DeepMind扩展了AlphaFold使用AI深度学习的方式,AlphaFold2参加了2020年底的CASP14大赛,AlphaFold2的预测模型的得分的中位分数达到92.5,远高于其他参赛选手。2021年7月,DeepMind公开了AlphaFold2模型论文《Highly accurate protein structure prediction with AlphaFold》以及详细解释模型的50页附件。

虽然AlphaFold2预测精确,人们关心在CASP14之外,AlphaFold2可以揭晓多少当前未知的蛋白质三级结构。大量的蛋白质线性氨基酸序列都是已知的,但人们只知晓很少这些蛋白质的三级结构。据调查,截止2021年1月,经过数千名科学家工作仅仅确定了20000个人类蛋白质中大约30%蛋白质的结构和280百万非人类蛋白质中仅0.01%的蛋白质结构。

AlphaFold2是否能够快速将已知的核苷酸和氨基酸序列转成更多蛋白质的三级结构?2021年7月22日,DeepMind公布了答案,在公司官网上DeepMind提供了全部20000人类蛋白质里98%的预测结构,同时也公布了在生物医学研究中使用的20种生物模型的365000个蛋白的预测结构,包括小鼠、果蝇、酵母和大肠杆菌,到2021年底DeepMind公布了地球上大约2.8亿个蛋白质中大约50%的预测结构,而与此对比的,在2021年7月1日前人类只知道0.01%。

再回看一下CASP的历史成绩,CASP13中基于深度学习的蛋白质预测模型大幅度超越了前面12届比赛中普遍使用的物理化学方法,全球生物信息研究者通过2018年CASP13所展示的算法表现对于如下结论应该都是非常清楚:基于深度学习的同源模型化方法是提高蛋白质预测性能的唯一方向,全部的问题只是在于到底该如何设计深度学习模型。

明确的题目、明确的方向、明确的方法,只是等着各路研究者给出答案,但结果却是,从2018年到2020年,其他所有AI预测算法都在原地打转,只有AlphaFold2的性能突飞猛进。在CASP14中,除了AlphaFold2之外最好的15个算法的平均误差都是3Å左右,所使用的深度学习模型大同小异,只有AlphaFold2的遥遥领先,达到了平均误差1Å的原子级别精度。

AI技术领域里,在模型体系方面不可能存在只有某个公司知道而其他公司都不知道的所谓“黑科技”,模型技术都是公开的知识,那么为什么只有AlphaFold2的性能如此之好?这里体现出来的就是DeepMind在AI工程能力上的极致水平。

第一,AlphaFold2的模型架构并不难理解,但工程落地难度非常大。AlphaFold2的模型架构=特征模块+Transformer Encoder+RNN Decoder,在特征模块里抽取特征,在Transformer Encoder里融合特征,在RNN Decoder里还原3D结构。Transformer是Google在2017年的论文里就提出的技术,应该说所有深度学习研究者都知道Transformer性能超越RNN与CNN,但是真正要在大型模型里把Transformer用起来就需要很强的工程能力了。更大的工程能力门槛在于对Transformer模型的优化设计能力Google提出的经典Transformer是“向量to向量”,蛋白质预测模型里的Transformer是“矩阵to矩阵”,此时Transformer就需要重新设计了(AlphaFold2将其新设计的Transformer命名为Evoformer),这个工程门槛就挡住了绝大多数其他研究团队了。

图片


第二,AlphaFold2的优化设计让人匪夷所思,AlphaFold团队的优化能力令人惊叹。AlphaFold2模型在总体架构之下实现非常多的优化,较大的模块级别优化有10项,AlphaFold2论文也公布了这10项优化对于整体性能贡献度的消融研究结果。众多研究者看到这些优化后,共同的疑问就是“AlphaFold2到底是怎么想到这些优化方法的?”,答案只能是——这些优化来自于AlphaFold团队长期工程化的试验、验证与积累。领导过大型工程项目的研究者会比较容易理解这一点,大型工程项目中的各个模块经过长期尝试与验证,会积累出一些模块性的优化成果,最后再经过整体性尝试与验证后,还会积累出一些整体性的优化成果,当最后把这些优化成果汇聚到一起的时候,这些一点一点积累起来的优化成果确实会显得“匪夷所思”,因为如果要事先设计,再优秀的设计者也无法事先就想到这些优化方向。

图片

第三,AlphaFold2展示了业界高水平的大型AI工程团队组织能力。AlphaFold2论文作者一共34人,其中并列第一作者18人,大约可以认为对AlphaFold2影响巨大的核心研究人员就高达18人。这个团队在2018年和2020年两次输出重大成果,大约可以认为这个团队在2020年底成功输出AlphaFold2之前存在了5年,看准方向后维持这样规模团队5年,这就是DeepMind超强的把握方向能力与超强的AI工程团队组织能力。DeepMind在连续拿出AlphaFoldAlphaZero这样高水准的AI工程后,研发管理领域目前最流行这样一句话“拿金牌的价值远远高于常常拿铜牌”,以AlphaFold2为例,AlphaFold2的每项优化都足以发论文出成果,但是,真正高水平的研发组织,不会把力量分散到各种散乱的方向上去“勉强拿铜牌”,必然会把力量集中到最重要的方向上去“全力拼金牌”,AlphaFold团队通过5年的努力拿下了已持续了50年的蛋白质折叠问题的“金牌”。

图片

AlphaFold2公开之后,很多人争相指出AlphaFold2的局限,比如AlphaFold2尚无法预测相邻的蛋白的互相影响,因为很多蛋白质单独是没有功能的,这些蛋白质要的不是单个蛋白质的形状而是蛋白质复合体的形状。但是,生物信息领域的真实情况是,在AlphaFold2公开之后,目前全球范围内能够把AlphaFold2论文复现出来的团队都寥寥无几,AlphaFold2自身固然还需要向更远大的目标迈进,但当前的AlphaFold2已经和绝大多数生物信息AI团队形成了代际的能力差距。

AlphaFold2这类系统的出现给AI工业界引入一种工程化门槛,在这个程化门槛之下的AI模型,只要论文公开,大家就能去应用;在这个程化门槛之上的AI系统,即使论文公开代码开源,绝大多数团队连去复现源版系统性能的AI工程化能力都没有。