Retrieved Sequence Augmentation for Protein Representation Learning (2024)

Abstract: Protein Language Models traditionally depend on Multiple Sequence Alignments (MSA) to incorporate evolutionary knowledge. However, MSA-based approaches suffer from substantial computational overhead and generally underperform in generalizing to de novo proteins. This study reevaluates the role of MSA, proposing it as a retrieval augmentation method and questioning the necessity of sequence alignment. We show that a simple alternative, Retrieved Sequence Augmentation (RSA), can enhance protein representation learning without the need for alignment and cumbersome preprocessing. RSA surpasses MSA Transformer by an average of 5% in both structural and property prediction tasks while being 373 times faster. Additionally, RSA demonstrates enhanced transferability for predicting de novo proteins. This methodology addresses a critical need for efficiency in protein prediction and can be rapidly employed to identify homologous sequences, improve representation learning, and enhance the capacity of Large Language Models to interpret protein structures.

## introduction proteins are fundamental yet complex components of life. they exhibit a diverse range of functions within organisms. the enigmatic characteristic of these macromolecules originates from the intricate interplay between their sequences, structures, and functions, which is influenced jointly by physics and evolution @xcite . protein language models @xcite @xcite capture the co-occurrence probability of amino acids observed in nature, thus encapsulating structural and evolutionary information within the resulting representations. while this approach has demonstrated its effectiveness (elnaggar et al., 2021; jumper et al., 1 code and data are available at this repo. 2021; @xcite @xcite , the evolutionary knowledge that can be extracted from a single sequence remains insufficient @xcite . in order to compensate for this limitation, multiple sequence alignments (msa) @xcite @xcite ) have been extensively used as a foundational protein feature engineering technique to extract protein evolutionary information in protein models @xcite @xcite . msa draws on the evolutionary principle that functional constraints of species govern the mutation rate, which in turn drives the convergence of sequences. therefore, key residues at functional sites tend to be conserved across protein families. msa primarily aligns these conserved regions across homologous proteins to identify critical functional residues, such as substrate binding sites @xcite . traditional approaches @xcite such as potts model @xcite , directly extracts statistical features from msa for structural prediction. in recent studies @xcite @xcite , mod-protein task: where is the protein mgsnkskpkdasqrrrslepaen located within a cell? p(y|x) = σp(y|r,x)p(r|x) the protein is located on the membrane. ## rsa motivated by the potential of pre-trained retrievers to identify proteins that are homologous or geometric similar, we propose a pipeline, rsa (retrieval sequence augmentation), to directly augment protein models on-the-fly. rsa follows the retrievethen-predict framework in eq. 1. it comprises of a neural sequence retriever p(r|x), and a protein model that combines both original input and retrieved sequence to obtain prediction p(y|x, r). rsa retriever is used for finding the sequences that are semantically close to the query. the similarity score f (x, r) is defined as the negative l2 distance between the embedding g of the two sequences. for protein retrieval, we aim to retrieve protein sequences that have similar structures or are homologous to the query sequence. motivated by the high accuracy of k-nearest neighbor retrieval with esm-1b @xcite pre-trained embeddings (as shown in training for downstream finetuning, we maximize p(y|x) by training on the retrieval augmented protein encoder. we freeze the retriever parameters during training. for a query sequence of length l with n retrieved proteins, suppose the length of retrieved proteins l ′ ≤ l the computation cost is n times the original model, o(n l 2 ) for a transformer encoder layer, which is as efficient as msa transformer with o(n l 2 ) + o(n 2 l) complexity. ## related work retrieval-augmented language models the integration of non-parametric retrieval and parametric models has been revolutionary for many problems @xcite @xcite @xcite . retrieval-augmentation introduces knowledge from memories and improve performance in complex reasoning @xcite and generalization @xcite . protein language models to model and further understand the protein sequence data, language models are introduced to train on mass data @xcite . large scale pre-training enables language models to learn structural and evolutionary knowledge @xcite @xcite . despite these successes, many important applications still require msas and other external knowledge @xcite @xcite @xcite . msas have been shown effective in improving representation learning, despite being extremely slow and costly in computation. @xcite and @xcite use dense retrieval to accelerate multiple sequence augmentation, while still dependent on alignment procedures. recent work @xcite @xcite explores msa-free language models though additional pretraining is involved. we take this step further to investigate retrieval-augmented protein language models. another line of work improves msa quality and generation speed by using generative models to augment and produce msas @xcite @xcite . we differ from this line of work as we discuss substitutes rather than augmentation of msas. ## conclusion this work has highlighted the limitations inherent in traditional msa-based approaches and proposed rsa as a substitute. through extensive evaluation, we have demonstrated a significant improvement not only in the speed of processing-being more than 300 times faster than msa-but also in enhancing predictive performance in downstream tasks with various models.