Efficient inference, training, and fine-tuning of protein language models
Open accessPublished version of a preprint
Abstract
Abstract Protein language models have shown great promise in predicting protein structure, function, and the effects of missense variants on protein fitness. However, their use has been limited by the substantial computational resources required. In this work, we focus on improving the computational efficiency of protein language models (PLMs), specifically the Evolutionary Scale Modeling (ESM) family, to increase the accessibility of PLMs. By implementing optimizations such as FlashAttention and Partition-Attention, a novel technique designed to handle proteins of variable length, we achieved a 16-fold speedup in inference time and reduced memory usage by 3 to 14 times for long proteins. Additionally, 4-bit quantization applied to billion-parameter models led to a 2 to 3 times reduction in memory consumption with minimal performance loss in the missense variant effect prediction task. Training efficiency was also improved, with a 6-fold reduction in runtime achieved through activation checkpointing and the DeepSpeed Zero-Offload strategy. For fine-tuning, we employed parameter-efficient methods, enabling state-of-the-art predictions of protein properties and functions by training only the model head or a small fraction of adapter weights. For instance, we achieved a Spearman’s correlation coefficient of 70% in melting point prediction and an 87% area under the precision-recall curve (AU-PRC) for transcription factor prediction. Our efficient ESM (ESME) implementation significantly lowers the barrier to using these powerful models, making them accessible to academic laboratories with limited computational resources. The ESME implementation is available on PyPI ( pypi.org/project/esm-efficient ) and GitHub ( github.com/uci-cbcl/esm-efficient ).
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Self-check
Percentages, sums and reported differences re-computed from the paper’s own numbers.
The self-check needs the paper’s full text.
Cites
Adam: A Method for Stochastic Optimization
Kingma, D. P., & Ba, J. (2014) · arXiv · doi:10.48550/arxiv.1412.6980 · Cited by 98,946 · see citing works
Highly accurate protein structure prediction with AlphaFold
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Žídek, A., Potapenko, A., et al. (2021) · Nature · doi:10.1038/s41586-021-03819-2 · Cited by 43,764 · see citing works
Accurate structure prediction of biomolecular interactions with AlphaFold 3
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T., Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J., et al. (2024) · Nature · doi:10.1038/s41586-024-07487-w · Cited by 15,861 · see citing works
Evolutionary-scale prediction of atomic level protein structure with a language model
Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z., Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y., et al. (2022) · doi:10.1101/2022.07.20.500902 · Cited by 5,523 · see citing works
LoRA: Low-Rank Adaptation of Large Language Models
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L., & Chen, W. (2021) · arXiv · doi:10.48550/arxiv.2106.09685 · Cited by 2,523 · see citing works
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022) · arXiv · doi:10.48550/arxiv.2205.14135 · Cited by 1,214 · see citing works
Disease variant prediction with deep generative models of evolutionary data
Frazer, J., Notin, P., Dias, M., Gomez, A., Min, J. K., Brock, K., Gal, Y., & Marks, D. S. (2021) · Nature · doi:10.1038/s41586-021-04043-8 · Cited by 791 · see citing works
Mutation effects predicted from sequence co-variation
Hopf, T. A., Ingraham, J. B., Poelwijk, F. J., Schärfe, C. P. I., Springer, M., Sander, C., & Marks, D. S. (2017) · Nature Biotechnology · doi:10.1038/nbt.3769 · Cited by 775 · see citing works
QLoRA: Efficient Finetuning of Quantized LLMs
Dettmers, T., Pagnoni, A., Holtzman, A., & Zettlemoyer, L. (2023) · arXiv · doi:10.48550/arxiv.2305.14314 · Cited by 507 · cited 2× · see citing works
“We conducted in-depth benchmarking of the original ESM and efficient ESME reimplementation for inference, training, and downstream fine-tuning tasks (Figure S1).”
ProGen2: Exploring the boundaries of protein language models
Nijkamp, E., Ruffolo, J. A., Weinstein, E. N., Naik, N., & Madani, A. (2023) · Cell Systems · doi:10.1016/j.cels.2023.10.002 · Cited by 480 · see citing works
Rao, R., Liu, J., Verkuil, R., Meier, J., Canny, J. F., Abbeel, P., Sercu, T., & Rives, A. (2021) · doi:10.1101/2021.02.12.430858 · Cited by 350 · see citing works
A Study of BFLOAT16 for Deep Learning Training
Kalamkar, D., Mudigere, D., Mellempudi, N., Das, D., Banerjee, K., Avancha, S., Vooturi, D. T., Jammalamadaka, N., Huang, J., Yuen, H., et al. (2019) · arXiv · doi:10.48550/arxiv.1905.12322 · Cited by 169 · cited 1× · see citing works
“Inference performance was bench-marked on NVIDIA A100 GPUs using ESM2 models ranging from 8 million to 15 billion parameters, with a batch size of 16 proteins in brain floating point (bfloat16) 49 for both original and efficient implementations.”
ZeRO-Offload: Democratizing Billion-Scale Model Training
Ren, J., Rajbhandari, S., Aminabadi, R. Y., Ruwase, O., Yang, S., Zhang, M., Li, D., & He, Y. (2021) · arXiv · doi:10.48550/arxiv.2101.06840 · Cited by 151 · see citing works
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Dao, T. (2023) · arXiv · doi:10.48550/arxiv.2307.08691 · Cited by 143 · cited 2× · see citing works
“46 FlashAttention is an efficient, IO-aware self-attention algorithm that optimizes memory usage and runtime for long sequences by minimizing read/write operations between high-bandwidth memory (HBM) and SRAM while fusing operations to reduce memory overhead (STAR Methods).”
Genomic language model predicts protein co-regulation and function
Hwang, Y., Cornman, A. L., Kellogg, E. H., Ovchinnikov, S., & Girguis, P. R. (2024) · Nature Communications · doi:10.1038/s41467-024-46947-9 · Cited by 99 · cited 1× · see citing works
“52 We explored weight quantization techniques to optimize inference by benchmarking 8-bit quantization with LLM.int8() 53 and 4-bit quantization using QLoRA 54 algorithms using bitsandbytes package.”
Flagged first, then by citations
- Authors
- Çelik, M. H., & Xie, X.
- Affiliations
- University of California, Irvine (US)
- Corresponding author
- Xie, Xiaohui
- Venue
- iScience
- Publisher
- Elsevier BV
- Published
- 2025-10
- Volume
- 28
- Issue
- 10
- Pages
- 113495
- ISSN
- 2589-0042 (ISSN-L)
- PMC
- PMC12481099
- OpenAlex
- W4413950477
- Cited by
- 7
- References
- 15
- Citations by year
- 2026: 5 · 2025: 2
- FWCI
- 2.10
- Type
- journal article
- Language
- en
- Open access
- gold
- doi.org
- Licence
- Locations
- iScience (publishedVersion, cc-by)
- bioRxiv (Cold Spring Harbor Laboratory) (acceptedVersion, cc-by)
- Europe PMC (PubMed Central) (publishedVersion, cc-by)
- PubMed (publishedVersion)
- PubMed Central (publishedVersion, cc-by)
- Topics
- Life Sciences › Biochemistry, Genetics and Molecular Biology › Molecular Biology › Machine Learning in Bioinformatics
- Life Sciences › Biochemistry, Genetics and Molecular Biology › Molecular Biology › Protein Structure and Dynamics
- Life Sciences › Biochemistry, Genetics and Molecular Biology › Molecular Biology › RNA and protein synthesis mechanisms
- Keywords
- protein language models, FlashAttention, parameter-efficient fine-tuning, adapters, protein property prediction, protein function prediction, transcription factor identification, computational efficiency
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
