[Paper Review] Endowing Protein Language Models with Structural Knowledge

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

Protein Structure and Dynamics13 citations

TL;DR

PST enhances a pretrained protein language model (ESM-2) by integrating a structure extractor in every self-attention block, enabling structure-aware representations that improve protein function and structure predictions with better parameter efficiency.

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.

Motivation & Objective

• Bridge sequence and structure in protein modeling by injecting structural information into transformer-based PLMs.
• Improve protein function and structure prediction accuracy while reducing parameter requirements.
• Demonstrate the versatility of structure-aware representations across diverse downstream tasks without heavy task-specific fine-tuning.
• Show that refining only structure extractors can yield substantial gains, and that partial pretraining strategies can maintain or enhance performance.

Proposed method

• Represent proteins as graphs with nodes for residues and edges for proximity within 8Å.
• Integrate a structure extractor (e.g., 2-layer GIN) into each self-attention block of a pretrained ESM-2 backbone.
• Modify Q, K, V with structure-aware embeddings produced by the extractor (Eq. 5).
• Pretrain the PST model on AlphaFoldDB SwissProt subset using the same MLM objective as ESM-2.
• Optionally update only the structure extractor parameters, or the whole model, during pretraining.
• Evaluate PST by fixed representations (MLP/linear heads) across tasks without task-specific fine-tuning.

Experimental results

Research questions

• RQ1 Does integrating structural information into transformer self-attention improve protein function and structure prediction beyond backbone PLMs?
• RQ2 What is the impact of using structure extractors on parameter efficiency and performance across model sizes?
• RQ3 How does PST perform on GO and EC function prediction, fold classification, and ProteinShake tasks compared to state-of-the-art sequence and structure models?
• RQ4 Is updating only the structure extractors during pretraining sufficient to achieve competitive representations, and can combining structural and sequence representations at inference boost results?

Key findings

• PST achieves state-of-the-art performance on function prediction tasks such as enzyme and gene ontology classification.
• PST representations are robust across tasks even with fixed representations, reducing the need for task-specific fine-tuning.
• PST consistently outperforms ESM-2, with larger gains for smaller ESM-2 backbones, highlighting parameter efficiency benefits.
• Pretraining the entire PST yields the best results, but updating only the structure extractors provides comparable performance with higher efficiency.
• Incorporating more subtle structural information improves pretraining accuracy but may decrease downstream task performance, suggesting the need for more advanced objectives.