AssembleFlow: Rigid Flow Matching with Inertial Frames for Molecular Assembly

Q: Related work [2].

Thank you for pointing out this paper by Bose [2]. We have already discussed this work (and relevant works) in the Sec A.3 in the first version and a simlar answer to reviewer 1PjM. We agree that decomposing SE(3) into SO(3) and R(3) is not novel, yet we would like to emphasize that such a decomposition on inertial frame and with LERP and SLERP is novel. We summarize the main differences below, and you can find more details in Sec A.4 in the revised version.

Besides, we further conduct an ablation study on COD-5000 below, and as observed, using SLERP is better than MELERP. Please check Sec H.2 for more details. BTW. We are still running experiments on COD-10000 and COD, and will update the results soon.

[2] SE(3)-Stochastic Flow Matching for Protein Backbone Generation, Bose et al., ICLR 2024

Q: What advantage does the eigendecomposition of the inertial tensor I have over the common approach of using centered X.

Thank you for raising this question. This is indeed one of the biggest contributions of our work: utilizing inertial frame for rigid modeling. Using $I$ enables rigid modeling in AssembleFlow while centered $X$ in baseline CrystalFlow cannot:

• Using centered $X$ cannot guarantee rigidity. This is because we are modeling the atom-level rotation from $X_0$ to $X_1$, which cannot guarantee the rigidity. The baselines CrystalFlow and its variants are simulating crystallization directly in this way, without using inertial frames, and you can see the performance gap between them.
• Using $I$ can guarantee rigidity. This is because instead of modeling all the atoms, AssembleFlow learns how the reference frame rotates from $X_0$ to $X_1$.
• To sum up, the key difference is that learning rotation on the refernece frame can guarantee rigidity.

BTW. Can you help list the papers that are using centered $X$ for handling the molecular crystallization or rigid modeling? We are happy to check in more details on these any other related works.

Thank you for your thorough review of our paper. We have carefully addressed the concerns you raised, and you can find the corresponding responses below.

First for notation, let us use $X_1$ and $X_0$ be the strongly correlated structures (stable) and weakly correlated structures (unstable), respectively.

Q: A notable limitation of this work is the claim to be the first to implement rigid generation in SE(3) space for molecular assembly.

We believe this claim is accurate, as our work is the first to implement SE(3)-based rigid generation specifically for molecular assembly, focusing on modeling dynamic processes rather than stable state distributions. While the ICML work [1] provides valuable contributions to equilibrium conformational sampling, our approach addresses a complementary aspect by capturing the dynamic processes and distributions associated with molecular assembly. We see these efforts as addressing distinct but equally important challenges in molecular modeling.

Q: The differences between our work and the ICML 2023 paper [1]

Thank you for pointing out this paper [1], we agree that this is relevant, and have added a section (Sec A.5) to discuss the difference between the paper and our work in the revised version. Meanwhile, we would like to emphasize that our work differs from the ICML 2023 paper [1] in three fundamental aspects:

1. Dynamic Transition Modeling vs. Stable Structure Sampling: Unlike [1] which focuses on stable structure sampling, our work models the dynamic transition process from weakly correlated (unstable) structures to strongly correlated (stable) structures. Namely, the objectives are different: [1] models $p(X_1)$ while AssembleFlow models $p(X_1|X_0)$ (Fig 1). Notably, dynamic transition modeling toward stability is identified as a nontrivial next step in the ICML work [1], which states 'An interesting but nontrivial next step would be ... a flow model for the positions that can handle ... phase transitions'.
2. Use of Inertial Frames for Rigid Modeling: The ICML paper [1] 'assumes that these bodies are rigid'. In contrast, we employ inertial frames to enable rigid transformations, including both translation and rotation. This approach is also suggested as a future direction in the ICML paper, which states 'exploiting the SE(3) symmetry of jointly moving all rigid bodies'. Achieving SE(3) equivariance for rigid transformations in materials is non-trivial, as it handling permutation, and applying a reparameterization technique to directly model SE(3) at time t.
3. Limitations of Normalizing Flow Models on Molecular Assembling: Normalizing flows require a Gaussian prior, which is unsuitable for our weakly correlated structures that typically do not follow a normal distribution. The ICML paper also faces challenges with interpolation in dynamic modeling, as normalizing flows do not inherently support interpolation. Normalizing flows require a tractable prior (like Gaussian), which is non-trivial for weakly correlated or unstable structures. The ICML paper also faces a similar challenge and it solves it by introducing a surrogate prior for sampling equilibrium molecules (not unstable) with MD simulation. On the contrary, our proposed AssembleFlow is built on the Flow Matching framework, which supports the interpolation for dynamics simulation. In AssembleFlow, the LERP and SLERP interpolation methods enable us to model such rigid transformations between two arbitrary states, including from an unstable state to a stable state.

[1] Rigid Body Flows for Sampling Molecular Crystal Structures, Köhler et al., ICML 2023

Thank you for your response. My key concerns are still unresolved. Namely,

1. Differences between [1] and the proposed work

The authors state:

Our work models the dynamic transition process from weakly correlated (unstable) structures to strongly correlated (stable) structures

However, the provided mathematical formulation $p(X_1|X_0)$ does not model the dynamic transition process. The term 'dynamic transition' typically refers to modeling the entire trajectory from $X_0$ to $X_1$, as in molecular dynamics [2, 3]. This work, however, simply adds a conditioning term $X_0$ to $p(X_1)$ without modeling the intermediate transition process.

Additionally, I have reviewed the construction of the COD-Cluster17 dataset and observed that $X_0$ is generated by randomly applying rototranslations to $X_1$, rather than simulating a trajectory (using molecular dynamics) from a randomly initialized $X_0$ to a stable state $X_1$. This raises concerns about whether this dataset is suitable for addressing the task of molecular assembly, as $X_0$ and $X_1$ may not accurately represent a true pair in the assembly process. If I have misunderstood, I kindly ask for clarification.

Normalizing flows require a Gaussian prior, which is unsuitable for our weakly correlated structures that typically do not follow a normal distribution.

This claim is inaccurate. Normalizing flows do not require a Gaussian prior, as evidenced in the cited paper [1]. While I recognize the advantages of flow matching over traditional normalizing flows, this statement raises concerns about whether the authors have fully understood the difference between their approach and [1].

Additionally, I noticed the absence of a reproducibility statement and a link to the code. Including these would be greatly helpful for better understanding the differences and ensuring the work's reproducibility.

2. Presentation of the paper.

Thank you for pointing out this paper by Bose [2]. We have already discussed this work (and relevant works) in the Sec A.3 in the first version and a simlar answer to reviewer 1PjM. We agree that decomposing SE(3) into SO(3) and R(3) is not novel, yet we would like to emphasize that such a decomposition on inertial frame and with LERP and SLERP is novel.

I am aware of the difference between this work and that of Bose. My concern lies with the presentation of the paper. While the authors acknowledge that the decomposition is not novel, statements such as "This decomposition also empowers..." and "Such decomposition allows the explicit enforcement of ..." give the impression that it is a novel approach and is their contribution. I believe the authors could improve the presentation to more clearly distinguish their contributions from existing work.

3. Advantage inertial frames over PCA.

Thank you for raising this question. This is indeed one of the biggest contributions of our work: utilizing inertial frame for rigid modeling.

Regarding my statement on centered $X$, I was asking about the advantage of inertial frames over common approaches of using PCA. I apologize for the ambiguity in my earlier question.

While I appreciate the authors' response to Reviewer 1PjM regarding the numerical instability of PCA, the experiment setup seems somewhat artificial and unclear (e.g., "Notice that since the MSE reconstruction is meaningless when it is too small (why?), so we only compare these two frames when at least one of them has reconstruction greater than or equal to a threshold θ."). Given the authors' claim that this is the biggest contribution, the significance of inertial frames would be more evident if they provided ablations comparing the performance using PCA versus inertial frames.

[1] Rigid Body Flows for Sampling Molecular Crystal Structures, Köhler et al., ICML 2023.
[2] Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics, Klein et al., NeurIPS 2023.
[3] Generative Modeling of Molecular Dynamics Trajectories, Jing et al., NeurIPS 2024.

Q. Differences between [1] and the proposed work

Thank you for the feedback. We were initially confused by [1], as it blends the preliminaries (Eqs. 2, 4, and 5) with the core objectives (Eqs. 3 and 7). By Gaussian, we meant the tractable priors in general. [1] discussed using Gaussian as priors but replaced it with samples drawn from equilibrium MD simulation.

Now after revisiting [1] for further details, we have fixed our statement in previous round and provided our updated comparison below. Additionally, we would like to emphasize that our work can be seen as an extension of [1], yet addressing more practical and challenging problems, including rigid modeling on arbitrary molecules, SE(3)-equivariant modeling, and interpolation modeling.

• [1] targets at modeling the transition, e.g., water molecules at different temperatures. The MD simulation can provide samples at the stable or equilibrium status. Our work models the transition from the unstable to stable state.
• [1] specifically models the frame for H-O-H (codes here) (no code or comments related to methane rigid modeling were found in the GitHub repository, only ice). In contrast, our approach is more generalizable, as the inertial frame can serve as a reference frame for any molecule. Thus, [1] cannot be directly applied to our dataset, but our work can be viewed as an extension of [1] to a more general setting.
• [1] estimates an upper bound of log-ratio for energy difference estimation between two equilibrium states, while ours is modeling density estimation directly for unstable to stable transition.
• [1] states it does not 'explore the SE(3) symmetry of moving all rigid bodies', and we solve this limitation by introducing the SE(3)-equivariant modeling from two granularities.
• [1] does modeling with an extra bijectivity constraint, limiting the model capacity (e.g., Flow++). The velocity function models in FM do not have this issue. The transition of a molecular cluster from weakly correlated (unstable) structures to strongly correlated (stable) structures is a special case of the general dynamics. The limitation here is the data insufficiency (lack of intermediate snapshots), not modeling.

Q. Presentation of the paper

We will improve our presentation as suggested.

Q. Advantage inertial frames over PCA

We have proven that using the inertial frame is more accurate than using PCA. This is the reconstruction experiment in reply to Reviewer 1PjM. Please feel free to check the details.

When $\delta_{\text{Inertial}}$ or $\delta_{\text{PCA}}$ gets too small, as on the order of 1e-12, it falls well below the sensitivity of the dataset, rendering the comparison meaningless. Therefore, we consider the threshold from 1e-3 to 1e-8 for clearer and more meaningful comparison.

Thank you! We are happy to share additional insights.

Q: Why COD-Cluster17.

Thank you for raising this question. We actually asked a similar question to the authors of the COD-Cluster17, and one of the authors, a crystallographer, confirmed that this geometrically optimized dataset closely aligns with what crystallographers would typically use in practice, for their wet-lab experiments.

Additionally, we would like to share a critical comment from authors of COD-Cluster17, which is related to your question, is: "What constitutes true packing dynamcis? Because we think the only true dynamics is what is happenging in reallity, and we are only able to capture that by a super-powerful electron microscope, which is not yet practical." They also stated, "From this aspect, we do not consider MD simulation to represent true dynamcis." This is the general viewpoint shared by them (especially from a wet-lab crystallographer), and we believe it holds some merit. Feel free to reach out to the authors for further insights if needed.

On the other hand, our algorithm is designed to generalize to other datasets with diverse molecular packing dynamics. An ideal dataset would include more snapshots along packing trajectories and encompass a wider variety of molecular systems. We hope our work can inspire more people into exploring this problem and contribute to the development of additional datasets.

BTW. If you have any recommendations for alternative datasets, we would greatly appreciate your input.

Q: Code release.

Yes, we will release the code to support open science.

Q: Why experiment on stability of frame construction.

The experiment we designed and conducted aims to measure the stability of using an inertial frame and PCA to construct the reference frame, representing the most straightforward approach for comparison. While we could have conducted complete experiments based solely on PCA, we didn't as the stability comparison is sufficient to demonstrate that the inertial frame is more stable than PCA.

Now that we have a few more days during the rebuttal period, we could try our best to add the complete experiment if you insist. However, before proceeding, we would like to double-check whether you feel the stability comparison alone is insufficient.

Thank you for recognizing our work as both theoretically sound and thoroughly evaluated through comprehensive comparisons. We also appreciate your insightful question regarding the DFT energy. We acknowledge this as a potential issue and have addressed it in the latest revision. Please find the details below.

Q: AssembelFlow does not outperform PackMol in terms of total energy.

We have provided the comparison of the two systems in terms of formation energy, which is a more direct indicator of a material's energetically favorable state. As shown in the revised Figure 4, AssembleFlow outperforms PackMol. More detail please see comments To All Reviewers.

Q: Could you plot the energy differences between PackMol and AssembleFlow for the same structures?

Certainly, we have provided the pairwise differences between the two systems in terms of formation energy in Sec H.3 in the revised version.

Typos

Thank you for pointing out the typos, we have fixed them in the revision.

Q: About periodicity.

Thank you for asking this question. The periodic structure will show up in the final structure (strongly-correlated and stable), while it won't show up in the initial structures (weakly-correlated and unstable). The molecular crystallization task we are focusing on now is modeling the process from weakly-correlated structures to strongly-correlated structures, i.e., from non-periodic clusters to periodic clusters.

Q: Rigid assumption.

Yes, we agree that rigidity is a fundamental assumption in our task, and we have emphasized this at the beginning of Section 3 in the revised version. Additionally, we would like to note that this rigid assumption is widely adopted in similar contexts, such as backbone structure generation tasks in protein engineering (e.g., AlphaFold2, FoldFlow, FrameFlow, etc.). In other words, we want to highlight that AssembleFlow for rigid modeling represents just a primary step, with the potential to evolve into a more powerful tool, similar to its applications in protein design.

Q: Why energy is higher?

Thank you for pointing out this question. We acknowledge that using DFT energy is not that accurate, and we have replaced it with the formation energy.

DFT total energy reflects the absolute energy of a material in a specific configuration. This measure does not necessarily indicate whether the material is in a more favorable state compared to other compounds. To address this issue, we have compared the two systems in terms of formation energy, a direct indicator of the relative favorability of a material's state. As shown in the revised Figure 4, AssembleFlow provides a more favorable energy state over PackMol. For additional details, please refer to the comments under To All Reviewers.

Q: Comparison between two versions.

Sure, they are the last two columns in Table 2, AssembleFlow-Atom and AssembleFlow-Molecule. The architecture details are listed in Sec E.2.

Q: Initial conformation.

The initial conformations were generated in the COD-Cluster17 dataset. It's not generated using RDKit, but using a geometric optimization method. So yes, in our paper, we just take both the initial and final conformations from their data. We have highlighted this at the beginning of Sec 3 in the revised version.

Thank you for recognizing the novelty of using the inertial frame as a reference, the meaningful metrics, the significant improvements over deep learning baselines, and the efficiency compared to numerical methods. We also appreciate your detailed questions, particularly those regarding the protein rigid modeling comparison. We have provided the responses below.

Q: Comparison of two interpolation methods.

Thank you for raising this critical question (we have a similar answer to reviewer 1PjM). These relevant papers (FoldFlow and FrameFlow) are using exponential map interpolation (EMLERP) for SO(3) interpolation.

SLERP and EMLERP are different interpolation methods on SO(3), and they do not have a clear methodological advantage over each other; however, the EMLERP considers more approximation tricks in implementation (FoldFlow and FrameFlow). We provide a detailed comparison below, and you can find more details in Sec A.4 in the revised version.

Besides, we further conduct an ablation study on COD-5000 below, and as observed, using SLERP is better than MELERP. Please check Sec H.2 for more details. BTW. We are still running experiments on COD-10000 and COD, and will update the results soon.

Q: More reviews on SE(3) decomposition.

Thank you for pointing this out. We had discussed this in Section A.3 of the initial version, but we appreciate your suggestion and are happy to include additional relevant papers if you can provide references. We have made the following updates in the revised version:

• We added more relevant works in Sec A.3.
• There are two core differences between our work and the existing works:
  1. The first difference is the reference frame construction, as we already discussed in Sec A.3.
  2. The second difference is the SO(3) interpolation (which we explained in the reply above), and we listed the details in Sec A.4 and H.2 in the revised version.

Q: Numerical stability.

Thank you for raising this question. This is indeed an important implementation question.

• We calculate a reconstruction error when calculating the inertial frame. More concretely, suppose we have weakly-correlated structures $X_0$ and strongly-correlated structures $X_1$, and we find the corresponding bases using inertial frame $I$. The two bases are marked as $B_0 \in \mathbb{R}^{3 \times 3}$ and $B_1 \in \mathbb{R}^{3 \times 3}$. We can obtain the rotation matrix with $R^T = B_0^T B_1$, and then we can rotate the whole molecular system as $\tilde X_1 = X_0 R^T$, and we are measuring the reconstruction errors as MSE$(X_1 - \tilde X_1)$, marked as $\delta$. If $\delta$ is bigger than a threshold, e.g., 1e-3, then we just ignore this molecule.
• We have listed relevant details in Sec C.4 in the revised version.
• When there is a tie, we just do a depth-first-search to find an ordering such that the reconstruction error $\delta$ is the smallest.

Q: Questions on sampling, dataset, and problem formulation.

Thank you for raising these questions.

• It is a generative model, and since we are using the LERP and SLERP, they are two deterministic paths. The baseline CrystalFlow-LERP is also desterministic.
• The other baselines (CrystalSDE-VE/VP, CrystalFlow-VE/VP) are using the diffusion path, which is random.
• For fair comparison, all the methods sample only once.
• The input weakly correalted structures are fixed. (More details can be found in the original paper) We highlighted this in the Problem Formulation at the beginning of Sec 3 in the revised version.

Q: Different energy distribution.

Thank you for raising this question. Previsouly we were reporting the DFT energy, which is problemantic, and now we fix this by replacing it with the formation energy, and here are the detailed explanation.

DFT total energy reflects the absolute energy of a material in a specific configuration. This measure does not necessarily indicate whether the material is in a more favorable state compared to other compounds. To address this issue, we have compared the two systems in terms of formation energy, a direct indicator of the relative favorability of a material's state. As shown in the revised Figure 4, AssembleFlow provides a more favorable energy state over PackMol. For additional details, please refer to the comments under To All Reviewers. The superior formation energy of AssembleFlow compared to PackMol aligns with the packing matching metric presented in Table 2, suggesting that packing matching is an effective metric for this molecular assembly task.

Q: Questions on permutation and ordering of molecules.

In the provided COD-Cluster17 dataset, the dataset stores the molecules in a fixed order. That is to say, the ordering of molecules in the initial positions is the same as the ordering of molecules in the final positions. Thus, we do not have this permutation issue.

Q: Comparison between SLERP and exponential map interpolation. Are they mathematically equivalent?

This is an interesting question. These relevant papers (FoldFlow and FrameFlow) are using exponential map interpolation (EMLERP) for SO(3) interpolation.

A brief answer is that SLERP and EMLERP are different interpolation methods on SO(3). These two methods do not have a clear methodological advantage over each other; however, the EMLERP considers more approximation tricks in implementation (FoldFlow and FrameFlow). We provide a detailed comparison below, and you can find more details in Sec A.4 in the revised version.

Besides, we further conduct an ablation study on COD-5000 below, and as observed, using SLERP is better than MELERP. Please check Sec H.2 for more details. BTW. We are still running experiments on COD-10000 and COD, and will update the results soon.

Typos Thank you for pointing out the typos, we have fixed them in the revision.

Thank you for acknowledging our work for its significant task design, novel use of quaternions for SLERP, well-organized structure, and extensive background. We also appreciate your careful reading of our paper and for pointing out the typos. We have addressed your comments in the revision and outlined them below.

Q: Canonical pose.

Thank you for raising this interesting question. First, we would like to highlight that there are two roles and one important property of the inertial frame and its eigenvectors in our approach:

1. Role 1: The three bases in inertial frames act as a reference or a canonical pose.
2. Role 2: The three bases enable modeling the velocity function in SO3 space.
3. Property: What's more important, we expect the three bases to be numerically stable.

Then getting back to the question, though both can help build up the reference frame or canonical pose, there are certain aspects we would like to emphasize when comparing PCA and Inertial frames.

• Canonical pose: The key question is defining a canonical pose. If we just do PCA on the point clouds, this cannot guarantee the group symmetry in SE(3). However, if we remove the center point, then the group symmetry can be guaranteed; then we can use SVD to obtain the three principal components. Till this step, you can find this is somewhat similar to the inertial frame construction (Sec 3.1).
• Difference between PCA and Inertial Frame as reference frames: Though both can be used for building up the reference frame or canonical poses, SVD (for PCA) on the set of point clouds $N \times 3$ (N is the number of atoms) can be less numerically unstable, while eigendecomposition Inertial Frame $3 \times 3$ can be more numerically stable. We have conducted an additional experiment to verify this.
  • Experiment setup: suppose we have weakly-correlated structures $X_0$ and strongly-correlated structures $X_1$, and we find the corresponding bases using either eigendecomposition on centered $X$ or inertial frame $I$. The two bases are marked as $B_0 \in \mathbb{R}^{3 \times 3}$ and $B_1 \in \mathbb{R}^{3 \times 3}$.
  • Objective: We can obtain the rotation matrix with $R^T = B_0^T B_1$, and then we can rotate the whole molecular system as $\tilde X_1 = X_0 R^T$, and we are measuring the reconstruction errors as MSE$(X_1 - \tilde X_1)$. We mark the MSE using inertial frame and PCA as $\delta_{\text{Inertial}}$ and $\delta_{\text{PCA}}$, respectively. If $\delta_{\text{Inertial}} < \delta_{\text{PCA}}$, then we can conclude that using the inertial frame is more stable than using PCA, and vice versa. Notice that since the MSE reconstruction is meaningless when it is too small, so we only compare these two frames when at least one of them has reconstruction greater than or equal to a threshold $\theta$.
  • Results: The comparison results are below (unit is %), and we can observe that in general, using the inertial frame is more numerically stable than PCA.

We added this discussion in Sec C.4 in the revised version.

Q: Kabsch algorithm.

Kabsch algorithm is one way to compute the optimal rotation matrix that minimizes the root-mean-square deviation (RMSD) between two sets of points (atoms in our case). However, it is guaranteed in the COD-Cluster17 dataset that the molecules in weakly correlated structures can rotate to molecules in strongly correlated structures; in other words, the RMSD can be approximately 0 if we use the Kabsch algorithm, which is equivalent to calculating the rotation matrix directly after we fix the poses. We show how to calculate the rotation matrix in the experiment above.

We added this discussion in Sec C.5 in the revised version.

Q: Random sampling baselines.

Thank you for raising this point.

• As far as we know, there are no deep learning baselines assuming the rigidity during crystallization simulation, which is the main motivation of our proposed AssembleFlow.
• Meanwhile, what you raised, random sampling on SO(3) space is interesting. So we haved added it for comparison, marked as Random. We present the results on COD-5000 below, with more comprehensive results available in Section H.1. As observed, the Random baseline performs exceptionally well across all three validity metrics; however, its packing matching (PM) is significantly worse by an order of magnitude.