Generating Highly Designable Proteins with Geometric Algebra Flow Matching

Summary

This paper tackles de-novo protein design, focusing on structure generation, where the goal is to generate novel protein backbones. The paper builds on top of FrameFlow, a well-established and popular method that uses flow matching, a residue frame-based protein backbone representation, and equivariant neural networks (invariant point attention and transformers). The proposed method, however, represents the oriented residue frames as elements of a projective geometric algebra. Different orientations and translations of the residue frames are essentially represented through reflections along a set of planes. The advantage of the approach is that it allows to use more complex operations on this algebra and design higher order messages in the message passing of the equivariant network. This can improve expressivity and performance, as the paper shows. The experiments are carried out on standard, relatively small scale, protein generation benchmarks, where the method demonstrates overall slightly favourable performance compared to baselines.

Strengths

**Clarity:** While the math underlying the paper's main contribution is not easy, the high level intuitions are explained well and the paper is overall written clearly and presented well. **Originality:** The paper essentially extends the Invariant Point Attention (IPA) layers used in previous works with ideas from geometric projective geometry. Such geometric projective geometry methods have been used before, but the integration into IPA and application for protein design seems to be new and original. **Quality:** The paper is of overall good quality. The paper is well written, the experiments and baseline comparisons carried out as well as the ablations and analyses seem to be appropriate. **Significance:** I think bringing novel ideas regarding new architectures and representations into the protein design field is timely. However, the results are only very incrementally better than the baselines and the scalability of these architectures is an open question (more below).

Weaknesses

Methodologically, the paper seems interesting and novel, but the results are somewhat weak. On the important designability metric, the proposed geometric algebra flow matching method is only incrementally better than the most relevant FrameFlow baseline. This becomes clear especially in Table 5 in the Appendix, where the authors re-trained FrameFlow. On diversity and novelty, the method is not better than any baselines, sometimes actually slightly worse. Hence, the claimed performance boosts seem to be barely significant. Moreover, the trained models are small and the generated protein backbones short. Whether the proposed approach is scalable remains an open question. The authors said they kept the number of parameters of the model similar to the FrameFlow baseline, but more details here would be helpful, i.e. which components of the model have how many parameters. Furthermore, it would also be nice to report other properties of the proposed model. Does the proposed architecture converge quicker or slower? How does it affect memory consumption (which affects scalability)? **Conclusions:** For these reasons, it is overall unclear how practically impactful the proposed method can be. These issues negatively affect the paper's significance. That said, despite these somewhat weak results, I think this is overall a well-written and interesting paper with insightful analyzes, and it proposes novel and original ideas. Hence, I am still leaning towards suggesting acceptance.

Questions

The authors do not in detail write down their basic training hyperparameters, like learning rates, batch sizes, etc. I assume these are all taken exactly from FrameFlow? It would be great if the authors could clarify this.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations and potential negative societal impact have been briefly, but appropriately discussed. I have no concerns.

Summary

This paper introduces a new architecture based on Geometric Algebra Transformer (GAT) for protein backbone design. By using adapting the protein design IPA architecture with GAT, they train a SE(3) flow matching method to generate protein backbone. The contributions of the papers are essentially on the architecture. They are able to generate designable and diverse (in terms of secondary structure) proteins. They evaluated their proposed architecture on the SCOPe dataset and against standard baseline.

Strengths

First of all, the paper is clearly written, clearly presented and easy to read. I really like the idea to revisit the IPA architecture with more recent transformer architectures that were developed for SE(3) data.

Weaknesses

The weaknesses I see are ont the performed experiments and ablation study. i) I think the authors should have run experiments on the PDB dataset because it makes the comparison between their method and other methods trained on the PDB datasets incomparable (Table 2). As some models were trained on longer proteins, their performance might not be optimal compared to GAFL. ii) I would like to see the model trained on longer proteins. Currently, the largest protein size is 128 which is very small. Some of the competitors were trained on proteins of much larger size. For instance, FrameDiff and FoldFlow were trained on size up to 384. I wonder if the proposed architecture would still have been performing on longer proteins especially on the generated secondary structure. Maybe it would be harder for it and that would be a valuable input about the limits of the current method. If the authors retrain their method on the PDB dataset, I would encourage them to increase the protein length up to 384. iii) I would like to see in the ablation study (Table 3) the use of different body order. The body order was set up at '3' and I would like to see the performance of when it is at least '2'. iv) Some important papers should be discussed. Recently [Mao et al.,] introduced a new architecture to improve IPA with frame-diff. The comparison of GAFL with this method would strengthen the proposed manuscript. [Mao et al] De novo Protein Design Using Geometric Vector Field Networks, ICLR 2024.

Questions

See weakness section. While the paper proposes a very interesting idea, I believe some additional experiments, comparisons and ablations are needed for it to pass the acceptance bar at NeurIPS. I am willing to improve my score if the reviewers are able to answer the limits I see in the current manuscript.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes.

Summary

This paper extends frame-based protein backbone generation with projective geometric algebra. It allows higher-order geometric tensors in frame modeling as seen in EGNNs. Based on FrameFlow, it demonstrates great designability with relatively small increases in computational consumption.

Strengths

- Although not the first work to introduce geometric algebra to proteins or molecules, it is the first to leverage its advantages in protein backbone generation, achieving competitive results. - Introducing concepts from EGNNs like MACE is innovative and effective. - The inclusion of DSSP distribution comparisons between different methods is informative.

Weaknesses

- It is unclear whether this Geometric Algebra Extension maintains SE(3) invariance/equivariance and permutation invariance. - In Tables 1 and 2, what is the length stride used for generation? - It would be interesting to see GAFL's performance when trained on the PDB dataset and its comparison with RFDiffusion. - GAFL appears to be more complex than the original IPA. Information on its training consumption and runtime is missing. Additionally, it is important to know if GAFL is sensitive to the timesteps used for generation.

Questions

It’s coherent to see GAFL has better designability and less diversity/novelty since EGNNs are well-suited to model the structure’s geometric nature. I would like to see how this architecture performs on protein function prediction or structure prediction.

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are discussed in the paper.

Summary

This work focuses on the protein backbone generation task and improves FrameFlow, a flow-based protein backbone generation framework, based on geometric algebra. More specifically, this work proposes to represent residue frames with the elements of projective geometric algebra, which allows for the usage of higher-order message passing based on bilinear geometric products. The results in terms of designability, diversity, and the distribution of secondary structure distribution are all good compared with various baselines.

Strengths

- The proposed model architecture is effective according to the experimental results. It is challenging to achieve both high designability and proper distribution of secondary structures, and this work has well addressed the problem. - The presentation of the paper is good. The algorithms between the original version of Framediff and the proposed version show the changes clearly.

Weaknesses

- The novelty of this work is limited. It seems that the application of geometric-algebra-based models on the protein backbone design task. The author might show the novelty of the architecture compared with the existing ones, though they may still not be used in this task. - Lack of theoretical analysis or understanding of the proposed architecture. It would be more solid if the author could theoretically show the superiority of the proposed architecture compared with the original IPA. - Lack of experiments. There are many other settings that are more useful in practical scenarios beyond unconditional backbone generation, such as motif scaffolding. The authors can refer to [1] for more details. - It would be better if the author could provide a more extensive background of geometric algebra considering the audience of this paper may hold a background in machine learning for protein engineering. This would also make this paper more self-containing. References: [1] Huguet, G., Vuckovic, J., Fatras, K., Thibodeau-Laufer, E., Lemos, P., Islam, R., Liu, C.H., Rector-Brooks, J., Akhound-Sadegh, T., Bronstein, M. and Tong, A., 2024. Sequence-Augmented SE (3)-Flow Matching For Conditional Protein Backbone Generation. arXiv preprint arXiv:2405.20313.

Questions

See weaknesses. Besides, it seems GAFL is trained on the smaller dataset (SCOPe-128). I wonder its performance if it is trained on the same dataset (the larger one) used in FrameDiff.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have discussed the limitations.

Clarification on FoldFlow values in Table 2 of the original submission

Note: The values reported in Table 2 in our original submission contained values for FoldFlow that were obtained without inference annealing. If evaluated with inference annealing, the backbones sampled by FoldFlow-OT are highly designable (99.7\%) but the model underperforms on all other metrics: The strand content is below 0.5\%, which we interpret as mode collapse towards sampling mostly helical (86\%) structures, which is also reflected in poor diversity (0.49) and novelty (0.82) scores. We will include these values in the final submission.

Thank you for the rebuttal

I have read the rebuttal. I thank the authors for their answer and I hope they will add the novel experiments and discussions to VFN to their manuscript. I have raised my score.

Reply to Rebuttal

I would like to acknowledge that I have read the authors' rebuttal. I appreciate the extensive discussion as well as the additional results, in particular the experiments on the PDB. One more thing I would suggest, though, is to also re-train a standard FrameFlow model on the PDB data for the final version and include the results. FrameFlow is the most relevant and comparable baseline here and I believe FrameFlow will perform very well on the PDB, too. Although I would guess that GAFL would still be slightly better. I have raised my score and suggest acceptance of the work.