迄今为止规模最大的蛋白质语言模型问世了!
一年前,DeepMind开源AlphaFold2连登Nature、Science,刷爆生物和AI学界。
一年后,Meta带着速度快一个数量级的ESMFold来了。
不光速度快,模型还足足有150亿个参数。
LeCun发推称赞,这是Meta-FAIR蛋白质团队的伟大新成果。
共同一作Zeming Lin透露,30亿参数的大模型在256个GPU上训练了3个星期,而ESMfold在128个GPU上用了10天。至于150亿参数的版本,目前还不清楚。
他还表示,代码随后肯定会开源,敬请关注!
又大又快!
今天,我们的主角是ESMFold,一个从蛋白质个体的序列,直接进行高准确度、端对端、原子层级结构预测的模型。
论文地址:https://www.biorxiv.org/content/10.1101/2022.07.20.500902v1
150亿个参数带来的好处不必多说——通过训练,如今的大模型可以在原子大小的精度上预测蛋白质的三维结构。
从准确度上看,ESMFold和AlphaFold2、RoseTTAFold差不多。
但是,ESMFold推测速度要比AlphaFold2快一个数量级!
一下说数量级可能不好理解三者之间速度的对比,看看下面这张图就懂了。
有啥区别?
虽说AlphaFold2和RoseTTAFold在原子分辨率结构预测问题上取得了突破性的成功,但它们也依赖于使用多序列比对(MSA)和类似的蛋白质结构模板来实现最佳性能。
相比之下,通过利用语言模型的内部表征,ESMFold只用一个序列作为输入就能生成相应的结构预测,从而大大加快了结构预测的速度。
研究人员发现,ESMFold对低复杂度序列的预测与当下最先进的模型相当。
而且,结构预测的准确性与语言模型的复杂度息息相关,也就是说,当语言模型能更好地理解序列时,便可以更好地理解结构。
目前,有数十亿结构和功能未知的蛋白质序列,其中许多来自元基因组测序。
利用ESMFold,研究人员只需6个小时,就能折叠完成100万个元基因组序列的随机样本。
其中很大一部分具有高置信度,并且与任何已知的结构不同(在数据库中没有记录)。
研究人员认为,ESMFold可以帮助理解那些超出现有认知的蛋白质结构。
此外,由于ESMFold的预测速度比现有的模型快一个数量级,因此研究人员便可借助ESMFold来协助填补快速增长的蛋白质序列数据库与进展缓慢的蛋白质结构和功能数据库之间的鸿沟。
150亿参数的蛋白质语言模型
接下来我们就来具体说说Meta这款全新的ESMFold。
ESM-2是一个基于Transformer的语言模型,并使用注意力机制来学习输入序列中成对氨基酸之间的相互作用模式。
相对于上一代模型ESM-1b,Meta对模型结构、训练参数进行了改进,并增加了计算资源和数据。同时,相对位置嵌入的加入,使模型能够推广到任意长度的序列。
从结果来看,具有1.5亿个参数的ESM-2模型比具有6.5亿个参数的ESM-1b模型表现得更好。
此外,在结构预测的基准上,ESM-2也超过了其他的蛋白质语言模型。这种性能的提高与大型语言建模领域建立的规律是一致的。
随着ESM-2规模的增加,可以观察到语言建模的精度有很大的提高。
端到端的单序列结构预测
SMFold和AlphaFold2的一个关键区别是,ESMFold使用语言模型表示,消除了对明确的同源序列(以MSA的形式)作为输入的需要。
ESMFold通过用一个处理序列的Transformer模块取代处理MSA的计算昂贵的网络模块,简化了AlphaFold2中的Evoformer。这种简化意味着ESMFold的速度大大提高,远高于基于MSA的模型。
折叠主干的输出接下来又被一个结构模块处理,它负责输出最终的原子级结构和预测的置信度。
研究人员将ESMFold与AlphaFold2和RoseTTAFold在CAMEO(2022年4月至2022年6月)和CASP14(2020年5月)测试集上进行比较。
当只给单一序列输入时,ESMFold的表现要比Alphafold 2好得多。
而当使用完整的管道时,AlphaFold2在CAMEO和CASP14上分别达到了88.3和84.7。ESMFold在CAMEO上取得了与RoseTTAfold相当的准确率,其平均TM分数为82.0。
结论
研究人员发现,以无监督学习为目标的语言模型在一个大型的进化多样化的蛋白质序列数据库中训练,能够对蛋白质结构进行原子级的分辨率预测。
将语言模型的参数扩大到15B,就可以系统地研究规模对蛋白质结构学习的影响。
我们看到,蛋白质结构预测的非线性曲线是模型规模的函数,并且观察到了语言模型对序列的理解程度与结构预测之间的强烈联系。
ESM-2系列的模型是迄今为止训练的最大的蛋白质语言模型,其参数仅比最近开发的最大文本模型少一个数量级。
而且,ESM-2比以前的模型有非常大的改进,即使在150M的参数下,ESM-2也比ESM-1代语言模型在6.5亿的参数下捕捉到更准确的结构图。
研究人员表示,ESMFold性能的最大驱动力是语言模型。由于语言模型的迷惑性和结构预测的准确性之间有很强的联系,他们发现当ESM-2能较好地理解蛋白质序列时,就可以获得与目前最先进的模型相当的预测结果。
ESMFold获得了准确的原子分辨率结构预测,推理时间还比AlphaFold2快了一个数量级。
在实践中,速度的优势甚至还要更大。因为ESMFold不需要搜索和进化相关的序列来构建MSA。
虽说有更快的方法可以减少搜索时间,但再怎么减少还是可能会很长。
而推理时间的极大缩短带来的利好不言自明——速度的提高将使绘制大型元基因组学序列数据库的结构空间成为可能。
除了基于结构的工具来识别远端同源性和保护性之外,用ESMFold进行快速准确的结构预测,还能在大量新序列集合的结构和功能分析中发挥重要作用。
在有限的时间内获得数以百万计的预测结构,有利于发现对天然蛋白质的广度和多样性的新认识,并能发现全新的蛋白质结构和蛋白质功能。
作者介绍
本文的共同一作是来自Meta AI的Zeming Lin。
据个人主页介绍,Zeming在纽约大学攻读博士学位,并在Meta AI担任研究工程师(访问),主要负责后端基础设施的工作。
他本硕都就读于弗吉尼亚大学,在那里,他和Yanjun Qi大佬一同做有关机器学习应用的研究,尤其是在蛋白质结构预测方面。
感兴趣的领域为深度学习、结构预测,以及信息生物学。