2024年12月11日,NeurIPS上发表论文:“ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention”。
蛋白质语言模型(PLMs)在各种蛋白质功能预测任务中表现出了卓越的能力。然而,尽管蛋白质功能与结构密切相关,大多数现有的蛋白质语言模型并未充分融入蛋白质结构信息。
为了解决这一问题,作者提出了
ProSST,一个基于Transformer的蛋白质语言模型,能够无缝地结合蛋白质序列和结构。
ProSST融合了一个结构量化模块和一个具有解耦注意力机制的Transformer架构。结构量化模块通过将蛋白质的三维结构先序列化为残基级别的局部结构,然后将其嵌入到密集向量空间中,进而利用一个预训练的聚类模型将这些向量量化为离散的结构令牌,从而为蛋白质结构提供有效的表示。这些令牌作为蛋白质结构的有效表征。此外,ProSST通过序列-结构解耦注意力机制显式地学习蛋白质残基令牌序列和结构令牌序列之间的关系。
作者使用掩蔽语言模型目标在数百万个蛋白质结构上预训练ProSST,使其能够学习全面的上下文表示。为了评估所提出的ProSST,作者在zero-shot突变效应预测和若干监督学习下游任务上进行了广泛的实验,结果表明,
ProSST在所有基准模型中达到了最先进的性能。
ProSST的模型框架如下:
ProSST主要包含两个模块:结构量化模块和具有序列-结构解耦注意力机制的Transformer模型。
图A展示了如何训练GVP编码器,GVP作为一种几何向量感知器,是专门设计用于处理几何特征的深度学习模型,用于编码一个残基及其局部结构中的领域信息。作者将GVP与位置感知多层感知器的解码器集成在一起,形成一个自动编码器模型。整个模型通过去噪预训练目标进行训练。在此过程中,还用3D高斯噪声扰动Ca坐标,并根据RF-Diffusion在旋转矩阵流形上使用布朗运动。模型的训练目标是恢复蛋白质结构至其原始、无噪状态。
图B展示了如何构建Local Structure Code Book用来实现将蛋白质局部结构的连续特征表示转化为离散标记。首先,利用结构编码器GVP将C.A.T.H数据集中所有残基的局部结构嵌入到连续潜在空间中。然后,然后应用k-means算法在潜在空间中识别K个质心,这些质心构成了结构码本。对于任何局部结构嵌入,它通过码本中最近的向量进行量化。
图C展示了如何将蛋白质的全局三维结构序列化,并将其量化为离散的结构标记。首先,对于蛋白质序列中每个位置的残基,基于每个残基的局部结构构建了图,然后使用结构编码器将其转化为一个连续向量。接着,利用构建好的结构码本为每个向量分配一个结构标记。
受DeBerta模型的启发,
作者使用了一种扩展形式的解耦注意力,将残基序列、结构序列以及相对位置信息的注意力结合起来。
具体来说,对于蛋白质序列中位置 i 的残基,它可以由以下三部分表示:
R
i
表示表示残基的隐藏状
态;
S
i
表示残基的局部结构的嵌
入;
P
i
|j
表示残基i相对于残基j的相对位置嵌入。
两个残基i和j的交叉注意力A可以分解为9个组件,公式如下:
模型主要关注于残基的上下文嵌入,像结构-结构、结构-位置等与残基无关的组件被移除。因此,最终只保留了5种类型的注意力。
ProSST通过结构条件掩码语言建模进行预训练。在这种方法中,每个输入序列x会通过将一部分氨基酸残基替换为特殊的掩码标记或其他随机残基的方式添加噪声。ProSST的目标是利
用被噪声化的序列和其对应的结构标记序列s作为上下文,来预测输入序列中原本的残基。
数据集:在ProteinGym数据集上进行了实验,并使用AlphaFold2生成序列的结构
研究表明,ProSST性能优越。
在 ProteinGYM 的零样本突变效应预测中,ProSST 的表现优于所有基线模型,并在稳定性、结合能力上达到最先进水平,尤其在稳定性预测中表现最佳。另外,
性能提升来源于结构信息建模:
移除结构信息的简化版 ProSST 的性能与其他基于序列的模型相近,证明模型的性能提升主要得益于高效的结构信息建模,而非其他因素。
下游任务:金属离子结合预测、蛋白质定位预测、蛋白质热稳定性预测、GO(Gene Ontology)注释预测。
研究表明,
ProSST在所有模型中表现最佳,在6个实验中获得了4项第一。
在DeepLoc和金属离子结合任务中,ProSST远超其他方法。
在热稳定性以及GO-BP和GO-CC任务中,ESM-GearNet的结果与ProSST相当,但其模型规模是ProSST的6倍以上。此外,
序列-结构模型表现更优:
ESM-GearNet、SaProt和ProSST等序列-结构结合模型的表现优于其他对比模型,这表明结构信息在蛋白质建模中的重要性。
此外,ProSST在整合蛋白质序列和结构信息方面比SaProt更具优势,验证了作者设计的有效性。
研究表明:
1.局部结构词汇表大小K的影响:随着K的增加,ProSST在所有指标上的性能持续提升,并在K=2048时达到最佳性能,因此将K=2048设置为默认值。
2.收敛性提升:K的增加使ProSST的收敛性逐步提高,表明引入结构提示能够有效增强模型的表示能力。
3.结构量化方法的有效性:在相同网络架构下,所提出的结构量化方法性能优于Foldseek和DSSP,证明了其设计的优势。
4.结构信息的重要性:ProSST(Foldseek)、ProSST(DSSP)和K>0的ProSST模型在所有指标上均显著优于K=0的模型,再次验证了结构信息对模型性能提升的关键作用。
5.参数增加非主要因素:K=1和K=0的ProSST性能几乎无差异,表明性能提升并非源于解耦注意力中参数的增加,而是归功于结构信息的有效建模。