Efficient inference, training, and fine-tuning of protein language models

Abstract

Protein language models (PLMs) have shown great promise in protein structure and function predictions, but their adoption is limited by computational cost. We address this challenge by enhancing the efficiency of evolutionary scale modeling (ESM). Using FlashAttention and sequence packing, we achieve 4–9× faster inference and 3–14× lower memory usage. Four-bit quantization of billion-parameter models further reduces memory by 2–3× while preserving accuracy for missense variant effect prediction. Training is also optimized, cutting runtime 6-fold with methods, such as activation checkpointing and DeepSpeed zero-offload. Parameter-efficient fine-tuning of a few adapter weights yields state-of-the-art performance at protein property and function predictions, resulting in 70% Spearman’s correlation for melting point and 87% AU-PRC for transcription factor identification. 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 code is available on GitHub.