- Physical Review Journals
All JournalsPhysics Magazine

PRX Life

Reuse & Permissions

It is not necessary to obtain permission to reuse this article or its components as it is available under the terms of the Creative Commons Attribution 4.0 International license. This license permits unrestricted use, distribution, and reproduction in any medium, provided attribution to the author(s) and the published article's title, journal citation, and DOI are maintained. Please note that some figures may have been included with permission from other third parties. It is your responsibility to obtain the proper permission from the rights holder directly for these figures.

Export citation

Export citation

Choose format for download:

Download Citation
  • Open Access

RAG-ESM: Improving Pretrained Protein Language Models via Sequence Retrieval

Damiano Sgarbossa* and Anne-Florence Bitbol†

  • *Contact author: damiano.sgarbossa@epfl.ch
  • †Contact author: anne-florence.bitbol@epfl.ch
PDF

PRX Life 3, 033013 – Published 19 August, 2025

DOI: https://doi.org/10.1103/db1b-hy16

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.

View figure in article

Physics Subject Headings (PhySH)

Collections

This article appears in the following collection:

Machine Learning and Structural Biology

Curated by the organizers of the Machine Learning for Structural Biology Workshop at NeurIPS, this collection features studies at the intersection of computer science and biology, showcasing how machine learning is impacting our understanding of biological molecules and their processes.

Every article published in this collection underwent a rigorous peer review process, adhering to the same high standards applied to all papers handled by the journal. The PRX Life editorial team managed the peer review and made all editorial decisions.

Article Text

Supplemental Material

References (47)

  1. The UniProt Consortium, UniProt: The universal protein knowledgebase in 2021, Nucl. Acids Res. 49, D480 (2021).
  2. M. Weigt, R. A. White, H. Szurmant, J. A. Hoch, and T. Hwa, Identification of direct residue contacts in protein-protein interaction by message passing, Proc. Natl. Acad. Sci. USA 106, 67 (2009).
  3. F. Morcos, A. Pagnani, B. Lunt, A. Bertolino, D. S. Marks, C. Sander, R. Zecchina, J. N. Onuchic, T. Hwa, and M. Weigt, Direct-coupling analysis of residue coevolution captures native contacts across many protein families, Proc. Natl. Acad. Sci. USA 108, E1293 (2011).
  4. D. S. Marks, L. J. Colwell, R. Sheridan, T. A. Hopf, A. Pagnani, R. Zecchina, and C. Sander, Protein 3D structure computed from evolutionary sequence variation, PLoS ONE 6, e28766 (2011).
  5. W. P. Russ, M. Figliuzzi, C. Stocker, P. Barrat-Charlaix, M. Socolich, P. Kast, D. Hilvert, R. Monasson, S. Cocco, M. Weigt, and R. Ranganathan, An evolution-based model for designing chorismate mutase enzymes, Science 369, 440 (2020).
  6. A. Hawkins-Hooker, F. Depardieu, S. Baur, G. Couairon, A. Chen, and D. Bikard, Generating functional protein variants with variational autoencoders, PLoS Comput. Biol. 17, e1008736 (2021).
  7. A. Elnaggar, M. Heinzinger, C. Dallago, G. Rehawi, Y. Wang, L. Jones, T. Gibbs, T. Feher, C. Angerer, M. Steinegger, D. Bhowmik, and B. Rost, ProtTrans: Towards cracking the language of life's code through self-supervised deep learning and high performance computing, in Proceedings of the IEEE Transactions on Pattern Analysis and Machine Intelligence (IEEE, New York, 2021), p. 1.
  8. J. Vig, A. Madani, L. R. Varshney, C. Xiong, R. Socher, and N. Rajani, BERTology meets biology: Interpreting attention in protein language models, in International Conference on Learning Representations (ICLR, 2021).
  9. A. Rives, J. Meier, T. Sercu, S. Goyal, Z. Lin, J. Liu, D. Guo, M. Ott, C. L. Zitnick, J. Ma, and R. Fergus, Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences, Proc. Natl. Acad. Sci. USA 118, e2016239118 (2021).
  10. A. Madani, B. McCann, N. Naik, N. S. Keskar, N. Anand, R. R. Eguchi, P.-S. Huang, and R. Socher, ProGen: Language modeling for protein generation, bioRxiv: 2020.03.07.982272.
  11. A. Madani, B. Krause, E. R. Greene, S. Subramanian, B. P. Mohr, J. M. Holton, J. L. Olmos, C. Xiong, Z. Z. Sun, R. Socher, J. S. Fraser, and N. Naik, Large language models generate functional protein sequences across diverse families, Nat. Biotechnol. 41, 1099 (2023).
  12. N. Ferruz, S. Schmidt, and B. Höcker, ProtGPT2 is a deep unsupervised language model for protein design, Nat. Commun. 13, 4348 (2022).
  13. J. Meier, R. Rao, R. Verkuil, J. Liu, T. Sercu, and A. Rives, Language models enable zero-shot prediction of the effects of mutations on protein function, Adv. Neural Inf. Proc. Syst. 34, 29287 (2021).
  14. P. Kantroo, G. P. Wagner, and B. B. Machta, Pseudo-perplexity in one fell swoop for protein fitness estimation, PRX Life (2025), doi:10.1103/zhx7-hcmm.
  15. T. Bepler and B. Berger, Learning protein sequence embeddings using information from structure, in International Conference on Learning Representations (ICLR, 2019).
  16. D. Sgarbossa, C. Malbranke, and A.-F. Bitbol, ProtMamba: A homology-aware but alignment-free protein state space model, bioRxiv: 2024.05.24.595730.
  17. S. Alamdari, N. Thakkar, R. van den Berg, A. X. Lu, N. Fusi, A. P. Amini, and K. K. Yang, Protein generation with evolutionary diffusion: Sequence is all you need, bioRxiv: 2023.09.11.556673.
  18. X. Wang, Z. Zheng, F. Ye, D. Xue, S. Huang, and Q. Gu, Diffusion language models are versatile protein learners, in Proceedings of the 41st International Conference on Machine Learning, Vol. 235 (PMLR, 2024), pp. 52309–52333.
  19. R. M. Rao, J. Liu, R. Verkuil, J. Meier, J. Canny, P. Abbeel, T. Sercu, and A. Rives, MSA transformer, in Proceedings of the 38th International Conference on Machine Learning, edited by M. Meila and T. Zhang, Proceedings of Machine Learning Research Vol. 139 (PMLR, 2021), pp. 8844–8856.
  20. J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, A. Bridgland, C. Meyer, S. A. A. Kohl, A. J. Ballard, A. Cowie, B. Romera-Paredes, S. Nikolov, R. Jain, J. Adler, T. Back, S. Petersen, D. Reiman, E. Clancy, M. Zielinski, M. Steinegger, M. Pacholska, T. Berghammer, S. Bodenstein, D. Silver, O. Vinyals, A. W. Senior, K. Kavukcuoglu, P. Kohli, and D. Hassabis, Highly accurate protein structure prediction with Alpha Fold, Nature (London) 596, 583 (2021).
  21. Z. Lin, H. Akin, R. Rao, B. Hie, Z. Zhu, W. Lu, N. Smetanin, R. Verkuil, O. Kabeli, Y. Shmueli, A. dos Santos Costa, M. Fazel-Zarandi, T. Sercu, S. Candido, and A. Rives, Evolutionary-scale prediction of atomic-level protein structure with a language model, Science 379, 1123 (2023).
  22. A. Elofsson, Progress at protein structure prediction, as seen in CASP15, Curr. Opin. Struct. Biol. 80, 102594 (2023).
  23. J. D. Thompson, B. Linard, O. Lecompte, and O. Poch, A comprehensive benchmark study of multiple sequence alignment methods: current challenges and future perspectives, PLoS ONE 6, e18093 (2011).
  24. T. F. Truong and T. Bepler, PoET: A generative model of protein families as sequences-of-sequences, in Proceedings of the 37th International Conference on Neural Information Processing Systems (Curran Associates, Red Hook, NY, 2023).
  25. P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W.-T. Yih, T. Rocktäschel, S. Riedel, and D. Kiela, Retrieval-augmented generation for knowledge-intensive nlp tasks, in Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20 (Curran Associates, Red Hook, NY, 2020).
  26. K. Guu, K. Lee, Z. Tung, P. Pasupat, and M.-W. Chang, Retrieval augmented language model pre-training, in Proceedings of the 37th International Conference on Machine Learning, ICML'20, Vol. 119 (PMLR, 2020), pp. 3929–3938.
  27. Z. Zheng, Y. Deng, D. Xue, Y. Zhou, F. Ye, and Q. Gu, Structure-informed language models are protein designers, in Proceedings of the 40th International Conference on Machine Learning, ICML'23, Vol. 202 (PMLR, 2023), pp. 42317–42338.
  28. See Supplemental Material at http://link.aps.org/supplemental/10.1103/db1b-hy16 for additional information on training, ablations, sequence generation method, and motif scaffolding model comparisons.
  29. E. Nijkamp, J. A. Ruffolo, E. N. Weinstein, N. Naik, and A. Madani, Progen2: exploring the boundaries of protein language models, Cell Syst. 14, 968 (2023).
  30. T. Hayes, R. Rao, H. Akin, N. J. Sofroniew, D. Oktay, Z. Lin, R. Verkuil, V. Q. Tran, J. Deaton, M. Wiggert, R. Badkundri, I. Shafkat, J. Gong, A. Derry, R. S. Molina, N. Thomas, Y. Khan, C. Mishra, C. Kim, L. J. Bartie, M. Nemeth, P. D. Hsu, T. Sercu, S. Candido, and A. Rives, Simulating 500 million years of evolution with a language model, Science 387, 850 (2025).
  31. N. Bhattacharya, N. Thomas, R. Rao, J. Dauparas, P. K. Koo, D. Baker, Y. S. Song, and S. Ovchinnikov, Interpreting Potts and transformer protein models through the lens of simplified attention, Biocomputing 2022 (World Scientific, Singapore, 2022), pp. 34–45.
  32. S. S. Sahoo, M. Arriola, Y. Schiff, A. Gokaslan, E. Marroquin, J. T. Chiu, A. Rush, and V. Kuleshov, Simple and effective masked diffusion language models, in The Thirty-eighth Annual Conference on Neural Information Processing Systems (NeurIPS, 2024).
  33. M. Mirdita, L. von den Driesch, C. Galiez, M. J. Martin, J. Söding, and M. Steinegger, Uniclust databases of clustered and deeply annotated protein sequences and alignments, Nucl. Acids Res. 45, D170 (2017).
  34. G. Ahdritz, N. Bouatta, S. Kadyan, L. Jarosch, D. Berenberg, I. Fisk, A. M. Watkins, S. Ra, R. Bonneau, and M. AlQuraishi, OpenProteinSet: Training data for structural biology at scale, in Proceedings of the 37th International Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Curran Associates, Red Hook, NY, 2023).
  35. S. Petti, N. Bhattacharya, R. Rao, J. Dauparas, N. Thomas, J. Zhou, A. M. Rush, P. K. Koo, and S. Ovchinnikov, End-to-end learning of multiple sequence alignments with differentiable Smith-Waterman, Bioinformatics 39, btac724 (2023).
  36. U. Lupo, D. Sgarbossa, and A.-F. Bitbol, Protein language models trained on multiple sequence alignments learn phylogenetic relationships, Nat. Commun. 13, 6298 (2022).
  37. Y. Zhang and J. Skolnick, Scoring function for automated assessment of protein structure template quality, Proteins: Structure, Proteins 57, 702 (2004).
  38. S. R. Eddy, Accelerated profile HMM searches, PLoS Comput Biol. 7, e1002195 (2011).
  39. J. Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. J. Ragotte, L. F. Milles, B. I. M. Wicky, A. Courbet, R. J. de Haas, N. Bethel, P. J. Y. Leung, T. F. Huddy, S. Pellock, D. Tischer, F. Chan, B. Koepnick, H. Nguyen, A. Kang, B. Sankaran, A. K. Bera, N. P. King, and D. Baker, Robust deep learning–based protein sequence design using ProteinMPNN, Science 378, 49 (2022).
  40. J. L. Watson, D. Juergens, N. R. Bennett, B. L. Trippe, J. Yim, H. E. Eisenach, W. Ahern, A. J. Borst, R. J. Ragotte, L. F. Milles, B. I. M. Wicky, N. Hanikel, S. J. Pellock, A. Courbet, W. Sheffler, J. Wang, P. Venkatesh, I. Sappington, S. V. Torres, A. Lauko, V. D. Bortoli, E. Mathieu, S. Ovchinnikov, R. Barzilay, T. S. Jaakkola, F. DiMaio, M. Baek, and D. Baker, De novo design of protein structure and function with RFdiffusion, Nature (London) 620, 1089 (2023).
  41. R. Wu, F. Ding, R. Wang, R. Shen, X. Zhang, S. Luo, C. Su, Z. Wu, Q. Xie, B. Berger, J. Ma, and J. Peng, High-resolution de novo structure prediction from primary sequence, bioRxiv: 2022.07.21.500999.
  42. X. Wang, Z. Zheng, F. Ye, D. Xue, S. Huang, and Q. Gu, DPLM-2: A multimodal diffusion protein language model, in The Thirteenth International Conference on Learning Representations (ICLR, 2025).
  43. G. Munsamy, R. Illanes-Vicioso, S. Funcillo, I. T. Nakou, S. Lindner, G. Ayres, L. S. Sheehan, S. Moss, U. Eckhard, P. Lorenz, and N. Ferruz, Conditional language models enable the efficient design of proficient enzymes, bioRxiv: 2024.05.03.592223.
  44. J. Su, C. Han, Y. Zhou, J. Shan, X. Zhou, and F. Yuan, SaProt: Protein language modeling with structure-aware vocabulary, in The Twelfth International Conference on Learning Representations (ICLR, 2024).
  45. M. Heinzinger, K. Weissenow, J. G. Sanchez, A. Henkel, M. Mirdita, M. Steinegger, and B. Rost, Bilingual language model for protein sequence and structure, NAR Genom. Bioinf. 6, lqae150 (2024).
  46. K. Michalewicz, M. Barahona, and B. Bravi, Integrating protein sequence embeddings with structure via graph-based deep learning for the prediction of single-residue properties, arXiv:2502.17294.
  47. D. Sgarbossa and A.-F. Bitbol, https://github.com/Bitbol-Lab/rag-esm (2025).

Outline

Information

PRX Life 3, 033013– Published 19 August, 2025
  • Received 6 April 2025
  • Accepted 30 July 2025
Crossmark - Check for updates button
Creative Commons Logo

Published by the American Physical Society under the terms of the Creative Commons Attribution 4.0 International license. Further distribution of this work must maintain attribution to the author(s) and the published article's title, journal citation, and DOI.

Published by the American Physical Society

Sign In to Your Journals Account

Filter

Filter

Article Lookup

Enter a citation