Energy Efficient Protein Language Models: Leveraging Small Language Models with LoRA for Controllable Protein Generation

TL;DR

This paper introduces small, energy-efficient protein language models using LoRA, capable of controllable protein generation with high structural accuracy, reducing training costs and computational resources significantly.

Contribution

The work presents novel small protein language models based on Llama-3-8B and Phi-3-mini, utilizing LoRA to enable controllable protein generation with reduced training time and energy consumption.

Findings

Achieved an average pLDDT score of 69.75 for protein structure generation.

Attained a TM-Score of 0.84 for controllable protein generation.

Reduced training time by 70% and model size by up to 60%.

Abstract

Large language models (LLMs) have demonstrated significant success in natural language processing (NLP) tasks and have shown promising results in other domains such as protein sequence generation. However, there remain salient differences between LLMs used for NLP, which effectively handle multiple tasks and are available in small sizes, and protein language models that are often specialized for specific tasks and only exist in larger sizes. In this work, we introduce two small protein language models, based on Llama-3-8B and Phi-3-mini, that are capable of both uncontrollable and controllable protein generation. For the uncontrollable generation task, our best model achieves an average pLDDT score of 69.75, demonstrating robust performance in generating viable protein structures. For the controllable generation task, in which the model generates proteins according to properties specified in the prompt, we achieve a remarkable average TM-Score of 0.84, indicating high structural similarity to target proteins. We chose 10 properties, including six classes of enzymes, to extend the capabilities of prior protein language models. Our approach utilizes the Low-Rank Adaptor (LoRA) technique, reducing trainable parameters to just 4% of the original model size, lowering computational requirements. By using a subset of the UniRef50 dataset and small models, we reduced the overall training time by 70% without compromising performance. Notably, Phi-3-mini reduced trainable parameters by 60%, decreasing training cost by 30% compared to Llama 3. Consequently, Phi-3 achieved a comparable TM-Score of 0.81, demonstrating that smaller models can match the performance of larger ones, like Llama 3. We also demonstrate the deployment of our models on the energy efficient ET-SoC-1 chip, significantly improving the TPS/W by a factor of 3.