Modelos de Linguagem de Proteínas Eficientes em Energia: Alavancando Modelos de Linguagem Pequenos com LoRA para Geração Controlável 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

Resumo

Os grandes modelos de linguagem (LLMs) têm demonstrado um sucesso significativo em tarefas de processamento de linguagem natural (NLP) e têm mostrado resultados promissores em outros domínios, como a geração de sequências de proteínas. No entanto, existem diferenças marcantes entre os LLMs usados para NLP, que lidam efetivamente com múltiplas tarefas e estão disponíveis em tamanhos pequenos, e os modelos de linguagem de proteínas que muitas vezes são especializados para tarefas específicas e existem apenas em tamanhos maiores. Neste trabalho, apresentamos dois pequenos modelos de linguagem de proteínas, baseados em Llama-3-8B e Phi-3-mini, que são capazes de gerar proteínas tanto de forma incontrolável quanto controlável. Para a tarefa de geração incontrolável, nosso melhor modelo alcança uma pontuação média de pLDDT de 69,75, demonstrando um desempenho robusto na geração de estruturas de proteínas viáveis. Para a tarefa de geração controlável, na qual o modelo gera proteínas de acordo com propriedades especificadas no estímulo, alcançamos uma notável pontuação média de TM-Score de 0,84, indicando alta similaridade estrutural com proteínas-alvo. Escolhemos 10 propriedades, incluindo seis classes de enzimas, para ampliar as capacidades dos modelos de linguagem de proteínas anteriores. Nossa abordagem utiliza a técnica de Adaptador de Baixo Rank (LoRA), reduzindo os parâmetros treináveis para apenas 4% do tamanho original do modelo, reduzindo os requisitos computacionais. Ao usar um subconjunto do conjunto de dados UniRef50 e modelos pequenos, reduzimos o tempo total de treinamento em 70% sem comprometer o desempenho. Notavelmente, o Phi-3-mini reduziu os parâmetros treináveis em 60%, diminuindo o custo de treinamento em 30% em comparação com o Llama 3. Consequentemente, o Phi-3 alcançou um TM-Score comparável de 0,81, demonstrando que modelos menores podem igualar o desempenho dos maiores, como o Llama 3. Também demonstramos a implantação de nossos modelos no chip de baixo consumo de energia ET-SoC-1, melhorando significativamente o TPS/W em um fator 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.