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

TL;DR

This work introduces two small protein language models, based on Llama-3-8B and Phi-3-mini, that are capable of both uncontrollable and controllable protein generation and utilizes the Low-Rank Adaptor (LoRA) technique, reducing trainable parameters to just 4% of the original model size, lowering computational requirements.

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.

Figures (9)

• Figure 1: Training pipeline - Llama 3 and Phi 3 models were trained in two stages. First, they were trained on protein sequences and then on protein-property pairs. LoRA was used to lower the training cost.
• Figure 2: Web interface for protein sequence generation and visualization. The user enters the starting amino acids, sequence length, and property for initiating the generation. The generated sequence and its structure shows up in the corresponding sections along with the performance metrics during inference.
• Figure 3: Evaluation pipeline - The generated sequence is passed to ESMFold for predicting its structure. Foldseek aligns this structure with known proteins and gives the TM-Score and RMSD value.
• Figure 4: Structural alignment of proteins generated by Llama 3 for four different classes with their closest natural counterpart ranked according to TM-Score.
• Figure 5: Comparison of best TM-Scores obtained by Llama 3 and Phi 3 on four properties