[Paper Review] ProLLaMA: A Protein Large Language Model for Multi-Task Protein Language Processing

Liuzhenghao Lv, Zongying Lin+6 more|arXiv (Cornell University)|Feb 26, 2024

Machine Learning in Bioinformatics12 citations

TL;DR

ProLLaMA transforms a general LLM into a multi-task ProLLM via a two-stage training framework using LoRA: continual learning for protein language and instruction tuning for PLP tasks, enabling unconditional generation, controllable design, and protein property prediction.

ABSTRACT

Recent advances in Protein Language Models (PLMs) have transformed protein engineering, yet unlike their counterparts in Natural Language Processing (NLP), current PLMs exhibit a fundamental limitation: they excel in either Protein Language Understanding (PLU) or Protein Language Generation (PLG), but rarely both. This fragmentation hinders progress in protein engineering. To bridge this gap, we introduce ProLLaMA, a multitask protein language model enhanced by the Evolutionary Protein Generation Framework (EPGF). We construct a comprehensive instruction dataset containing approximately 13 million samples with over 11,000 superfamily annotations to facilitate better modeling of sequence-function landscapes. We leverage a two-stage training approach to develop ProLLaMA, a multitask LLM with protein domain expertise. Our EPGF addresses the mismatch between statistic language modeling and biological constraints through three innovations: a multi-dimensional interpretable scorer, hierarchical efficient decoding, and a probabilistic-biophysical joint selection mechanism. Extensive experiments demonstrate that ProLLaMA excels in both unconditional and controllable protein generation tasks, achieving superior structural quality metrics compared to existing PLMs. Additionally, ProLLaMA demonstrates strong understanding capabilities with a 67.1% exact match rate in superfamily prediction. EPGF significantly enhances the biological viability of generated sequences, as evidenced by improved biophysical scores (+4.3%) and structural metrics (+14.5%). The project is available at https://github.com/PKU-YuanGroup/ProLLaMA.

Motivation & Objective

• Motivate the need for multi-task ProLLMs in Protein Language Processing (PLP) beyond single-task protein generation.
• Develop a universal training framework that preserves natural language capabilities while acquiring protein-language understanding.
• Enable multi-task PLP through instruction tuning to support generation, controllable design, and property prediction.
• Demonstrate scalability to additional PLP tasks with minimal training cost.

Proposed method

• Leverage a pre-trained LLaMA2 as the base model and perform continual learning on UniRef50 to acquire protein language while freezing original parameters and training LoRA adapters.
• Insert LoRA adapters into multiple decoder weights and include both Embed and Generation Head layers to preserve distinct protein-language and natural-language representations.
• Stage 1: protein-language continual learning to form ProLLaMA; Stage 2: instruction-tuning on a multi-task PLP dataset to enable instruction following and multi-task PLP abilities.
• Use a two-stage objective: CLM loss for protein language and combined CLM with instructions for multi-task PLP (Equation 1 and Equation 5 in the paper).
• Adopt low-rank adapters with relatively high LoRA rank for effective protein-language learning, and lower rank during instruction tuning to reduce training cost and avoid forgetting.
• Extend ProLLaMA to new tasks by additional instruction tuning with minimal resource use due to LoRA.

Experimental results

Research questions

• RQ1 Can a general LLM be transformed into a multi-task ProLLM capable of handling diverse PLP tasks?
• RQ2 Does two-stage training (continual protein-language learning + instruction tuning) preserve natural-language abilities while acquiring PLP skills?
• RQ3 How does ProLLaMA perform on unconditional protein generation, controllable protein generation, and protein-property prediction compared with existing ProLLMs?
• RQ4 Is the framework scalable to additional PLP tasks with low training overhead?

Key findings

• ProLLaMA achieves state-of-the-art results in unconditional protein sequence generation (pLDDT, TM-score, RMSD).
• In controllable protein generation, ProLLaMA designs novel proteins with functionalities matching given instructions (e.g., SAM-MT, Trx).
• In protein-property prediction, ProLLaMA attains average ~72% accuracy with near-100% accuracy in many superfamilies.
• ProLLaMA demonstrates strong natural-language abilities, albeit slightly below LLaMA2, indicating preserved NL capabilities after protein-language learning.
• The framework enables easy extension to more PLP tasks with low additional cost due to LoRA, and shows robust long sequences generation (stable pLDDT and TM-scores for longer proteins).