密苏里大学许东团队提出S-PLM模型，通过序列和结构对比学习的语言模型进行蛋白质预测

蛋白质在各种生化过程中起着至关重要的作用。大型蛋白质语言模型（PLM）通过加速蛋白质功能的确定和具有所需功能的蛋白质的设计，在重塑蛋白质研究方面具有良好的潜力。PLM的预测和设计能力依赖于从蛋白质序列中获得的表示。然而，大多数PLM缺乏关键的3D结构信息，这限制了PLM在各种应用中的预测能力，特别是在高度依赖3D结构的应用中。

2024年12月12日，美国密苏里大学许东教授团队在Advanced Science上发表文章S-PLM: Structure-Aware Protein Language Model via Contrastive Learning Between Sequence and Structure。

作者提出了S-PLM，这是一种3D结构感知的蛋白质语言模型（PLM），它利用多视图对比学习来对齐协调潜在空间中蛋白质的序列和3D结构。S-PLM利用Swin-Transformer在AlphaFold预测的蛋白质结构上嵌入结构信息，并将其融合到蛋白质语言模型ESM2的序列嵌入中。实验证明了S-PLM的性能超越了现有方法。

如图1所示，S-PLM预训练架构由两个编码器组成（图1a），一个编码蛋白质序列，另一个编码三维蛋白质结构。本研究采用单字母氨基酸序列作为蛋白质序列的输入。由于残基间距离包含了蛋白质结构的全面和必要的信息，因此使用主链Cα接触图来表示蛋白质的三维结构。在预训练过程中，S-PLM同时输入氨基酸序列和主链Cα接触图。通过序列编码器将蛋白质序列信息转换为残基级嵌入，通过结构编码器（Swin-Transformer）将接触图信息转换为蛋白质级嵌入。然后，将序列和接触图转换为单独的蛋白质水平嵌入。最后，使用对比学习训练S-PLM模型，以最小化同一批序列和接触图的对比损失。S-PLM模型的目标是最大化同一蛋白质序列和结构的嵌入对齐，并明确分离不同蛋白质之间的序列和结构的嵌入去对齐。受SimCLR方法的启发，本工作采用了CLIP方法进行对比语言图像预训练。除了CLIP跨不同模态的对齐和去对齐之外，模型还考虑了同一模态内的去对齐。例如，如图1a所示，模型还强调不同蛋白质的序列（s1与s2）和接触图（c1与c2）的嵌入之间的不相似性。

图1 S-PLM结构图

S-PLM的序列编码器通过将结构感知模块集成到ESM2模型中来实现（图1b）。有几种方法可以实现这个结构感知模块。一种选择是通过调整预训练的权值来微调ESM2模型，这在集成结构信息的同时保留了ESM2编码器的原始架构。或者，可以使用适配器调优，其中适配器模块集成到ESM2模型的top-K Transformer层中。这些适配器模块作为结构感知组件，仅在预训练过程中进行训练。如图1b所示，结构感知模块由瓶颈结构和跳过连接组成，在一个Transformer层中定位两次：在多个头部注意力投影之后和在两个前馈层之后。使用适配器调优来实现S-PLM的序列编码器具有几个优点。首先，集成的适配器模块非常紧凑。它比ESM2原有的Transformer模块包含更少的参数，减轻了培训负担。其次，它允许连续训练以添加新的蛋白质特征（例如蛋白质功能），以用于模型的未来扩展，而不会忘记先前学习的特征，因为S-PLM预训练保留了ESM2模型的序列表示能力，而ESM2主干模型（架构及其权重）完好无损。

在推理阶段，S-PLM可以灵活地接受序列或接触图作为输入，并根据特定的下游任务产生相应的各级嵌入。这种多功能性使S-PLM能够根据给定任务的特定输入数据和要求进行调整并提供合适的表示。在本文和后续研究中，S-PLM模型主要从蛋白质序列中生成序列嵌入。因此，预训练的S-PLM序列编码器在部署之前生成残差级嵌入，用于下游任务。整个序列编码器可以完全冻结或可学习。为了充分利用S-PLM在监督蛋白预测任务中的潜力，作者基于S-PLM的序列编码器开发了几种轻量级调优策略，所有这些策略都被纳入轻量级调优工具箱，包括微调顶层、适配器调优（图1c）和LoRA调优（图1d）。

模型的序列编码器是基于预训练的ESM2模型开发的。考虑到计算资源和模型容量的限制，作者选择ESM2-t33-650M-UR50D作为基础PLM模型，该模型有6.5亿个参数。首先对每个氨基酸进行独热编码，然后使用33层Transformer编码器对输入蛋白序列进行标记。每个位置的嵌入维数为1280。在此过程中，将BEGIN令牌和END令牌添加到序列中，并与氨基酸令牌一起通过Transformer，并且使用PAD令牌填充序列。通过Transformer层，输出是每个残基1280维向量的张量，嵌入BEGIN和END标记以及填充序列的嵌入。每个残基的嵌入用于下游任务中蛋白质的残基水平表示。每个残基嵌入的平均值（不包括填充标记）用于对比学习训练和下游任务的蛋白质水平嵌入。然后，将两个投影层应用于蛋白级嵌入，将该维度转化为最终输出的蛋白级嵌入，输出维数为256。

由于接触图包含完整的蛋白质结构信息，具有固有的不变性，且易于实现，因此采用接触图来表示三维蛋白质结构。因此，作者的结构编码器是专门设计来编码蛋白质接触图的，应用Swin-Transformer (swinv2-tinty-patch4-window8-256)，因为它能够更有效地从接触图表示中提取特征。为了满足图像网络需要三个输入通道的要求，将接触映射转换为具有三个通道的表示。通过计算一个序列中每个氨基酸的Cα原子之间的坐标距离，生成原始接触图。通常，如果两两距离在选定的阈值内，则接触映射值为1，表示残基之间的接触；否则，它的赋值为0。对每个通道应用距离阈值，并将原始接触映射转换为连续相似矩阵，每个通道的距离阈值为22埃，与AlphaFold2中使用的值相同。通过对Swin-Transformer表示层的嵌入进行平均，得到结构的蛋白级嵌入，用于对比学习训练。然后，对蛋白级嵌入应用两个投影层，将其维数变换为最终输出的蛋白级嵌入，其维数为256维，与序列中最终输出的蛋白级嵌入相同。

本研究中对比学习的目的是使来自同一蛋白质的序列嵌入和结构嵌入更加接近，并进一步排斥来自不同蛋白质的潜伏空间嵌入。为了实现这一点，对序列和结构编码器的最后一个投影层获得的蛋白质水平嵌入应用了多视图对比损失函数SimCLR，其中的NT-Xent对多视图对比损失函数进行了改进，该损失函数有助于最大化同一个蛋白质两个视图中蛋白质嵌入的对齐，并尽量减少蛋白质嵌入与其他蛋白质嵌入的相似性。

训练数据库从Swiss-Prot文库中获取蛋白质的氨基酸序列，并保存为FASTA格式。蛋白质的三维结构从AlphaFold2库中获得。基于AlphaFold2预测的3D结构，使用内部Python代码确定了单个蛋白质的Cα-Cα接触图。从Swiss-Prot文库中随机选择50万个蛋白进行训练，选择41500个蛋白进行验证。验证集中与训练集中相似的序列没有被删除。考虑到Swiss-Prot文库的庞大规模（总共有542378个蛋白质），这种相似性可以忽略不计。在预训练过程中，截断大于512个残基的输入蛋白，只对序列和接触图谱的前512个残基进行再训练。

本文中的轻量级调优策略是指对预先存在的模型进行特定且通常是参数有效的修改的训练方法，与完全微调训练模型相比，减少了所需的计算资源和内存。在S-PLM的序列编码器上实现微调顶层、LoRA和适配器调优，如图1所示，用于下游蛋白质序列预测任务。每项策略的详情如下：

• 轻量调整序列编码器微调顶层：ESM2主干模型共有33层Transformer。这里的“顶层微调”是指只微调K≤33层的变压器顶层，其余的全部冻结。这里，K是构型中的一个超参数。
• 序列编码器-LoRA的轻量级调优：LoRA指的是低秩自适应，它冻结了预训练的模型权重，并将可训练的秩分解矩阵注入Transformer体系结构的每一层。它可以显著减少下游任务的可训练参数的数量。
• 适配器调优：将适配器模块集成到ESM2模型的Transformer层中。适配器模块在ESM2的一个Transformer层中定位了两次：在自注意力投影之后和在两次前馈之后。每个适配器模块由一个瓶颈结构和一个跳跃连接组成。瓶颈结构将输入数据压缩到降维空间，然后对数据进行重构，使其恢复到原始输入维。瓶颈结构使适配器模块相对于原始Transformer中的注意层和前馈层具有很少的参数。ESM2模型与结构感知模块的集成如图1b所示。与最初的适配器调优不同，作者将适配器模块应用到所有Transformer层中，这些适配器模块专门插入到ESM2的top-K Transformer层中。

对于受监督的下游任务，适配器调优是通过将一组额外的并行适配器集成到结构感知模块中来实现的，每个适配器专用于一个特定的任务，如图1c所示。这些适配器模块中的每个模块都与结构感知模块中的模块共享相同的体系结构，并捕获针对各种下游任务定制的输入特性的不同方面。这些模块独立处理相同的输入特征，提取独特的表示，然后组合。为了训练新的下游任务，将一个新的可训练适配器模块添加到并行适配器列表中，同时冻结所有先前添加的并行适配器。

作者将S-PLM与一些具有代表性的方法进行了比较，如表1所示，在蛋白质折叠分类、蛋白质超家族分类、蛋白质家族分类任务上，无论是否微调模型，无论top K取值多少，S-PLM准确度都高于对应的基线语言模型ESM2，而参数量却相差不大，因此效率也相差不大。结果表明，S-PLM超越了现有方法。

表1 与其他方法对比

作者设计了消融实验来验证模型设计的有效性。如表2所示，作者使用蛋白质功能预测中的三个数据集（生物过程BP，细胞组件CC，分子功能MF）对比了不同微调策略的影响，说明了作者设计的适配器微调的有效性。S-PLM在经过适配器微调和LoRA微调后准确度均高于使用ESM2。与适配器策略相比，LoRA策略的准确度虽然稍有降低，但其参数量显著减小，适用于轻量化部署的场景。以上结果说明了S-PLM的优越性。

表2 消融实验

作者还进行了案例分析。为了研究S-PLM是否可以将结构信息注入序列潜在空间，作者评估了CATH蛋白结构域的序列表示。由于本实验优先考虑结构信息，作者从每个CATH超家族中选择一个具有代表性的序列，以提供清晰的可视化。作者利用了CATHS40数据集，其蛋白质具有最大40%的序列同一性。作者的分析主要集中在CATH层次结构的类、体系结构和拓扑级别，不包括最后一级的同源超家族，后者主要由序列相似性驱动。作者将S-PLM生成的序列表示可视化并对仅依赖序列信息的模型进行基准测试，包括ESM2；S-PLM这两个语言模型，另外两个结构感知模型PromptProtein和ProstT5通过预测序列的3D结构或3D结构标记进行预训练（图3a）。图3a的每一行显示了一个层次结构中最具代表性的五个类别的蛋白质嵌入的t-SNE可视化。结果表明，S-PLM的序列表示比其他模型的嵌入更能清晰地分离CATH结构类。

作者进一步利用Calinski-Harabasz指数（CHI）来定量评估不同方法得到的嵌入在区分CATH结构类别方面的能力。CHI分数量化了簇间离散度之和与簇内离散度之和之间的比率。作者应用CATH分类来定义基真聚类，使用序列嵌入来计算聚类之间和聚类内部的离散度。如图2b所示，对于类、体系结构和拓扑级别，S-PLM的CHI比ProsT5高30%左右，比ESM2高300%左右。鉴于这些CATH分类是使用蛋白质结构建立的，该分析表明，所开发的S-PLM产生的序列嵌入具有固有的蛋白质结构意识，在有效区分具有不同结构特征的蛋白质方面超过了其他模型。仅使用序列的ESM2表现不佳，也表明其在明确获取蛋白质结构知识方面存在局限性。

图2 案例分析

在这项工作中，作者提出了S-PLM，一种通过蛋白质序列和3D结构之间的对比学习进行预训练的结构感知蛋白质语言模型。S-PLM序列编码器通过在预训练的ESM2模型上使用结构感知对比学习模块，生成融合三维结构信息的序列表示，同时保持ESM2原有的序列表示能力。重要的是，这种基于适配器的体系结构是可扩展的；未来如果需要加入新的蛋白质属性，如蛋白质功能、蛋白质-蛋白质相互作用等，现有的ESM2和结构感知对比学习模块可以保持不变。新属性可以通过简单地添加一个并行适配器并对更新的数据进行训练来集成。实验结果表明了S-PLM及其轻量级调优策略作为蛋白质分析和预测任务的替代PLM的潜力。

在未来，通过扩展使用更全面的蛋白质结构和序列数据库，S-PLM有望得到进一步改进。此外，作者提出的适配器调优方法为蛋白质语言模型的持续学习提供了坚实的基础。通过不断完善和扩展训练数据集，作者的目标是为S-PLM提供对不同生物背景下蛋白质结构的广泛理解。这种数据增强和模型改进的迭代过程有望推动基于序列的蛋白质表示学习的界限，使S-PLM在更广泛的蛋白质分析和预测任务中实现更高的效率。

参考文献:

Wang et al. S-PLM: Structure-Aware Protein Language Model via Contrastive Learning Between Sequence and Structure. Advanced Science. 2024