Modelos de Lenguaje de Proteínas Eficientes Energéticamente: Aprovechando Modelos de Lenguaje Pequeños con LoRA para la Generación Controlable de Proteínas

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

November 8, 2024

Autores: Aayush Shah, Shankar Jayaratnam

cs.AI

Resumen

Los grandes modelos de lenguaje (LLMs) han demostrado un éxito significativo en tareas de procesamiento de lenguaje natural (NLP) y han mostrado resultados prometedores en otros ámbitos como la generación de secuencias de proteínas. Sin embargo, existen diferencias destacadas entre los LLMs utilizados para NLP, que manejan eficazmente múltiples tareas y están disponibles en tamaños pequeños, y los modelos de lenguaje de proteínas que a menudo están especializados para tareas específicas y solo existen en tamaños más grandes. En este trabajo, presentamos dos pequeños modelos de lenguaje de proteínas, basados en Llama-3-8B y Phi-3-mini, que son capaces de generar tanto proteínas de forma incontrolada como controlada. Para la tarea de generación incontrolada, nuestro mejor modelo logra un puntaje pLDDT promedio de 69.75, demostrando un rendimiento sólido en la generación de estructuras de proteínas viables. Para la tarea de generación controlada, en la cual el modelo genera proteínas de acuerdo con propiedades especificadas en la indicación, logramos un notable puntaje TM-Score promedio de 0.84, indicando una alta similitud estructural con las proteínas objetivo. Elegimos 10 propiedades, incluidas seis clases de enzimas, para ampliar las capacidades de los modelos de lenguaje de proteínas anteriores. Nuestro enfoque utiliza la técnica de Adaptador de Rango Bajo (LoRA), reduciendo los parámetros entrenables a solo el 4% del tamaño original del modelo, disminuyendo los requisitos computacionales. Al utilizar un subconjunto del conjunto de datos UniRef50 y modelos pequeños, redujimos el tiempo total de entrenamiento en un 70% sin comprometer el rendimiento. Especialmente, Phi-3-mini redujo los parámetros entrenables en un 60%, disminuyendo el costo de entrenamiento en un 30% en comparación con Llama 3. En consecuencia, Phi-3 logró un puntaje TM-Score comparable de 0.81, demostrando que los modelos más pequeños pueden igualar el rendimiento de los más grandes, como Llama 3. También demostramos la implementación de nuestros modelos en el chip energéticamente eficiente ET-SoC-1, mejorando significativamente el TPS/W en un factor 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.