Modèles de langage protéique écoénergétiques : Exploiter de petits modèles de langage avec LoRA pour la génération contrôlable de protéines

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

Résumé

Les grands modèles de langage (LLM) ont démontré un succès significatif dans le traitement automatique du langage naturel (TALN) et ont montré des résultats prometteurs dans d'autres domaines tels que la génération de séquences protéiques. Cependant, il existe des différences saillantes entre les LLM utilisés pour le TALN, qui gèrent efficacement plusieurs tâches et sont disponibles en petites tailles, et les modèles de langage protéique qui sont souvent spécialisés pour des tâches spécifiques et n'existent qu'en tailles plus grandes. Dans ce travail, nous introduisons deux petits modèles de langage protéique, basés sur Llama-3-8B et Phi-3-mini, capables à la fois de générer des protéines de manière incontrôlée et contrôlée. Pour la tâche de génération incontrôlée, notre meilleur modèle atteint un score pLDDT moyen de 69,75, démontrant des performances robustes dans la génération de structures protéiques viables. Pour la tâche de génération contrôlée, dans laquelle le modèle génère des protéines selon les propriétés spécifiées dans la consigne, nous obtenons un remarquable score TM-Score moyen de 0,84, indiquant une grande similarité structurelle avec les protéines cibles. Nous avons choisi 10 propriétés, dont six classes d'enzymes, pour étendre les capacités des modèles de langage protéique précédents. Notre approche utilise la technique de l'Adaptateur à Faible Rang (LoRA), réduisant les paramètres entraînables à seulement 4% de la taille initiale du modèle, réduisant les exigences computationnelles. En utilisant un sous-ensemble de l'ensemble de données UniRef50 et de petits modèles, nous avons réduit le temps global d'entraînement de 70% sans compromettre les performances. Notamment, Phi-3-mini a réduit les paramètres entraînables de 60%, diminuant le coût d'entraînement de 30% par rapport à Llama 3. En conséquence, Phi-3 a obtenu un score TM-Score comparable de 0,81, démontrant que des modèles plus petits peuvent égaler les performances de modèles plus grands, comme Llama 3. Nous démontrons également le déploiement de nos modèles sur la puce écoénergétique ET-SoC-1, améliorant significativement le TPS/W d'un facteur de 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.