Title: Sequence-Augmented SE(3)-Flow Matching For Conditional Protein Generation
Venue: NeurIPS 2024
Accepted: yes
Keywords: Proteins

Flow Matching

Generative Models
Abstract: Proteins are essential for almost all biological processes and derive their diverse functions from complex $3 \rm D$ structures, which are in turn determined by their amino acid sequences. 
In this paper, we exploit the rich biological inductive bias of amino acid sequences and introduce FoldFlow++, a novel sequence-conditioned $\text{SE}(3)$-equivariant flow matching model for protein structure generation. FoldFlow++ presents substantial new architectural features over the previous FoldFlow family of models including a protein large language model to encode sequence, a new multi-modal fusion trunk that combines structure and sequence representations, and a geometric transformer based decoder. To increase 
diversity and novelty of generated samples -- crucial for de-novo drug design -- we
train FoldFlow++ at scale on a new dataset 
that is an order of magnitude 
larger than PDB datasets of prior works, containing both known proteins in PDB and high-quality synthetic structures achieved through filtering. We further demonstrate the ability to align FoldFlow++ to arbitrary rewards, e.g. increasing secondary structures diversity, by introducing a Reinforced Finetuning (ReFT) objective. We empirically observe that FoldFlow++ outperforms previous state-of-the-art protein structure-based generative models, improving over RFDiffusion in terms of unconditional generation across all metrics including designability, diversity, and novelty across all protein lengths, as well as exhibiting generalization on the task of equilibrium conformation sampling. Finally, we demonstrate that a fine-tuned FoldFlow++ makes progress on challenging conditional design tasks such as designing scaffolds for the VHH nanobody.


Peer reviews:
comments: SUMMARY:
The authors proposed FOLDFLOW++, which is built on top of FOLDFLOW [ICLR 2024]. It adds a joint structure and sequence representation and a transformer-based geometric decoder, enabling folding and inpainting applications.

STRENGTHS:
The tasks the authors are attempting to solve seem very interesting and important for drug discovery.

WEAKNESSES:
Considering that the theoretical novelty of the paper is somewhat limited and its main contribution lies in introducing certain architectures, the experiments conducted for the new tasks (other than unconditional generation) are also somewhat limited. Please note that I am not very familiar with the topic and do not know what potential experiments could be included.

QUESTIONS:
- I was wondering why it is not possible to calculate diversity and novelty for other tasks.

- It is not clear what the main component is that helps the model improve in terms of diversity and novelty compared to FOLDFLOW on the unconditional task.

- Can this be used for inverse folding as well?

- I was wondering if the inpainting task is not similar to seed optimization or editing tasks. For example, LaMBO-2 [NeurIPS 2023] or GFNSeqEditor [GenBio 2023]? If yes, why one cannot compare with those kinds of models? 

- What are the oracles to evaluate the generated samples?



--- after rebuttal ---

I have read the authors' responses to my questions. They have addressed most of my concerns, and I believe the paper has merit to be accepted at NeurIPS.

LIMITATIONS:
I think the model's performance might depend on various components such as ESM-2 and the structure encoder.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

comments: SUMMARY:
The paper proposes a new model FOLDFLOW++ for Conditional Protein Backbone Generation. It incorporates several techniques including sequence model, finetuning strategies, and high-quality synthetic structures to improve its performance on various tasks. The experimental results suggest the method achieves SOTA performance on various protein-related generation tasks.

STRENGTHS:
- The paper is well-written and can be used in many real-world scenarios. 

- The proposed method gains sota performance on unconditional generation, motif scaffolding, folding, fine-tuning to improve secondary structure diversity, and equilibrium conformation sampling from molecular dynamics trajectories.

- Considering the sequence embedding and ReFT is reasonable for improving the base model's performance.

WEAKNESSES:
- This paper has limited technical novelty. The core components are mostly proposed by previous works. 
- Fusing the well-trained sequence model(ESM) into the backbone generation method (FoldFlow) intuitively can improve the structure generation[1]. Therefore, we cannot see the insightful discussion and surprising conclusion from the paper.
- Some baselines concerning MD may be missing: EIGENFOLD[2], STR2STR[3],CONFDIFF[4].

[1] A Hierarchical Training Paradigm for Antibody Structure-sequence Co-design

[2]EigenFold: Generative Protein Structure Prediction with Diffusion Models

[3]Str2Str: A Score-based Framework for Zero-shot Protein Conformation Sampling

[4]Protein Conformation Generation via Force-Guided SE(3) Diffusion Models

QUESTIONS:
- Is the model framework the same as FlowFold If the sequence input is fully masked?
- Can the model compare with the AlphaFold2/3?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a protein generative model FoldFlow++ augmented with protein language model embeddings upon FoldFlow. The model is trained with sequence and structure information to learn embedding projections in SE3 space. Experiments on unconditional generation show a favorable performance of FoldFlow++ over SOTA method RFdiffusion. FoldFlow++ has the capability to be aligned to arbitrary awards like secondary structure diversity through reinforce finetuning, as well as the capability to motif scaffolding and conformation sampling.\n\nSTRENGTHS:\n**Originality**\n\nThe paper is an excellent piece of work implementing protein language model embedding-guided protein flow matching model. The design of the network is reasonable and novel. Training with half-time sequence masking introduces the capability of protein folding and design at the same time. Overall the model is carefully designed and shows wonderful protein generative modeling potential.\n\n**Quality**\n\nThe submission is technically sound, with much of the mathematical foundations explained in the previous FoldFlow paper. Various aspects of protein generative models are tested, e.g. unconditional sampling, protein folding, motif scaffolding, conformation sampling.\n\n**Clarity**\n\nThe paper is easy to comprehend and figures are well-designed and clear.\n\n**Significance**\n\nThis work integrates protein language model into a protein flow matching framework to make its protein modeling and design ability more versatile. It can perform various kinds of tasks in protein design and has a great potential as a foundational model for protein researcher.\n\nWEAKNESSES:\n1. What are the diversity of structures for motif scaffolding task? Please include some visualization of generated structures for motif scaffolding benchmarks and statistical results.\n2. Alphafold2 also has conformation sampling ability. Did you benchmark it?\n\nQUESTIONS:\n1. Are the results reported for model trained with synthetic dataset or PDB dataset?\n2. Could you explain how is the loss function implemented in the code?\n3. For protein folding and inpainting test, what kind of noise is the structure input?\n4. how is the AF2 high-confidence structures further distilled?\n\nLIMITATIONS:\nAlthough FoldFlow++ has various types of protein generation capabilities, performance on some of the tasks like conformation sampling and protein folding is not impressive, though I believe further finetuning on more specific datasets can benefit the model on this.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces FoldFlow++, a sequence-conditioned SE(3)-equivariant flow matching model designed for protein structure generation. FoldFlow++ builds upon previous FoldFlowmodels by incorporating a protein language model to encode sequences, a multi-modal fusion trunk to integrate structure and sequence representations, and a geometric transformer-based decoder. The model is trained on a large dataset of both known proteins and high-quality synthetic structures, demonstrating substantial improvements over previous state-of-the-art models in terms of designability, diversity, and novelty. FoldFlow++ also excels in conditional design tasks, such as designing scaffolds for VHH nanobodies.\n\nSTRENGTHS:\n- The paper includes a detailed ablation study examining architectural components, different flow matching schedules, and more.\n- The model surpasses previous state-of-the-art generative models in terms of designability, diversity, and the novelty of generated protein structures.\n- The authors propose a large and diverse dataset, including high-quality synthetic structures, which enhances the model's generalizability and robustness.\n- The paper explores several meaningful settings, such as Reinforced FineTuning, Motif Scaffolding, and Zero-shot Equilibrium Conformation Sampling.\n\nWEAKNESSES:\n1. The proposed pipeline appears to be a special case of Multi-Flow.\n2. The model architecture remains very similar to previous work (e.g., Genie, FrameFlow/FrameDiff), leaving it unclear which specific parts of the algorithm drive the observed improvements.\n\nQUESTIONS:\n1. In section 3.1, it’s claimed that the invariant point attention encoder (IPA) is SE(3)-equivariant. Shouldn’t it be invariant?\n2. When the sequence is chosen to be masked in training, how is the mask applied? Is it applied to the input space or the embedding space (after ESM2)?\n3. The common embedding in ESM2 is a single representation. What is the pair representation in ESM2 mentioned in Section 3.1?\n4. In Section 3.1, it’s claimed that adding a skip-connection between the encoder and decoder is essential for good performance. Is there any ablation study regarding this point?\n5. In synthetic dataset processing, how is the “masking low confidence residue” handled? Is the masking applied to both structure and sequence? Does the model leave a position for the masked residue or just ignore it?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}