ESM2和ESM fold：语言模型进行原子级结构预测

机器学习领域的最新进展利用了多序列比对中的进化信息来预测蛋白质结构。这里作者则展示了如何使用大语言模型从蛋白质的一级序列直接推断出完整的原子级蛋白质结构。随着蛋白质序列语言模型的参数扩展到150亿，蛋白质结构的原子分辨率图像在学习到的表征中逐渐显现。这使得高分辨率结构预测的速度提升了一个数量级，从而能够对蛋白质进行大规模的结构表征。作者运用这一能力构建了ESM宏基因组图谱，对超过6.17亿个宏基因组蛋白质序列的结构进行了预测，其中包括超过2.25亿个高置信度的预测结果，这让后续的研究得以一窥天然蛋白质的广泛多样性。

来自：Evolutionary-scale prediction of atomic-level protein structure with a language model, Science, 2023

背景概述

在进化尺度上，蛋白质序列蕴含着生物结构和功能的信息。蛋白质的生物学特性限制了其序列中通过进化选择的突变。因此，蛋白质的结构和功能可以从序列模式中推断出来。从经典方法到深度学习的引入，再到如今的高精度结构预测，这一认知一直是计算结构预测取得进展的核心。

语言模型有潜力学习蛋白质序列在进化过程中的模式。这一想法推动了对进化尺度语言模型的研究，在这类模型中，基础模型学习反映潜在生物学特征的表征，并且随着表征能力的增强，能够在低分辨率下捕捉蛋白质的二级结构和三级结构。 从香农的文本熵模型开始，复杂度不断提升的语言模型得以开发，最终发展出现代基于注意力机制的大规模架构。尽管文本语言模型的训练目标很简单，比如填充缺失单词或预测下一个单词，但它们展现出了涌现能力，这种能力会随着计算能力、数据量和参数数量的增加而发展。现代语言模型拥有数百亿甚至数千亿参数，并展现出诸如少样本语言翻译、常识推理和数学问题解决等能力，而且这些都无需明确的监督。

作者认为，在进化过程中填补蛋白质序列中缺失氨基酸的任务，需要语言模型理解产生这些序列模式的潜在结构。随着语言模型的表征能力以及其在训练中接触到的蛋白质序列多样性的增加，有关蛋白质序列生物学特性的深层次信息将会显现，因为正是这些特性产生了人们在序列中观察到的模式。为了研究这种涌现现象，作者将语言模型的参数规模从800万扩展到150亿。发现原子分辨率的结构在语言模型中逐渐显现，并且在参数规模跨越四个数量级的过程中持续改进。语言模型对蛋白质序列的理解程度（困惑度）与结构预测的准确性之间存在很强的相关性，这揭示了语言建模与结构学习之间的紧密联系。

作者证明了语言模型能够直接从序列实现快速的端到端原子分辨率结构预测。该方法利用语言模型捕获的进化模式来生成精确的原子级预测。这消除了当前最先进的结构预测流程中成本高昂的部分，既无需进行多序列比对（MSA），又极大地简化了用于推理的神经架构。这使得仅推理前向传递的速度就提高了多达60倍，同时还完全省去了相关蛋白质的搜索过程。使用AlphaFold和RoseTTAFold的高灵敏度流程时，搜索过程可能需要超过10分钟，即使是最近灵敏度较低的快速流程，搜索过程也是计算成本的重要组成部分。在实际应用中，这意味着与最先进的预测流程相比，速度提升了一到两个数量级。

这种速度优势使得将结构预测扩展到宏基因组规模的数据集成为可能。在过去十年中，人们致力于通过宏基因组采样，将对蛋白质序列的认知扩展到地球上微生物的巨大自然多样性领域。这些努力促使蛋白质序列数据库的规模呈指数级增长，如今数据库中已包含数十亿种蛋白质。最近，对人类蛋白质组中约2万种蛋白质以及UniProt中约2亿种已编目的蛋白质的计算结构表征工作已经完成，但宏基因组蛋白质的庞大数量对结构表征来说是一个更大的挑战。宏基因组结构的范围和多样性仍是未知的，这是生物学知识的前沿领域。

作者对宏基因组蛋白质进行了进化尺度的结构表征，几乎对MGnify90数据库中的所有序列（超过6.17亿种蛋白质）进行了折叠预测。作者在由2000个GPU组成的异构集群上，仅用2周就完成了这一表征工作，这表明该方法可扩展应用于更大规模的数据库。作者对超过2.25亿个结构进行了高置信度预测，揭示并描绘了与现有认知相距甚远的宏基因组空间区域。这些预测的结构可在ESM宏基因组图谱（https://esmatlas.com）这一开放科学资源中获取。

在基于蛋白质序列训练的语言模型中，原子分辨率的结构逐渐显现

作者训练了一系列基于Transformer的蛋白质语言模型ESM-2，其参数规模从800万到150亿不等。与上一代模型ESM-1b相比，ESM-2在架构（带RoPE的BERT）、训练参数方面进行了改进，并增加了计算资源和数据量。由此产生的ESM-2模型系列在参数数量相当的情况下，性能优于之前最先进的ESM-1b。并且在结构预测benchmark上优于其他蛋白质语言模型。

ESM-2 被训练用于预测从蛋白质序列中随机被屏蔽掉的氨基酸的类别：$${\mathcal{L}}_{MLM}=-\sum _{i\in M}\log p(x_i|x_{/M})$$其中，对于随机生成的掩码$M$，它涵盖了序列$x$中15%的位置。模型的任务是根据周围未被屏蔽的上下文$x_{/M}$（不包括被mask的位置），预测掩码中氨基酸$x_i$的类别。这种掩码语言建模目标促使模型学习氨基酸之间的依赖关系。尽管训练目标本身简单且无需监督，但要在数百万个在进化上具有多样性的蛋白质序列上完成这一任务，需要模型学习跨越进化历程的序列模式。作者预期这种训练会使生物结构在语言模型中具象化，因为生物结构与序列模式息息相关。ESM-2是在UniRef蛋白质序列数据库中的序列上进行训练的。在训练过程中，从约1.38亿条UniRef90序列中，以均匀权重对约4300万个UniRef50训练簇进行序列采样。因此，在整个训练过程中，该模型会接触到约6500万条序列。

随着将ESM-2的参数规模从800万增加到150亿（8M、35M、150M、650M、3B、15B），作者观察到其对蛋白质序列建模的保真度有了显著提高。这种保真度可以用困惑度来衡量，完美模型的困惑度为1，而随机预测的模型困惑度为20。直观地说，困惑度描述了模型在序列的每个位置上可供选择的平均氨基酸数量。从数学角度讲，困惑度被定义为序列负对数似然的指数。本节中所有模型均在训练27万步时进行比较。随着参数增加，直至最大参数规模的模型，其保真度持续提升。800万参数的模型困惑度为10.45，而150亿参数的模型困惑度达到6.37，这表明随着模型规模的扩大，其对蛋白质序列的理解有了很大的提升。

这种训练还使模型中出现了结构信息。由于ESM-2仅在序列上进行训练，所以模型所获得的任何有关结构的信息必然是对序列模式进行表征的结果。已知通过掩码语言建模训练的Transformer模型会形成与蛋白质残基-残基接触图相对应的注意力模式。作者研究了这种低分辨率的蛋白质结构信息是如何随模型规模的变化而出现的。

使用线性投影从语言模型的注意力模式中提取接触图。预测的前L个（L为蛋白质长度）接触的精度（长程接触精度）衡量了注意力模式与蛋白质结构的对应程度。ESM-2中形成了与三级结构相对应的注意力模式（图1A），并且扩大模型规模能显著提升对结构的理解（图1B）。预测接触的准确性会随着训练集中进化相关序列数量的不同而变化。训练集中具有更多相关序列的蛋白质，其学习轨迹相对于模型规模的增长更为陡峭（图1C）。因此，对于进化深度高的序列，在模型规模较小时其性能提升就会趋于饱和；而对于进化深度低的序列，随着模型规模的增大，性能会持续提升。

• 图1A：PDB 3LYW的预测接触概率（右下角）和实际接触精度（左上角）。对于长度为L的序列，如果一个接触在最可能的前L个接触范围内，则该接触被视为正预测。
• 图1B：根据在训练集中搜索时MMseqs命中次数对性能进行分类。更大规模的ESM-2模型表现更好。
• 图1C：具有不同MMseqs命中次数的序列，其性能随模型规模增加的提升轨迹。

对于单个蛋白质，经常观察到，随着模型规模的增大，接触预测的准确性呈现出非线性提升。绘制每次向更大规模模型转变时长程接触精度分布的变化情况，结果显示整体分布向更好的性能方向偏移（图1D）。接触图预测的准确性与困惑度相互关联，接触图准确性变化较大的蛋白质，其困惑度也会发生较大变化。这种关联表明，语言建模目标与注意力图中折叠结构信息的显现直接相关。

为了识别模型中的原子分辨率信息，作者使用等变Transformer从语言模型的内部表征中投影出每个原子的空间坐标。这一投影是利用蛋白质数据库（PDB）中通过实验确定的蛋白质结构来进行拟合的，并在194个CAMEO蛋白质和51个CASP14蛋白质上进行评估。TM分数的取值范围为0到1，用于衡量投影结果相对于真实结构的准确性，0.5这个数值对应的是正确预测蛋白质折叠的阈值。

• 图1D：从左至右展示了参数从800万到150亿的模型，通过无监督接触精度将较小的模型（x轴）与下一个更大的模型（y轴）进行比较。这些点代表蛋白质数据库（PDB）中的蛋白质，根据较小和较大模型之间序列困惑度的变化进行着色。接触预测性能变化较大的序列，在以困惑度衡量的语言模型理解方面也表现出较大变化。
• 图1E：在CASP14和CAMEO联合测试集上的TM分数。预测是通过在语言模型之上仅使用结构模块头进行的。各点根据模型之间困惑度的变化进行着色。

这种投影的准确性会随着语言模型规模的增大而提高。参数规模达150亿的模型在CAMEO测试集上的TM分数达到0.72，在CASP14测试集上达到0.55，相较于参数规模为1.5亿的ESM-2模型，分别提高了14%和17%（图1E）。每次模型规模增大时，都有一部分蛋白质的预测准确性会发生显著变化。例如，当模型参数规模从3500万增加到1.5亿时，蛋白质7QQA的均方根偏差（RMSD）从7.0 Å降至3.2 Å；当参数规模从30亿增加到150亿时，CASP目标蛋白T1056的均方根偏差从4.0 Å降至2.6 Å（图1F）。这些发现表明，语言建模能力的提升与低分辨率（接触图）和高分辨率（原子水平）结构信息的增加相关联。

• 图1F：在所有ESM-2模型规模下，对CAMEO结构7QQA和CASP target 1056进行结构预测，预测结果按pLDDT值上色（粉色表示值低；青色表示值高）。对于7QQA，在模型达到1.5亿参数规模时，预测精度有所提升。对于T1056，在模型达到150亿参数规模时，预测精度有所提升。（灰色表示真实结构）

利用语言模型加速高精度原子分辨率结构预测

语言模型极大地加快了当前最先进的高分辨率结构预测进程。语言模型在序列中学习到了与结构相关的进化模式，从而无需外部进化数据库、多序列比对（MSA）以及模板。作者发现，ESM-2语言模型能够直接从蛋白质一级序列生成最先进的三维（3D）结构预测结果，这使得结构预测速度提升一个数量级，同时保持高分辨率精度。

通过为ESM-2训练一个折叠头，开发出了ESMFold，这是一款完全端到端的单序列结构预测工具（图2A）。在进行预测时，将蛋白质序列输入到ESM-2中。该序列经过语言模型的前馈层处理，表征被传递至折叠头。折叠头起始于一系列折叠模块。每个折叠模块在更新Seq Rep和Pair Rep之间交替进行。这些模块的输出被传递至一个等变Transformer结构模块，在输出最终的原子级结构和预测置信度之前，会执行三步循环操作。与当前最先进的结构预测模型相比，这种架构大大简化。

• 图2A：ESM-Fold模型架构。箭头展示了网络中的信息流，从语言模型流向折叠主干，再到输出三维坐标和置信度的结构模块。LM即语言模型。ESM-fold包含两个结构，一个是折叠模块（folding module），输入语言模型ESM-2（3B参数量版本）生成的representation；另一个是结构模块(structure module)，一个等变Transformer输入上述模块输出，并输出3D原子坐标。
• 图2B：ESMFold 可产生准确的原子分辨率预测。散点图将 ESMFold（x 轴）的预测结果与 AlphaFold2（y 轴）进行比较，并按语言模型的困惑度着色。困惑度低的蛋白质，其预测结果与 AlphaFold2 相似。
• 图2C：在 CAMEO 数据集上，模型预测的 pLDDT 与真实 LDDT 的对比（左）以及与 AlphaFold 相比的相对性能（右）。pLDDT 是对预测准确性的一种校准良好的估计。
• 图2D：成功案例：上图展示了 T1057 在测试集上的预测结果，ESMFold（左）和 AlphaFold2（右）。颜色表示两个模型预测的 LDDT（ESMFold 高置信度为青色；AlphaFold2 高置信度为绿色；两者低置信度为粉色）。真实结构以灰色显示。下面两张图展示了对二聚体（PDB 编号：7LQM）和四聚体（PDB 编号：7QYM）的复合物预测；ESMFold 的预测结果按链 ID 着色，并叠加在真实结构（灰色）上。报告了相互作用的 DockQ分数；对于四聚体 7QYM，该分数是相互作用链对分数的平均值。
• 图2E：失败案例：T1074 在测试集上的预测结果，ESMFold（左）和 AlphaFold2（右）。颜色表示两个模型预测的 LDDT（ESMFold 高置信度为青色；AlphaFold2 高置信度为绿色；两者低置信度为粉色）。真实结构以灰色显示。ESMFold 的 TM 分数明显低于 AlphaFold2 的 TM 分数。该失败序列的困惑度为 16.6，这意味着语言模型无法理解输入序列。

ESMFold显著提高了预测速度。在单块英伟达V100 GPU上，ESMFold对一个包含384个残基的蛋白质进行预测仅需14.2秒，比单个AlphaFold2模型快6倍。对于更短的序列，速度提升可达约60倍。在约25000个簇上训练折叠头，这些簇涵盖了来自蛋白质数据库（PDB）的总共约325000个经实验测定的结构，此外还使用了一个由作者用AlphaFold2预测得到的约1200万个结构组成的数据集进行增强。该模型采用与AlphaFold相同的损失函数进行训练。

作者将ESMFold在评估数据集上的结果与AlphaFold2和RoseTTAFold进行比较（图2B）。ESMFold在CAMEO数据集上的平均TM分数达到0.83，在CASP14数据集上为0.68。当通过去除多序列比对（MSA）来评估AlphaFold2和RoseTTAFold对单序列的预测性能时，它们的表现大幅下降，远低于ESMFold。这是一种人为设定的情况，因为AlphaFold2并未针对单序列进行明确训练；然而，单序列场景最近在蛋白质设计中变得很重要。

图2B中，困惑度与结构预测之间的关系表明，语言模型性能的提升将转化为单序列结构预测准确性的提高，这与scaling分析的观察结果一致（图1D和图1E）。这意味着语言模型对某个序列的困惑度可用于预测ESMFold结构预测的质量。

宏基因组的进化尺度结构表征

这种快速且高分辨率的结构预测能力使得对宏基因组蛋白质进行大规模结构表征成为可能。作者对来自MGnify90数据库的超过6.17亿条序列进行了结构折叠预测。这些序列涵盖了长度在20到1024之间的所有序列，占MGnify90数据库中全部序列的99%。总体而言，此次表征产生了约3.65亿个置信度良好的预测结果（平均预测局部距离差异测试值（pLDDT）> 0.5且预测TM分数（pTM）> 0.5），这相当于该数据库中约59%的序列；同时还产生了约2.25亿个高置信度的预测结果（平均pLDDT > 0.7且pTM > 0.7），这约占所折叠的全部结构的36%（图3）。

• 图3A：针对宏基因组序列，ESMFold与AlphaFold2的校准对比。x轴表示平均预测局部距离差异测试值（pLDDT），y轴表示与相应AlphaFold2预测结果的局部距离差异测试值（LDDT）。图中展示的分布是对来自MGnify数据库约4000条序列子样本的密度估计。
• 图3B：对来自MGnify数据库的约6.17亿个ESMFold预测结构，每个结构所计算出的平均pLDDT值的分布。
• 图3C：从100万个随机抽取的高置信度（平均pLDDT > 0.7且预测TM分数（pTM）> 0.7）结构中，每个结构与最相似的蛋白质数据库（PDB）结构的TM分数分布。
• 图3D：使用UMAP算法将100万个高置信度蛋白质结构样本在二维空间中可视化，并根据与最近的PDB结构的距离进行着色，其中与已知结构相似度低的区域用深蓝色表示。图中给出了示例蛋白质结构及其在序列图谱中的位置。
• 图3E：另一张UMAP图，其中100万个序列根据与图D相同的坐标绘制，但根据blastp搜索结果，按照与UniRef90中最相似条目的序列一致性进行着色。

许多宏基因组结构预测结果具有较高的置信度（图3B），并且在现有的结构数据库中没有对应的代表结构（图3C - E）。在随机抽取的100万个高置信度结构样本中，76.8%（767,580个）的蛋白质与UniRef90数据库中的任何序列的序列一致性都低于90%，这表明这些蛋白质与现有的UniRef90序列不同（图3E）。其中，有3.4%（33,521个蛋白质）在UniRef90中根本找不到匹配序列。我们使用Foldseek将预测结构与蛋白质数据库（PDB）中的已知结构进行比较。当TM分数阈值分别设定为0.7和0.5时，Foldseek报告分别有25.4%（253,905个蛋白质）和12.6%（125,765个蛋白质）的预测结构找不到匹配项（图3C和图3D）。有2.6%（25,664个）的预测结构既与已知结构的结构相似度较低（TM分数 ≤ 0.5），又没有相近的序列同源物（序列一致性 > 30%）。根据这些子样本的估计结果，大约有2800万个蛋白质（在2.25亿个高置信度预测结果中占12.6%）的预测结果置信度高，且与已知蛋白质结构的TM分数 < 0.5。这些结果表明，ESMFold能够有效地对与现有认知差异较大的蛋白质空间区域进行表征。

对于从头蛋白设计而言，大多的蛋白都是自然界中没有存在的，更没有MSA，因此ESMfold将比AF2更适用于此类不存在足够MSA的蛋白（从头设计蛋白和孤儿蛋白）的结构预测。

• 图4：宏基因组序列的ESMFold结构预测示例。
• A) 六个不同宏基因组序列的预测结构示例。每个子图左侧：预测结构与AlphaFold2的预测结果（淡绿色）一同展示。每个子图右侧：预测结构与根据TM分数由Foldseek确定的最相近的蛋白质数据库（PDB）结构一同展示。
• B) 和 C) 两个ESMFold预测结构的示例，它们与PDB中的实验结构吻合度良好，但与UniRef90中的任何序列的序列一致性较低。B) MGYP000936678158的预测结构与一种细菌核酸酶的实验结构（浅棕色，PDB编号：3H4R）对齐，而 C) MGYP004000959047的预测结构与一种细菌固醇结合域的实验结构（浅棕色，PDB编号：6BYM）对齐。

图4B和图4C主要说明了ESMFold在识别序列相似性缺失情况下的结构相似性方面的能力，以及这种能力对于挖掘蛋白质功能信息的重要意义，具体如下：
揭示远缘同源关系
传统基于序列相似性的方法（如在UniRef90数据库中搜索或使用jackhmmer参考蛋白质组搜索）可能无法检测到某些蛋白质之间的关系，但ESMFold能够通过预测结构发现远缘同源性。例如，MGYP000936678158和MGYP004000959047这两个序列在UniRef90中找不到匹配项，但ESMFold预测的结构却与PDB中的实验结构存在显著相似性，表明这些蛋白质虽然在序列上差异较大，但在结构层面存在进化上的关联。
助力功能预测
蛋白质的结构与其功能密切相关。图4B和图4C中的例子表明，通过ESMFold发现的结构相似性可以为蛋白质功能的预测提供线索。如MGYP000936678158的预测结构与细菌核酸酶的结构相似，那么它可能也具有类似核酸酶的功能；MGYP004000959047的预测结构与细菌固醇结合域相似，提示其可能参与脂质结合相关的生物学过程。这使得研究人员能够在序列信息有限的情况下，借助结构信息深入了解蛋白质的潜在功能。
拓展蛋白质结构和功能认知
这些结果显示了ESMFold在挖掘新的蛋白质结构和功能方面的潜力。它可以发现那些被传统序列分析方法遗漏的蛋白质关系，有助于填补我们对蛋白质结构和功能多样性认识的空白。