RAG-ESM: Improving Pretrained Protein Language Models via Sequence Retrieval
This article has been Reviewed by the following groups
Discuss this preprint
Start a discussion What are Sciety discussions?Listed in
- Evaluated articles (Arcadia Science)
Abstract
Protein language models are significantly advancing the modeling of sequence-function relationships. However, most of them are not directly informed of homology and evolutionary relationships between protein sequences. Here, we propose a method to make them homology aware. We introduce RAG-ESM, a retrieval-augmented-generation (RAG) framework that allows us to condition pretrained ESM2 protein language models on homologous sequences, using a minimal number of additional cross-attention parameters and minimal computational cost. We show that RAG-ESM models outperform larger ESM2 models for masked amino acid prediction. We find that sequence alignment capabilities spontaneously emerge in specific cross-attention heads of RAG-ESM. By using a discrete diffusion objective for training, and by conditioning on homologs during inference, RAG-ESM reaches state-of-the-art performance for conditional protein sequence generation and motif scaffolding, among sequence-based models. Our method thus possesses strong potential for scalable, efficient and controlled protein engineering.
Article activity feed
-
Discrete diffusion objective: we experiment with two different masking techniques, the first is the standard discrete diffusion objective where the masking fraction is sampled from a uniform distribution over (0, 1), in the second we sample the masking fraction 80% of the time from a β(3, 9) distribution and 20% of the time from a uniform distribution over (0, 1). This approach, adapted from (Hayes et al., 2024), aims to balance representation and generation capabilities. It allows the model to observe masking fractions across (0, 1), with an average
. Both these objectives improve the effectiveness for iterative denoising during sequence generation with respect to standard MLMOther objectives could have been chosen beyond ease of implementation, why in particular this objective? Why not a hybrid objective? what abt retrieved …
Discrete diffusion objective: we experiment with two different masking techniques, the first is the standard discrete diffusion objective where the masking fraction is sampled from a uniform distribution over (0, 1), in the second we sample the masking fraction 80% of the time from a β(3, 9) distribution and 20% of the time from a uniform distribution over (0, 1). This approach, adapted from (Hayes et al., 2024), aims to balance representation and generation capabilities. It allows the model to observe masking fractions across (0, 1), with an average
. Both these objectives improve the effectiveness for iterative denoising during sequence generation with respect to standard MLMOther objectives could have been chosen beyond ease of implementation, why in particular this objective? Why not a hybrid objective? what abt retrieved neighbor training?
-
Here, we propose a method to make them homology-aware. We introduce RAG-ESM, a retrieval-augmented framework that allows to condition pretrained ESM2 protein language models on homologous sequences, using a minimal number of additional cross-attention parameters and minimal computational cost
This is an interesting idea. I wonder what the scaling looks like and what the efficacy of the augmentation is with respect to context window size and the quality of retrieval.
-
-
