Paper: SE(3)-Stochastic Flow Matching for Protein Backbone Generation (ICLR-2024)

Reviewer comment (Experiments):
Ablation studies could be improved to provide a full picture on the proposed techniques/tricks (see Q2).

Author response:
We thank the reviewer for their helpful comments and positive feedback.
We have now included all ablation studies, including those with auxiliary losses, in the updated Table 6 of Appendix J. Overall, we observe the following trends: 

 - Stochasticity improves robustness and the ability of the model to generate novel proteins in $7/8$ settings, as observed in the 
    fraction of novel and designable proteins (Novelty-fraction).
 - Optimal transport improves the designability of the model by reducing the variance in the training objective in $7/8$ settings.
 - The auxiliary losses improve the designability of the models in $7/8$ settings.
 - Inference annealing improves the performance of all $8/8$ FoldFlow models in all metrics.

We believe that these results significantly strengthen our empirical claims and thank the reviewer for their suggestion.

Reviewer comment (Experiments):
Is it possible to analyze how the sampling steps affect the quality of the generated samples? Presenting such results would give the readers better insight how the generation process works for this flow-matching based model, especially compared with diffusions (e.g., FrameDiff).

Author response:
This is a great question. We first note that FoldFlow-Base and FoldFlow-OT learn deterministic flows which are ODEs while FoldFlow-SFM and FrameDiff are stochastic and thus correspond to SDEs. However, the SDEs learned during training differ between FoldFlow-SFM and FrameDiff. In a nutshell, FrameDiff centers the diffusion process at the origin and adds time-scaled Brownian noise while FoldFlow-SFM constructs a path between two endpoints where the mean of the IGSO(3) distribution is the McCann (optimal transport) interpolant. Thus, during (sampling) inference, the learned vector field between FoldFlow-SFM and FrameDiff is different and is one of the reasons for the differences in their performance.

Furthermore, when visually inspecting the frames as we move in time during inference we noticed that for large parts of the trajectory, the generated protein lacked a coherent structure. Furthermore, noticeable protein structures—e.g. Helices, sheets—appear near the end of the trajectory. Such visual inspection motivated us to consider our inference annealing trick of $i(t) = ct$ which has the intended effect of speeding up the inference path at the beginning—i.e. when we are starting from noise—and slowing (because of the $t$ factor) near the end of inference. The impact of slowing down allows the model to “dwell” longer near the exact data distribution which helps create higher fidelity samples. We also note that inference annealing helps control the norm on the rotation flow such that it does not explode near the end of inference. 

We thank the reviewer for their interesting idea of analyzing the performance of the models at different sampling steps. We are currently running experiments to obtain quantitative results on the performance and will update the manuscript with the results as soon as they become available.

Reviewer comment (Experiments):
For the ablation study, it would be great to have the ablate the four parts (Stochas., OT, Aux loss, inf. annealing) on top of the full version, besides the incremental approach adopted in this paper. That is, it would be interesting to see how the performance would be if we remove one of the four parts respectively from the full version with all four techniques.

Author response:
We appreciate the reviewer's comment regarding the value of including ablation studies with all the combinations of the 4 features of the model (Stochasticity, OT, Aux. loss, Inf. Annealing). We politely refer the reviewer to our global response regarding the additional ablation experiments we are now including. In summary, we observe that, as expected, OT improves designability, Stochasticity improves robustness and generalization, and inference annealing improves the overall performance of all of the models. The results of our ablation experiments are included in Table 5 in the appendix. We are also now re-training our models without the auxiliary loss and will include those results in the final manuscript as well. 

We thank the reviewer again for their suggestions that helped us strengthen our ablation study. We hope that our rebuttal responses fully addressed all the important questions raised by the reviewer and we kindly ask the reviewer to potentially upgrade their score if the reviewer is satisfied with our responses. We are also more than happy to answer any further questions that arise.

Reviewer comment (Experiments):
The paper’s main empirical claim is that using flow matching instead of diffusion results in better generated protein structure backbones, but it does not sufficiently explain why this improvement occurs or demonstrate, with as many other parameters held constant as possible, that the flow matching objective is responsible for the improvements in generation quality.

Author response:
We thank the reviewer for their interesting question regarding the motivation behind using flow matching. We have extended our Appendix H to further discuss the theoretical and practical advantages of flow matching over diffusion and have also further motivated the OT and stochastic approaches in sections 3.2 and 3.3. We provide a summary of these points below:
- Flow matching-based approaches enjoy the property of transporting any source distribution to any target distribution. This is in contrast to diffusion where one typically needs a Gaussian-like source distribution.
- Flow matching approaches are readily compatible with optimal transport due to the same property of being able to transport and source to any target. Optimal transport which itself has the advantage of providing faster training with a lower variance training objective and reducing the numerical error in inference due to straighter paths. Diffusion models by themselves are not amenable to optimal transport but instead one can do entropic regularized OT. In Euclidean space, this corresponds to a Schrodinger bridge but this is not an optimal transport path as it is stochastic. However, we note that SDEs can leverage OT formulations in the form of entropy-regularization.
- In general, simulating an ODE is much more efficient than simulating an SDE during inference. Conditional flow-matching and OT-conditional flow matching [3, 4] both learn ODEs as the learned flow corresponds to a continuous normalizing flow. Diffusion models on the other hand are SDEs and while being more robust to noise in higher dimensions require more challenging inference.
In addition to the points above, an advantage of all of our models over the diffusion approach is that it does not rely on the costly calculation of the pdf of the IGSO(3) distribution, which FrameDiff requires for sampling and computing the score. FoldFlow simply needs to sample from the IGSO(3) which is fast to do.

Reviewer comment (Theory):
While the theoretical contribution lies in enabling stochastic flow matching for protein backbone generation, there is no theoretical or qualitative discussion of why flow matching should result in more efficient training and higher-quality generations.

Author response:
We thank the reviewer for their interesting question regarding the motivation behind using flow matching. We have extended our Appendix H to further discuss the theoretical and practical advantages of flow matching over diffusion and have also further motivated the OT and stochastic approaches in sections 3.2 and 3.3. We provide a summary of these points below:
- Flow matching-based approaches enjoy the property of transporting any source distribution to any target distribution. This is in contrast to diffusion where one typically needs a Gaussian-like source distribution.
- Flow matching approaches are readily compatible with optimal transport due to the same property of being able to transport and source to any target. Optimal transport which itself has the advantage of providing faster training with a lower variance training objective and reducing the numerical error in inference due to straighter paths. Diffusion models by themselves are not amenable to optimal transport but instead one can do entropic regularized OT. In Euclidean space, this corresponds to a Schrodinger bridge but this is not an optimal transport path as it is stochastic. However, we note that SDEs can leverage OT formulations in the form of entropy-regularization.
- In general, simulating an ODE is much more efficient than simulating an SDE during inference. Conditional flow-matching and OT-conditional flow matching [3, 4] both learn ODEs as the learned flow corresponds to a continuous normalizing flow. Diffusion models on the other hand are SDEs and while being more robust to noise in higher dimensions require more challenging inference.
In addition to the points above, an advantage of all of our models over the diffusion approach is that it does not rely on the costly calculation of the pdf of the IGSO(3) distribution, which FrameDiff requires for sampling and computing the score. FoldFlow simply needs to sample from the IGSO(3) which is fast to do.

Reviewer comment (Writing):
The paper presents a large amount of technical background from both protein generation and Riemannian flow matching, but with the page limit, readers from outside these fields may have trouble following the more technical exposition.

Author response:
We understand the reviewer’s point regarding the amount of complex background and theory that we had to cover in this paper. We tried our best to make the information as accessible as possible, especially by including more background material in the appendix and adding references to seminal textbooks on Lie groups and Riemannian geometry. We would be more than happy to include additional details on any of the sections that the reviewer found lacking. Finally, we thank the reviewer for pointing out the potential redundancy between our “main approach” section and contribution list. We have combined and shortened these sections in the updated PDF.

Reviewer comment (Novelty):
The usefulness of generating from the equilibrium distribution over conformations after training on molecular dynamics (MD) data is unclear, as it seems potentially superfluous if MD has already been run for that particular model, unless it generalizes to new structures.

Author response:
We acknowledge the reviewer's comment regarding the utility of a generative model for simulating trajectories for proteins. Indeed, as the reviewer correctly points out ideally, the generative model should generalize to MD trajectories of new proteins, which is an active part of our research. We further note that we compare our inference against MD frames not in the training set as well. Nevertheless, considering how expensive MD simulations are, the ability to cheaply simulate such trajectories with a generative model trained on relatively few frames is still of great interest. There has been a substantial amount of work in the Botlzmann generator community on this particular problem (e.g. see [1, 6]), and reducing the cost has practical applications for computational biologists. This also dovetails with the larger goal of AI4Science which seeks to amortize the cost of running expensive simulations with AI-accelerated approximations to achieve faster inference.

Reviewer comment (Presentation):
Some important information is placed in the supplement rather than the main text, such as the fact that the same neural architecture as FrameDiff is used but with a slightly different training set, which is important for interpreting empirical results.

Author response:
We thank the reviewer for their comment. To enable the fairest comparison possible, we have now retrained FrameDiff using this new training set (which we call FrameDiff-Retrained) and have included the updated results in the main text. In section 4 under Architecture, we mention that we follow Yim et al. (2023b) in our architecture. We have further highlighted this by explicitly mentioning FrameDiff in the text.

Reviewer comment (Presentation):
The two paragraphs beginning "Our approach" and the final list of main contributions at the end of the introduction contain redundant content, and reducing such redundancy could aid clarity.

Author response:
We understand the reviewer’s point regarding the amount of complex background and theory that we had to cover in this paper. We tried our best to make the information as accessible as possible, especially by including more background material in the appendix and adding references to seminal textbooks on Lie groups and Riemannian geometry. We would be more than happy to include additional details on any of the sections that the reviewer found lacking. Finally, we thank the reviewer for pointing out the potential redundancy between our “main approach” section and contribution list. We have combined and shortened these sections in the updated PDF.

Reviewer comment (Writing):
In section G.1.1, the text “Here we find that FOLDFLOW is over 2x faster than FoldFlow-Improved per step” should likely read “Here we find that FOLDFLOW is over 2x faster than FrameDiff-Improved per step.”

Author response:
**Typo** We thank the reviewer for pointing out the typo in G.1.1. It has been fixed.

Reviewer comment (Experiments):
The claim that FoldFlow outperforms FrameDiff in particular and diffusion in general would be stronger if the training set and compute were matched between FrameDiff-improved and FoldFlow, as small differences in datasets and hyperparameters like batching weaken the comparison.

Author response:
We thank the reviewer for their comment regarding the discrepancy between FoldFlow and FrameDiff training sets. As mentioned above, we have included new results on FrameDiff trained with this new dataset (called FrameDiff-retrained). Additionally, we inherit as many hyperparameters as possible from FrameDiff as we have used their code and architecture as our base code. Given that the experimental setup of FrameDiff-retrained and the FoldFlow models resemble one another as much as possible, we believe that we have provided the fairest comparison possible.

Reviewer comment (Experiments):
The section on equilibrium conformation generation does not clearly show that FoldFlow is better than (1) diffusion from a simple prior trained on the same data or (2) flow matching from a simple prior trained on the same data.

Author response:
We thank the reviewer for their suggestion on including additional ablations for FoldFlow in our equilibrium conformation sampling experiment. We have now included 2 new baselines presented in Table 3 of the main paper. These are FoldFlow-Rand and FrameDiff, both of which use a random prior as opposed to FoldFlow which starts from the empirical dataset. We discuss these results in our global response to all reviewers. In summary, our quantitative results using the Wasserstein-2 distance, as observed in the new Table 3, strengthen our original claim as we find FoldFlow to be the best model and, equally importantly, both FoldFlow and FoldFlow-Rand outperform FrameDiff. We hope that these additional results help answer the reviewer's great question and allow the reviewer to more enthusiastically endorse our experimental findings.

Reviewer comment (Evaluation):
In Table 2, the Diversity and novelty metrics should include uncertainties, and it would also be valuable to include scRMSD in the table.

Author response:
We appreciate the reviewer’s suggestions for improving Table 2. This table now provides a more complete picture of our results, as we have included the scRMSD, an additional metric for measuring novelty as well as standard errors for novelty and designability. Regarding the scRMSD, we can see that they have a similar trend to the designability metric that we originally reported. However, we would like to note that the average scRMSD is mostly affected by a few long proteins in the dataset which is why we originally only included the one that measures the fraction of proteins with scRMSD < 2. We kindly refer the reviewer to our global response for a discussion of the additional novelty metric.

Reviewer comment (Experiments):
Performing ablations of FoldFlow-OT and FoldFlow-SFM without the auxiliary loss or inference annealing would strengthen claims that these components are universally helpful.

Author response:
We thank the reviewer for their suggestion to include ablations of FoldFlow-OT and -SFM without inference annealing and/or the auxiliary loss. We politely refer the reviewer to our global response regarding the additional ablation experiments we are now including. In summary, we observe that as expected, OT improves designability, Stochasticity improves robustness and generalization, and inference annealing improves the overall performance of all of the models. We are also now re-training our models without the auxiliary loss and will include those results in the final manuscript as well. The results of our additional ablation experiments are included in Table 5 in the appendix.

Reviewer comment (Experiments):
FoldFlow-Base and FoldFlow-SFM perform poorly on designability for longer proteins, possibly due to the deduplicated dataset containing very few longer proteins.

Author response:
We agree with the reviewer that the inferior performance of FoldFlow models on longer proteins may be due to a lack of data. We believe that trying out larger proteins is an exciting direction for future work.

Reviewer comment (Novelty):
Are there applications for the arbitrary starting distribution other than sampling from the conformational ensemble via generation?

Author response:
We thank the reviewer for their question. Other biological applications can also benefit from starting from a different source distribution. One important example is protein docking where one knows the starting point of the trajectory [5]. Thus, in this case, we can set the source distribution to be the separate unbound conformations, with the target being the true bound conformation.

Reviewer comment (Evaluation):
Is there a reason to prefer scRMSD instead of scTM, as used in some earlier work?

Author response:
We thank the reviewer for this interesting question. We chose to report the scRMSD instead of scTM since previous research [2] has shown it to be a more stringent metric and more correlated to the experimental performance of protein designs.

Reviewer comment (Novelty):
The FOLDFLOW-SFM does not have much motivation provided (besides empirical results e.g. noting that they have higher novelty).

Author response:
We value the reviewer's concern regarding the insufficient motivation for FoldFlow-SFM in comparison to FoldFlow-Base and FoldFlow-OT. We would like to point the reviewer to our global response to all reviewers in which we give a detailed motivation on the utility of FoldFlow-SFM, and note that we have also highlighted the motivation behind the stochastic approach in section 3.3 of the paper. In summary, as previous research suggests, we expect stochasticity to help with robustness, particularly in high dimensions. This expectation is validated by our updated inference and new results, which demonstrate FoldFlow-SFM’s superior performance in novelty, compared to the deterministic variants of FoldFlow.

Reviewer comment (Experiments):
The experiment on equilibrium confomer generation utilizing the informed base distribution does not compare against an uninformed base distribution. Thus, it is hard to judge exactly how helpful this is (although intuitively it seems like it would be very helpful!).

Author response:
We appreciate the reviewer’s thoughtful comments and continued participation in this discussion. To highlight the significance of our equilibrium confirmation generation sampling experiments, we have—as the reviewer rightly encouraged:

 - Added Ramachandran plots for FoldFlow, FoldFlow-Rand, and FrameDiff.
 - Added the associated KL divergence for residue 56 to the test set, all evaluated at 10k training steps.
 - Added further context to the 2-Wasserstein numbers by adding evaluations of the distance of the training set, random prior, and AlphaFold2 structures to the test set.

We observe that the FoldFlow model trained with an informed prior achieves a 2-Wasserstein of $4.379$. This is only slightly worse than that of the 2-Wasserstein of the train set to the test set, which is $4.140$, and significantly better than the other models including FoldFlow-Rand at $4.446$. We hope that this clarifies the significance of starting from an informed prior. These results are in an updated Appendix G.2 in Table 5 and Figure 12.

Analyzing the Ramachandran plot results in Figure 12 we find that FoldFlow, in comparison to FoldFlow-Rand and FrameDiff, is already starting to capture both modes of the distribution, whereas the models with an uninformed prior only focus on one of the modes. Measuring this quantitatively, we find the following KL divergences for residue $56$: $0.441 < 1.205 < 3.051$ for FoldFlow, FoldFlow-Rand, and FrameDiff, respectively. 

We thank the reviewer again for allowing us to improve our paper and we hope that these results, in conjunction with the totality of all the rebuttal updates, allow the reviewer to more enthusiastically endorse our paper and potentially upgrade their score.

Reviewer comment (Writing):
I find the term “stochastic flow” a bit confusing as typically flow refers to deterministic dynamics - could the authors please clarify?

Author response:
We appreciate the reviewer's feedback regarding the usage of the term stochastic flow for our FoldFlow-SFM model. Our naming choice is motivated by the fact that FoldFlow-SFM utilizes the (OT) geodesic path between two samples $r_0 \sim \rho_0$ and $r_1 \sim \rho_1$. This path is the OT Flow between these two points and is also the mean of the IGSO(3) distribution. By sampling a noisy point from this IGSO(3) distribution we get a stochastic path which we dub as the $\textit{stochastic flow}$. Note the flow is still the mean of the distribution and the stochasticity simply refers to the injection of IGSO(3) noise to the mean path (flow). We hope this helps clear up our name choice.

Reviewer comment (Theory):
Why should optimal transport be a good prior for protein backbone generation? Section 3.2 on FoldFlow-OT provides little explanation why this should overall improve the quality of the generated samples. Could you elaborate on that? Similarly, it would be great if you could provide an intuition for motivating the introduction of the stochastic version. The results (Table 2) also do not suggest that there is any benefit in introducing three variants.

Author response:
## Motivations for using Optimal Transport
The reviewer raises an important question regarding the importance of OT as a prior for protein backbones. There are several theoretical reasons why considering OT paths is interesting in the generative modeling setup. First note, that any coupling that builds paths between two samples $r_0 \sim \rho_0$ and $r_1 \sim \rho_1$ that minimizes the OT cost necessarily constructs a “shorter” path than a random coupling. Globally, this means the average path length of the flow between $\rho_0$ and $\rho_1$ is shorter, and more importantly, none of the paths cross each other. This is a universal property of OT paths and can be applied to any generative modeling problem—including the protein backbone generation problem we consider in this paper. 
In the specific context of flow matching using OT, existing work in Euclidean space has already demonstrated the empirical benefits of OT paths [1, 2]. Moreover, in the same prior work, OT was shown to lead to faster training with a lower variance training objective in Euclidean spaces [1,2]. This is particularly relevant for protein backbones as we model them as objects in $\text{SE(3)}$ which is the semi-direct product of $\mathbb{R}^3$ and $\text{SO(3)}$. Moreover, for proteins having shorter inference directly corresponds to higher quality and more designable generated samples. Consequently, we wish to inherit and extend the benefits of OT for flow matching in Euclidean space to flow-matching on $\text{SE(3)}$ which is the providence for our FoldFlow family of models. We hope this helps address the reviewer's concerns regarding the motivations for using OT and we hope this matches expectations given the strong performance of FoldFlow-OT on the designability of proteins.
We have further highlighted the benefits of OT in section 3.2 of the paper.

## Motivations for FoldFlow-SFM
We appreciate the reviewer's comment regarding the motivation behind the SFM model and how this gets translated into the experiments. We would like to point the reviewer to our global response in which we give a detailed motivation on the utility of FoldFlow-SFM. In summary, as previous research suggests, we expect stochasticity to help with robustness, particularly in high dimensions. This expectation is validated by our updated results, which demonstrate FoldFlow-SFM’s superior performance in both of our novelty metrics, compared to the deterministic variants of FoldFlow. These results firmly establish FoldFlow-SFM as the most novel and designable model which is a key goal in de novo protein design and AI-accelerated drug discovery at large.
We have also extended the discussion on the advantages of the stochastic approach in section 3.3 of the paper.

Reviewer comment (Evaluation):
You claim that FoldFlow models are designed to be fast, stable, and efficient training. Why are the runtimes of the baselines in Table 2 missing? Regarding Table 3, what are the time differences between -base, -OT, and -SFM? What does FoldFlow in Table 3 refer to? You claim in Appendix B that you "want to get straighter flows for faster inference and more stable training". Is this something you observe?

Author response:
## Training Speed and Runtimes 
- In Table 2, we only include iters / second comparison to models with the same architecture. We have made a note of this in the manuscript. RFDiffusion and Genie baselines both use a different architecture confounding this comparison. We also note that the training code for RFDiffusion is not publicly available rendering this comparison impossible. 
- In Table 2 we show that the three FoldFlow models have approximately the same training speed, therefore, to avoid redundancy, in Table 4 we have used the term FoldFlow as an umbrella term for all the models. We thank the reviewer for pointing out how this could be confusing and have changed it to be more clear in the updated manuscript.

Reviewer comment (Novelty):
Using flow matching methods on protein structures is not new. It would be important to cite previous work here and highlight the differences [1].

Author response:
## Citation
- We thank the reviewer for bringing this important reference to our attention and have added a citation to it. We would like to however politely point out that this paper is concerned with protein docking in Euclidean, which is a different problem than the one considered here. Nevertheless, we agree with the reviewer about the relevance of this paper and we have included it in our related work section.

Reviewer comment (Evaluation):
The current experimental validation raises concerns, particularly regarding the absence of a pattern in performance of different FoldFlow variants and the chosen reported metrics and comparisons.

Author response:
Dear Reviewer,

Thank you for your time! We hope that we have addressed all of your concerns in our response. 

We have also updated the manuscript with additional analysis for the equilibrium conformation generation (Table 5, Figure 12) as well as a complete ablation experiment (Table 6). We believe that the additional ablation results show a clear pattern of performance in different FoldFlow variants. More, specifically, we observe the following trends:

 - Stochasticity improves robustness and the ability of the model to generate novel proteins in $7/8$ settings, as observed in the 
   fraction of novel and designable proteins (Novelty-fraction).
 - Optimal transport improves the designability of the model by reducing the variance in the training objective in $7/8$ settings.
 - The auxiliary losses improve the designability of the models in $7/8$ settings.
 - Inference annealing improves the performance of all $8/8$ FoldFlow models in all metrics.

The end of the discussion period is fast approaching and we would love to still have the opportunity to engage with you if there are any lingering questions. We thank you again for your time and valuable feedback, which have helped us make our experimental results more convincing.