Energieeffiziente Protein-Sprachmodelle: Nutzung kleiner Sprachmodelle mit LoRA für kontrollierbare Protein-Generierung

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

November 8, 2024

Autoren: Aayush Shah, Shankar Jayaratnam

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben signifikanten Erfolg bei natürlicher Sprachverarbeitung (NLP) gezeigt und vielversprechende Ergebnisse in anderen Bereichen wie der Generierung von Proteinsequenzen erzielt. Dennoch bestehen wesentliche Unterschiede zwischen LLMs, die für NLP verwendet werden, mehrere Aufgaben effektiv bewältigen und in kleinen Größen verfügbar sind, und Protein-Sprachmodellen, die oft für spezifische Aufgaben spezialisiert sind und nur in größeren Größen existieren. In dieser Arbeit stellen wir zwei kleine Protein-Sprachmodelle vor, basierend auf Llama-3-8B und Phi-3-mini, die sowohl unkontrollierte als auch kontrollierte Protein-Generierungsfähigkeiten besitzen. Für die unkontrollierte Generierungsaufgabe erreicht unser bestes Modell einen durchschnittlichen pLDDT-Wert von 69,75 und zeigt eine robuste Leistung bei der Generierung lebensfähiger Proteinstrukturen. Für die kontrollierte Generierungsaufgabe, bei der das Modell Proteine gemäß den in der Eingabe angegebenen Eigenschaften generiert, erreichen wir einen bemerkenswerten durchschnittlichen TM-Score von 0,84, was eine hohe strukturelle Ähnlichkeit zu Zielproteinen anzeigt. Wir wählten 10 Eigenschaften, darunter sechs Klassen von Enzymen, um die Fähigkeiten früherer Protein-Sprachmodelle zu erweitern. Unser Ansatz nutzt die Low-Rank Adaptor (LoRA) Technik, um die trainierbaren Parameter auf nur 4% der Originalmodellgröße zu reduzieren und den Rechenaufwand zu verringern. Durch die Verwendung eines Teils des UniRef50-Datensatzes und kleiner Modelle konnten wir die Gesamttrainingszeit um 70% verkürzen, ohne die Leistung zu beeinträchtigen. Insbesondere reduzierte Phi-3-mini die trainierbaren Parameter um 60%, was die Trainingskosten im Vergleich zu Llama 3 um 30% senkte. Folglich erreichte Phi-3 einen vergleichbaren TM-Score von 0,81 und zeigt, dass kleinere Modelle die Leistung größerer Modelle wie Llama 3 erreichen können. Wir zeigen auch die Implementierung unserer Modelle auf dem energieeffizienten ET-SoC-1 Chip, wodurch die TPS/W um den Faktor 3 verbessert wird.

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.