Retrieved Sequence Augmentation for Protein Representation Learning

TL;DR

This paper introduces Retrieved Sequence Augmentation (RSA), a retrieval-based method that enhances protein representation learning by linking sequences to similar ones in a database, improving prediction accuracy and speed without relying on multiple sequence alignments.

Contribution

The paper presents RSA, a novel retrieval-augmented approach for protein modeling that outperforms traditional MSA-based methods in accuracy and efficiency, especially on de novo proteins.

Findings

RSA improves structure and property prediction accuracy by 5%.

RSA is 373 times faster than MSA Transformer.

RSA transfers better to new protein domains.

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.