エネルギー効率の高いタンパク質言語モデル：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用に使用されるLLMsと、複数のタスクを効果的に処理し、小規模で利用可能なLLMsとは異なり、しばしば特定のタスクに特化したタンパク質言語モデルと、より大きなサイズでのみ存在するという顕著な違いが残っています。本研究では、Llama-3-8BとPhi-3-miniに基づく2つの小規模タンパク質言語モデルを紹介し、制御不能および制御可能なタンパク質生成の両方が可能です。制御不能な生成タスクでは、最良のモデルが平均pLDDTスコア69.75を達成し、実用的なタンパク質構造を生成する堅牢なパフォーマンスを示しています。プロンプトで指定された特性に従ってタンパク質を生成する制御可能な生成タスクでは、驚異的な平均TMスコア0.84を達成し、目標タンパク質との高い構造類似性を示しています。私たちは、6つの酵素クラスを含む10の特性を選択し、従来のタンパク質言語モデルの機能を拡張しました。私たちのアプローチは、Low-Rank Adaptor（LoRA）技術を活用し、トレーニング可能なパラメータを元のモデルサイズのわずか4％に削減し、計算要件を低減しています。UniRef50データセットのサブセットと小規模モデルを使用することで、全体のトレーニング時間を70％削減し、パフォーマンスを損なうことなく実現しました。Phi-3-miniは、Llama 3と比較してトレーニングコストを30％削減し、トレーニング可能なパラメータを60％削減したため、0.81のTMスコアを達成し、小さなモデルでも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.