2Bits of Protein: Efficient Protein Language Models at the Scale of 2-bits
Keywords: pLM, language model, quantization, encoder, low precision
TL;DR: We train a ternary precision pLM and show competetive performance with the mixed-precision equivalent, while delivering significant energy, memory, and compute usage improvements.
Abstract: Protein language models have become an increasingly popular tool across various biological tasks, from variant effect prediction to novel sequence generation. However, state-of-the-art models often have up to billions of parameters. Such large model architectures restrict usage to groups with the necessary compute infrastructure or necessitate the use of cloud computing, incurring substantial costs and raising data privacy concerns. In this work, we investigate a ternary protein language model, which uses low-precision weights to reduce model size, energy demand, and computational requirements, making it suitable for operation on edge devices such as laptops. This addresses privacy concerns by ensuring data remains on-device and eliminates the costs associated with cloud services. We train a ternary protein language model and benchmark it against ESM-2 (8M) using the ProteinGym benchmark, demonstrating that our model achieves comparable performance while being more suitable for edge deployment. A discussion is provided on ways to improve the ternary model to outperform ESM-2 in terms of accuracy.
Submission Number: 52
Loading