Endowing Protein Language Models with Structural Knowledge

TL;DR

This work tackles the challenge of linking protein sequence, structure, and function by enhancing pretrained protein language models with explicit structural knowledge. It introduces the Protein Structure Transformer (PST), which injects a structure extractor into every self-attention block of an ESM-2 backbone and is pretrained on a 542K structure dataset using the masked language modeling objective $ \mathcal{L}_{\mathrm{MLM}}$. PST achieves superior parameter efficiency and consistently outperforms state-of-the-art sequence models on function and structure prediction benchmarks, with notable gains for smaller base models and robust zero-shot and residue-level performance. The findings highlight the potential of structure-aware PLMs to leverage extensive structure databases for scalable, effective protein modeling, and point to future work on larger structure corpora and richer training objectives beyond MLM.

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.

Figures (7)

• Figure 1: Overview of the proposed Protein Structure Transformer. A protein 3D structure is converted to an ordered graph with nodes representing amino acid types and edges linking any pair of residues within a specified threshold (8Å). Then, we sample a fraction of nodes and mask them using a special mask token. The output of Protein Structure Transformer (PST) is fed into a linear prediction head to predict the original amino acid types. Compared to ESM-2, PST uses a GNN to extract local structural features around each node before computing the self-attention at each transformer layer. The PST model is initialized with the pretrained ESM-2 weights.
• Figure 2: Performance of PST models trained with and without distance information used as edge attributes on ProteinShake tasks.
• Figure 3: Performance of PST and ESM-2 across varied model sizes on ProteinShake datasets as well as DeepFRI and VEP datasets.
• Figure 4: Effect of pretraining strategies on model performance. "Full" refers to the strategy where one updates the full model during pretraining, including both ESM-2 and structure extractor weights. "Struct Only" refers to the strategy where only the structure extractor weights are being updated during training. "Struct Only + Seq" is an extension of "Struct Only" at inference. By bypassing the structure extractors, the PST model is capable to obtain the same sequence representations as the base ESM-2 model. Averaging both structure and sequence representations leads to "Struct Only + Seq".
• Figure 5: Overview of the pretraining pipeline of the propose Protein Structure Transformer.