Энергоэффективные языковые модели белков: использование небольших языковых моделей с LoRA для управляемой генерации белков.

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

November 8, 2024

Авторы: Aayush Shah, Shankar Jayaratnam

cs.AI

Аннотация

Большие языковые модели (LLM) продемонстрировали значительный успех в обработке естественного языка (NLP) и показали многообещающие результаты в других областях, таких как генерация последовательностей белков. Однако существуют существенные различия между LLM, используемыми для NLP, которые эффективно обрабатывают несколько задач и доступны в небольших размерах, и языковыми моделями белков, которые часто специализируются на конкретных задачах и существуют только в больших размерах. В данной работе мы представляем две небольшие языковые модели белков, основанные на Llama-3-8B и Phi-3-mini, способные к как неконтролируемой, так и контролируемой генерации белков. Для задачи неконтролируемой генерации наша лучшая модель достигает среднего показателя pLDDT в 69.75, демонстрируя надежную производительность в создании жизнеспособных структур белков. Для задачи контролируемой генерации, при которой модель создает белки в соответствии с указанными свойствами в запросе, мы достигаем выдающегося среднего показателя TM-Score в 0.84, указывающего на высокую структурную схожесть с целевыми белками. Мы выбрали 10 свойств, включая шесть классов ферментов, для расширения возможностей предыдущих языковых моделей белков. Наш подход использует технику Low-Rank Adaptor (LoRA), снижая количество обучаемых параметров до всего лишь 4% от исходного размера модели, снижая вычислительные требования. Используя подмножество набора данных UniRef50 и небольшие модели, мы сократили общее время обучения на 70% без ущерба производительности. Заметно, что Phi-3-mini сократил количество обучаемых параметров на 60%, снизив стоимость обучения на 30% по сравнению с Llama 3. Следовательно, Phi-3 достиг сопоставимого показателя TM-Score в 0.81, демонстрируя, что более маленькие модели могут соответствовать производительности более крупных, таких как Llama 3. Мы также демонстрируем развертывание наших моделей на энергоэффективном чипе ET-SoC-1, значительно улучшая TPS/W в 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.