Retrieved Sequence Augmentation for Protein Representation Learning

Abstract

Protein language models have excelled in a variety of tasks, ranging from structure prediction to protein engineering. However, proteins are highly diverse in functions and structures, and current state-of-the-art models including the latest version of AlphaFold rely on Multiple Sequence Alignments (MSA) to feed in the evolutionary knowledge. Despite their success, heavy computational overheads, as well as the de novo and orphan proteins remain great challenges in protein representation learning. In this work, we show that MSAaugmented models inherently belong to retrievalaugmented methods. Motivated by this finding, we introduce Retrieved Sequence Augmentation(RSA) for protein representation learning without additional alignment or pre-processing. RSA links query protein sequences to a set of sequences with similar structures or properties in the database and combines these sequences for downstream prediction. We show that protein language models benefit from the retrieval enhancement on both structure prediction and property prediction tasks, with a 5% improvement on MSA Transformer on average while being 373 times faster. In addition, we show that our model can transfer to new protein domains better and outperforms MSA Transformer on de novo protein prediction. Our study fills a much-encountered gap in protein prediction and brings us a step closer to demystifying the domain knowledge needed to understand protein sequences. Code is available on https://github.com/HKUNLP/RSA.

Bilingual Language Model for Protein Sequence and Structure

“…Along the same lines, predicted 3Di can help speed up recently developed structural phylogenetics based on 3Di 71 . Deriving embeddings from structures will also expand the power of embedding-based alignments 72,73 , and retrieval Transformers 74 . Our proposed integration of 3D information into pLMs constitutes the first step toward building truly multi-model pLMs that capture the multiple facets of protein structure, function and evolution.…”

“…Constant speed-ups of LLMs [57] also decrease translation latency which in turn renders the translation from AA to 3Di an attractive alternative to searching large metagenomic datasets at the sensitivity of structure comparison while avoiding the overhead of actually having to predict 3D structures. Deriving embeddings from structures will also expand the power of embedding-based alignments [67], [68], and retrieval Transformers [69]. Our proposed integration of 3D information into pLMs may only constitute the first step towards building truly multi-model pLMs that capture the multiple facets of protein structure, function and evolution.…”

Feature Reuse and Scaling: Understanding Transfer Learning with Protein Language Models

“…Second, while a variety of other pretraining tasks have been proposed for protein transfer learning, and different pretraining tasks could potentially learn different aspects of protein biology, we remain uncertain if they will result in significant differences from MLMs. Many pretraining tasks still aim to reconstruct natural sequences (He et al, 2021; Notin et al, 2022; Tan et al, 2023; Ma et al, 2023) and so are also likely to primarily learn coevolutionary patterns. Other tasks use structure as an additional input or target, but they generally make only modest improvements on function prediction tasks (Mansoor et al, 2021; Wang et al, 2022; Yang et al, 2023; Su et al, 2023).…”