에너지 효율적인 단백질 언어 모델: 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

초록

대형 언어 모델(LLMs)은 자연어 처리(NLP) 작업에서 상당한 성공을 거두었으며 단백질 서열 생성과 같은 다른 영역에서 유망한 결과를 보여주었습니다. 그러나 NLP에 사용되는 LLM과 종종 여러 작업을 효과적으로 처리하고 작은 크기로 사용 가능한 모델과 달리 종종 특정 작업에 특화된 대형 크기의 단백질 언어 모델 사이에는 현저한 차이가 남아 있습니다. 본 연구에서는 Llama-3-8B와 Phi-3-mini를 기반으로 하는 두 개의 소형 단백질 언어 모델을 소개합니다. 이 모델들은 조절 가능 및 비조절 가능한 단백질 생성이 가능합니다. 비조절 생성 작업에서 최고의 모델은 평균 pLDDT 점수인 69.75를 달성하여 실용적인 단백질 구조 생성에서 견고한 성능을 보여주었습니다. 조절 가능한 생성 작업에서는 모델이 프롬프트에서 지정된 속성에 따라 단백질을 생성하는 작업으로, 우리는 높은 구조적 유사성을 나타내는 평균 TM-Score가 0.84인 놀라운 성과를 달성했습니다. 우리는 6가지 효소 클래스를 포함한 10가지 속성을 선택하여 이전 단백질 언어 모델의 기능을 확장했습니다. 우리의 접근 방식은 Low-Rank Adaptor (LoRA) 기술을 활용하여 훈련 가능한 매개변수를 원래 모델 크기의 4%로 줄여 연산 요구 사항을 낮추었습니다. UniRef50 데이터셋의 하위 집합과 소형 모델을 사용함으로써 전반적인 훈련 시간을 70% 줄이면서 성능을 희생하지 않았습니다. 특히, Phi-3-mini는 훈련 가능한 매개변수를 60% 줄여 Llama 3에 비해 훈련 비용을 30% 절감했습니다. 결과적으로 Phi-3은 Llama 3와 같은 대형 모델의 성능을 맞출 수 있는 0.81의 TM-Score를 달성했습니다. 또한 우리의 모델을 에너지 효율적인 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.