Title: MSAGPT: Neural Prompting Protein Structure Prediction via MSA Generative Pre-Training
Venue: NeurIPS 2024
Accepted: yes
Keywords: Computational Biology

Protein Language Model

Protein Structure Prediction

MSA Generative Pre-Training
Abstract: Multiple Sequence Alignment (MSA) plays a pivotal role in unveiling the evolutionary trajectories of protein families. The accuracy of protein structure predictions is often compromised for protein sequences that lack sufficient homologous information to construct high-quality MSA. Although various methods have been proposed to generate high-quality MSA under these conditions, they fall short in comprehensively capturing the intricate co-evolutionary patterns within MSA or require guidance from external oracle models. Here we introduce MSAGPT, a novel approach to prompt protein structure predictions via MSA generative pre-training in a low-MSA regime. MSAGPT employs a simple yet effective 2D evolutionary positional encoding scheme to model the complex evolutionary patterns. Endowed by this, the flexible 1D MSA decoding framework facilitates zero- or few-shot learning. Moreover, we demonstrate leveraging the feedback from AlphaFold2 (AF2) can further enhance the model’s capacity via Rejective Fine-tuning (RFT) and Reinforcement Learning from AF2 Feedback (RLAF). Extensive experiments confirm the efficacy of MSAGPT in generating faithful and informative MSA (up to +8.5% TM-Score on few-shot scenarios). The transfer learning also demonstrates its great potential for the wide range of tasks resorting to the quality of MSA.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes MSAGPT, a novel method for generating MSAs. Utilizing 2D evolutionary positional encoding, MSAGPT reformalizes MSA generation as a one-dimensional sequence generation task optimized with a simple GPT objective. The model incorporates feedback from AlphaFold2 to reduce hallucinations during MSA generation via DPO fine-tuning. Experimental results on curated datasets demonstrate that MSAGPT enhances protein structure prediction in low-MSA scenarios, achieving improved structural reconstruction scores.\n\nSTRENGTHS:\n- The paper is well-written, easy to follow, and the proposed framework is simple yet effective, and is straightforward to implement.\n\n- The use of 2D positional encoding to re-formalize MSA generation as a one-dimensional sequence generation task is innovative and allows for zero- or few-shot MSA generation under a flexible in-context learning framework. It points out a possible direction of playing 2D sequences with novel positional encoding.\n\nWEAKNESSES:\n- **Efficiency Concerns:** Flattening 2D sequences into 1D for interaction with self-attention increases time complexity, even with FlashAttention. While Figure 8 shows MSAGPT's generation time is lower than the AF2 search pipeline, a comparison with other MSA generative models' efficiency is necessary.\n\n- **More Comparative Analysis:** The paper should include a comparison with diffusion-based models for generating protein sequences, such as EvoDiff MSA. Also, could you also report the RMSD (GTD_TS) score.\n\n- **Limited Use Case:** The practical use of generating virtual MSAs is limited to models that utilize MSAs, such as MSA Transformer or AlphaFold.\n\nQUESTIONS:\n1. I feel like most innovations are credited to 2D RoPE positional embedding, limiting the scope of novelty. Also, could you provide detailed explanation of this method in English? The reference is in Chinese.\n\n2. As mentioned in weakness, please provide a detailed comparison of the efficiency of MSAGPT with other MSA generative models.\n    \n3. How does the model perform in MSA-abundant conditions? This should also be evaluated.\n    \n4. For the “prediction accuracy” on lines 218 and 121, could you specify the metrics used (e.g., TM-score, RMSD, lDDT, pLDDT)?\n    \n5. minor: How does pLDDT selection help in finding structurally similar sequences, as mentioned on line 329?\n    \n6. minor: On line 230, do you mean the DPO dataset contains 11k samples?\n\nLIMITATIONS:\nPlease refer to the Weaknesses and Questions sections. I hope the authors can address concerns regarding efficiency, provide more comparative analysis with baseline models, and offer further explanation on the use of 2D RoPE positional embedding.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nProtein structure prediction tools such as alpha fold take a query protein sequence, expand it to a multiple sequence alignment (MSA) of related natural sequences, and then feed this alignment into the model. The first expansion step isn't possible, however, for proteins that don't have many natural relatives, and there are many such 'orphans'. This paper demonstrates that the MSA can be replaced with a set of virtual sequences sampled from a generative model. One of the key contributions of the paper is showing how to find tune those generative model explicitly for the sake of improving the downstream protein structure prediction performance using LM preference optimization techniques such as DPO.\n\nSTRENGTHS:\nProduces impressive performance improvements for protein structure prediction in the important regime where the query protein is far from other natural proteins. \n\nDemonstrates that good performance can be achieved without modeling techniques specific to multiple sequence alignments (such as axial attention). Using a vanilla setup is nice because it enables using off-the-shelf LM systems that have improvements such as flash attention, etc.\n\nUses modern tricks of the trade for improving generative models using DPO, etc.\n\nThe evaluation compares to multiple recent papers for generating virtual MSAs.\n\nDoes a systematic investigation into how to select generated virtual sequences to include in an MSA (L296).\n\nWEAKNESSES:\nThe evaluation sets are very very small (see below). This is why I gave 'soundness' a score of 2. \n\nThe advanced fine tuning techniques don't uniformly improve eval metrics (see below).\n\n==Update after author's response==\nI have raised my score to 'accept' and raised my soundness score, since the response adequately addressed my concerns about evaluation.\n\nQUESTIONS:\n**Tiny evaluation sets\nThe evaluation metrics are reported on very tiny evaluation sets (L247; 8 from cameo, 13 from CASP14&15, 179 from PDB). This makes me worry that the differences in performance are not statistically significant, or don't reflect performance on the broad distribution of orphan proteins that users may want to perform prediction for. \n\nI see two ways forward: (1) demonstrate that the proteins in these tiny eval sets are somewhat representative and that the differences in metrics are statistically significant or (2) change the eval setup to use a bigger eval set.\n\nI think (2) is much easier. To do this, can't you turn any example in the full eval set (e.g., the PDB) into an orphan? You could do the zero-shot eval using just the original sequence, with an MSA containing only virtual sequences. With this, can you report performance on the full eval sets?\n\n**Fine tuning decreases pLDDT\n\nI understand the argument that pLDDT decreased and TM score increased because the fine tuning targeted improvements to TM score. However, it's unfortunate that pLDDT increased. If you fine tune for pLDDT, would that increase at the expense of TM? Could you extend the fine tuning to improve both (such as using a composite reward function based on both)? \n\n**Minor points\nL24 is inaccurate: \"The remarkable success of AF2 can be attributed to its innovative use of co-evolutionary information supported by the Multiple Sequence Alignment (MSA).\" Coevolution had been used for structure prediction for many years before AF. AF got better performance because of general model scale and processing the MSA end-to-end instead of using pre-computed features of the MSA.\n\nL113: what does 'homogenous' mean in this context?\n\nL149: As far as I can tell, there are no semantics to the ordering of the rows. The row positional encoding is basically a unique id for which row it belongs to. Perhaps it would be better to use a different representation for the row index that doesn't reflect linear ordering so much.\n\nL180: What order do you use for flattening the axes? Have you tried both? \n\nI'm curious if there is any noticeable difference in the sequences generated after DPO vs. the base model. Are there any sort of degeneracies due to reward hacking, etc?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper proposes a method to generate multiple sequence alignments for a given protein sequence. To model the co-evolutionary information, the paper proposes 2d evolutionary positional encoding. After pretraining on the alignment sequences, the models are fine-tuned with AlphaFold2 annotations to avoid hallucinations.

STRENGTHS:
* The studied problem, generating multiple sequence alignment is interesting and novel. 
* The proposed method is technically sound.
* The paper is well-written and well-structured.

WEAKNESSES:
* Missing important experimental details. The paper omits crucial experimental details, particularly those related to pretraining. This omission significantly impacts the study's credibility and reproducibility. Especially, it's important to explain the process of hyperparameter selection is not explained in detail.
* Evaluation is not rigorous. There is no clear explanation of steps taken to prevent data leakage or the inclusion of data similar to the evaluation set in the training data. The absence of structural or temporal splits is particularly problematic, as these are crucial for assessing a model's performance in truly novel scientific scenarios. 
* Evaluation is limited. It would be interesting to design studies to directly investigate the evolutionary patterns learned in the multiple sequence alignment algorithm. Also the paper fails to include many important protein function tasks in its evaluation.
* The results presented in the paper lack error bars, which is particularly problematic for results that are close, such as Table 2 and 3.

QUESTIONS:
See above

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13