MSAGPT: Neural Prompting Protein Structure Prediction via MSA Generative Pre-Training

Summary

The 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.

Strengths

- The paper is well-written, easy to follow, and the proposed framework is simple yet effective, and is straightforward to implement. - 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.

Weaknesses

- **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. - **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. - **Limited Use Case:** The practical use of generating virtual MSAs is limited to models that utilize MSAs, such as MSA Transformer or AlphaFold.

Questions

1. 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. 2. As mentioned in weakness, please provide a detailed comparison of the efficiency of MSAGPT with other MSA generative models. 3. How does the model perform in MSA-abundant conditions? This should also be evaluated. 4. For the “prediction accuracy” on lines 218 and 121, could you specify the metrics used (e.g., TM-score, RMSD, lDDT, pLDDT)? 5. minor: How does pLDDT selection help in finding structurally similar sequences, as mentioned on line 329? 6. minor: On line 230, do you mean the DPO dataset contains 11k samples?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Please 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.

Summary

Protein 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.

Strengths

Produces impressive performance improvements for protein structure prediction in the important regime where the query protein is far from other natural proteins. Demonstrates 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. Uses modern tricks of the trade for improving generative models using DPO, etc. The evaluation compares to multiple recent papers for generating virtual MSAs. Does a systematic investigation into how to select generated virtual sequences to include in an MSA (L296).

Weaknesses

The evaluation sets are very very small (see below). This is why I gave 'soundness' a score of 2. The advanced fine tuning techniques don't uniformly improve eval metrics (see below). ==Update after author's response== I have raised my score to 'accept' and raised my soundness score, since the response adequately addressed my concerns about evaluation.

Questions

**Tiny evaluation sets The 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. I 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. I 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? **Fine tuning decreases pLDDT I 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)? **Minor points L24 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. L113: what does 'homogenous' mean in this context? L149: 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. L180: What order do you use for flattening the axes? Have you tried both? I'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?

Rating

7

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes

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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

NA