[论文解读] Endowing Protein Language Models with Structural Knowledge

Dexiong Chen, Philip Hartout+3 more|arXiv (Cornell University)|Jan 26, 2024

Protein Structure and Dynamics被引用 13

一句话总结

PST 通过在每个自注意力块中集成结构提取器来增强一个预训练的蛋白质语言模型（ESM-2），从而实现结构感知的表示，提升蛋白质功能和结构预测，同时具备更高的参数效率。

ABSTRACT

Understanding the relationships between protein sequence, structure and function is a long-standing biological challenge with manifold implications from drug design to our understanding of evolution. Recently, protein language models have emerged as the preferred method for this challenge, thanks to their ability to harness large sequence databases. Yet, their reliance on expansive sequence data and parameter sets limits their flexibility and practicality in real-world scenarios. Concurrently, the recent surge in computationally predicted protein structures unlocks new opportunities in protein representation learning. While promising, the computational burden carried by such complex data still hinders widely-adopted practical applications. To address these limitations, we introduce a novel framework that enhances protein language models by integrating protein structural data. Drawing from recent advances in graph transformers, our approach refines the self-attention mechanisms of pretrained language transformers by integrating structural information with structure extractor modules. This refined model, termed Protein Structure Transformer (PST), is further pretrained on a small protein structure database, using the same masked language modeling objective as traditional protein language models. Empirical evaluations of PST demonstrate its superior parameter efficiency relative to protein language models, despite being pretrained on a dataset comprising only 542K structures. Notably, PST consistently outperforms the state-of-the-art foundation model for protein sequences, ESM-2, setting a new benchmark in protein function prediction. Our findings underscore the potential of integrating structural information into protein language models, paving the way for more effective and efficient protein modeling Code and pretrained models are available at https://github.com/BorgwardtLab/PST.

研究动机与目标

• 通过将结构信息注入基于变换器的蛋白质语言模型，桥接序列与结构在蛋白质建模中的关系。
• 在降低参数需求的同时提升蛋白质功能和结构预测的准确性。
• 展示结构感知表示在不同下游任务中的通用性，而无需大量针对特定任务的微调。
• 证明仅微调结构提取器就能得到显著提升，并且部分预培训策略可以维持或提升性能。

提出的方法

• 将蛋白质表示为图，节点表示残基，边表示在 8Å 内的邻近关系。
• 将结构提取器（如两层 GIN）集成到预训练的 ESM-2 主干的每个自注意力块中。
• 用提取器产生的结构感知嵌入修改 Q、K、V（Eq. 5）。
• 在 AlphaFoldDB SwissProt 子集上对 PST 模型进行预训练，使用与 ESM-2 相同的 MLM 目标。
• 在预训练过程中可选择仅更新结构提取器参数，或更新整个模型。
• 通过固定表示（MLP/线性头）在各任务上对 PST 进行评估，且不进行任务特定微调。

实验结果

研究问题

• RQ1 将结构信息融入变换器自注意力是否能在 backbone PLMs 之上改善蛋白质功能和结构预测？
• RQ2 在不同模型规模下使用结构提取器对参数效率和性能有何影响？
• RQ3 与最先进的序列与结构模型相比，PST 在 GO 和 EC 功能预测、折叠分类以及 ProteinShake 任务上的表现如何？
• RQ4 在预训练中仅更新结构提取器是否足以获得具有竞争力的表示？在推断时将结构和序列表示结合是否能提升结果？

主要发现

• PST 在功能预测任务上达到最新研究水平，例如酶功能和基因本体分类。
• PST 表示在固定表示下在多任务中具有鲁棒性，降低了对任务特定微调的需求。
• PST 始终优于 ESM-2，对于较小的 ESM-2 主干具有更大提升，突出参数效率的优势。
• 对整个 PST 进行预训练可获得最佳结果，但仅更新结构提取器在效率更高的同时也提供了可比的性能。
• 引入更细粒度的结构信息可以提高预训练准确性，但可能降低下游任务的性能，提示需要更高级的目标函数。