Modelli di Linguaggio Proteico Energeticamente Efficienti: Sfruttare Modelli di Linguaggio Ridotti con LoRA per la Generazione Controllabile di Proteine

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

November 8, 2024

Autori: Aayush Shah, Shankar Jayaratnam

cs.AI

Abstract

I grandi modelli linguistici (LLM) hanno dimostrato un notevole successo nel trattamento del linguaggio naturale (NLP) e hanno mostrato risultati promettenti in altri ambiti come la generazione di sequenze proteiche. Tuttavia, rimangono differenze significative tra i LLM utilizzati per l'NLP, che gestiscono efficacemente molteplici compiti e sono disponibili in dimensioni ridotte, e i modelli linguistici proteici che sono spesso specializzati per compiti specifici e esistono solo in dimensioni più grandi. In questo lavoro, presentiamo due piccoli modelli linguistici proteici, basati su Llama-3-8B e Phi-3-mini, capaci di generare sia proteine in modo incontrollato che controllato. Per il compito di generazione incontrollata, il nostro miglior modello raggiunge un punteggio medio pLDDT di 69,75, dimostrando una prestazione robusta nella generazione di strutture proteiche valide. Per il compito di generazione controllata, in cui il modello genera proteine in base alle proprietà specificate nel prompt, otteniamo un notevole TM-Score medio di 0,84, indicando un'elevata similarità strutturale con le proteine target. Abbiamo selezionato 10 proprietà, tra cui sei classi di enzimi, per ampliare le capacità dei precedenti modelli linguistici proteici. Il nostro approccio utilizza la tecnica Low-Rank Adaptor (LoRA), riducendo i parametri addestrabili al solo 4% delle dimensioni del modello originale, riducendo i requisiti computazionali. Utilizzando un sottoinsieme del dataset UniRef50 e modelli ridotti, abbiamo ridotto il tempo complessivo di addestramento del 70% senza compromettere le prestazioni. In particolare, Phi-3-mini ha ridotto i parametri addestrabili del 60%, diminuendo il costo di addestramento del 30% rispetto a Llama 3. Di conseguenza, Phi-3 ha ottenuto un TM-Score comparabile di 0,81, dimostrando che modelli più piccoli possono eguagliare le prestazioni di quelli più grandi, come Llama 3. Dimostriamo inoltre il dispiegamento dei nostri modelli sul chip ad alta efficienza energetica ET-SoC-1, migliorando significativamente il TPS/W di un fattore 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.