Energiezuinige eiwit-taalmodellen: Het benutten van kleine taalmodellen met LoRA voor controleerbare eiwitgeneratie

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

November 8, 2024

Auteurs: Aayush Shah, Shankar Jayaratnam

cs.AI

Samenvatting

Grote taalmodellen (LLM's) hebben aanzienlijk succes aangetoond in natuurlijke taalverwerking (NLP) taken en hebben veelbelovende resultaten laten zien in andere domeinen zoals de generatie van eiwitsequenties. Er blijven echter opvallende verschillen bestaan tussen LLM's die worden gebruikt voor NLP, die effectief meerdere taken aankunnen en beschikbaar zijn in kleine formaten, en eiwit-taalmodellen die vaak gespecialiseerd zijn voor specifieke taken en alleen in grotere formaten bestaan. In dit werk introduceren we twee kleine eiwit-taalmodellen, gebaseerd op Llama-3-8B en Phi-3-mini, die in staat zijn tot zowel ongecontroleerde als controleerbare eiwitgeneratie. Voor de ongecontroleerde generatietaak behaalt ons beste model een gemiddelde pLDDT-score van 69.75, wat robuuste prestaties aantoont in het genereren van levensvatbare eiwitstructuren. Voor de controleerbare generatietaak, waarbij het model eiwitten genereert volgens eigenschappen gespecificeerd in de prompt, behalen we een opmerkelijke gemiddelde TM-Score van 0.84, wat wijst op een hoge structurele gelijkenis met doeleiwitten. We hebben 10 eigenschappen gekozen, waaronder zes klassen van enzymen, om de mogelijkheden van eerdere eiwit-taalmodellen uit te breiden. Onze aanpak maakt gebruik van de Low-Rank Adaptor (LoRA) techniek, waardoor het aantal trainbare parameters wordt teruggebracht tot slechts 4% van de oorspronkelijke modelgrootte, wat de rekenvereisten verlaagt. Door gebruik te maken van een subset van de UniRef50 dataset en kleine modellen hebben we de totale trainingsduur met 70% verkort zonder prestatieverlies. Opmerkelijk genoeg heeft Phi-3-mini het aantal trainbare parameters met 60% verminderd, waardoor de trainingskosten met 30% zijn verlaagd in vergelijking met Llama 3. Als gevolg hiervan behaalde Phi-3 een vergelijkbare TM-Score van 0.81, wat aantoont dat kleinere modellen de prestaties van grotere modellen, zoals Llama 3, kunnen evenaren. We demonstreren ook de implementatie van onze modellen op de energie-efficiënte ET-SoC-1 chip, waarbij de TPS/W significant wordt verbeterd met een factor 3.

English

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.