MOFFlow: Flow Matching for Structure Prediction of Metal-Organic Frameworks

Dear reviewer 1rvy,

We deeply appreciate your thorough review and constructive comments on our manuscript. In the revised manuscript, we have marked the changes in blue. Below, we address the reviewer's concerns in detail.

W1. The discussion on periodicity in MOFs is insufficient (e.g., periodic CoM, periodic edge connections). Additional details on modeling periodicity would enhance the paper.

We appreciate the reviewer’s insightful comment. MOFFlow addresses periodicity implicitly through learning similar to [1]. We attempted to directly model periodicity before submission but encountered training instability and poor performance. We elaborate our findings below.

As noted by the reviewer, periodicity on Cartesian coordinate system can be achieved by constructing periodic invariant edges and using model architectures that update periodic invariant features ([3, 4, 5]). However, implementing this approach in our framework revealed two challenges:

1. Training instability: We use a joint prediction framework, where the lattice parameter (L) is not fixed during periodic edge construction. This contrasts with previous works that either use ground-truth L (Matformer[3]) or generate L in a separate conditional step before constructing periodic edges (Symat[4], MOFDiff[5]). The absence of a fixed L likely caused substantial training instability.
2. Limited expressivity of the model: Architectures designed for periodic invariance (i.e., those updating invariant edges instead of explicit coordinates) performed poorly, likely due to limited model expressivity. This aligns findings in MCF[6] and AlphaFold3[7], where hard-coded inductive biases can degrade performance or provide no benefits.

These challenges highlight the difficulty of directly modeling periodicity within our framework, though we remain committed to addressing this problem. In particular, we find that the periodic CoM approach [2] mentioned by the reviewer presents a promising direction for future work.

W2. Design of backbone model lacks clarity. It is unclear whether atom-level and block-level modules are updated sequentially or alternatively. A figure for overall architecture would be beneficial.

We apologize for the lack of clarity. The update process is sequential: the atom-level module generates building block embeddings followed by iterative updates of the block-level module. In response to your suggestion, we have re-drawn Figure 3 in the revised manuscript to clearly illustrate the overall architecture. We appreciate your feedback, which helped improve the presentation of our model.

Q1. How many different types of building blocks can be decomposed from the dataset?

The BW-DB dataset contains approximately 2 million building blocks [5]. This large number arises because building blocks with the same atomic composition $\mathbf{a}_m$ can have slightly different 3D coordinates $\mathbf{Y}_m$ due to variations in torsion angles of rotatable bonds.

Q2. By generating only translations and rotations instead of fine-grained structures, the proposed method assume building blocks as rigid bodies. Could this assumption potentially limit the model's performance?

Thank you for the thoughtful question. We believe the rigid-body assumption is sufficient for the structure prediction task, as many commonly used MOF building blocks are known to remain rigid during the assembly process [8].

That said, extending our model to predict the dynamics of MOFs would likely benefit from relaxing this assumption and explicitly modeling torsional angles. This would be an exciting direction for future work.

Reference

[1] Köhler et al., Rigid Body Flows for Sampling Molecular Crystal Structures, ICML 2023.
[2] Lin et al., Equivariant Diffusion for Crystal Structure Prediction., ICML 2024.
[3] Yan et al., Periodic Graph Transformers for Crystal Material Property Prediction, NeurIPS 2022.
[4] Luo et al., Towards Symmetry-Aware Generation of Periodic Materials, NeurIPS 2023.
[5] Fu et al., MOFDiff: Coarse-grained Diffusion for Metal-Organic Framework Design, ICLR 2024.
[6] Wang et al., Swallowing the Bitter Pill: Simplified Scalable Conformer Generation, ICML 2024.
[7] Abramson et al., Accurate structure prediction of biomolecular interactions with AlphaFold3, Nature 2024.
[8] Yusuf et al., Review on Metal–Organic Framework Classification, Synthetic Approaches, and Influencing Factors: Applications in Energy, Drug Delivery, and Wastewater Treatment, ACS Omega 2022.

Q1. Are the lattice parameters actually treated as global features of a finite 3D graph in the proposed framework? A simple "yes" is fine, as I understand that incorporating periodicity is a challenging issue.

Yes, they are treated as global features of a finite 3D graph. Thank you for understanding that incorporating periodicity is a significant challenge within our framework.

Q2. How is the rigid structure of each building block determined during the generation process?

Thank for you raising this question. We assume that the structure of each building block is known during the generation process. This is a valid assumption given the availability of extensive libraries of MOF building blocks with detailed structural information [1, 2]. For building blocks not found in these libraries, their 3D structure can be reliably determined using conformer generation tools, as MOFs typically consist of simple and rigid building blocks [3].

We hope that our responses have addressed all your concerns thoroughly. If you have any further questions or require additional clarifications, we would be happy to provide them!

[1] Boyd, Peter G., et al. "Data-driven design of metal–organic frameworks for wet flue gas CO2 capture." Nature 576.7786 (2019): 253-256.
[2] Gibaldi, Marco, et al. "The HEALED SBU library of chemically realistic building blocks for construction of hypothetical metal–organic frameworks." ACS Applied Materials & Interfaces 14.38 (2022): 43372-43386.
[3] Yusuf et al., Review on Metal–Organic Framework Classification, Synthetic Approaches, and Influencing Factors: Applications in Energy, Drug Delivery, and Wastewater Treatment, ACS Omega 2022.

W4. (Minor) Paper mainly focuses on properties relevant to structure prediction accuracy (e.g., match rate and RMSE). Explore how MOFFlow's predictions perform in real-world settings that are essential for the practical use of MOFs.

Thank you for the suggestion. MOFFlow is designed to be directly applicable to real-world workflows that support experimental synthesis. For instance, experimentalists can use MOFFlow to (1) predict a MOF structure from building blocks, (2) assess whether the predicted structure is desired, and (3) proceed with synthesis. This workflow parallels AlphaFold [5], saving the costs for experimental structure determination.

References

[1] Miller et al., FlowMM: Generating Materials with Riemannian Flow Matching., ICML 2024.
[2] Lin et al., Equivariant Diffusion for Crystal Structure Prediction., ICML 2024.
[3] Corso et al., DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking, ICLR 2023.
[4] Yim et al., SE(3) diffusion model with application to protein backbone generation, ICML 2023.
[5] Jumper et al., Highly accurate protein structure prediction with AlphaFold, Nature 2021.

Dear reviewer 6goZ,

We deeply appreciate your thorough review and constructive comments on our manuscript. In the revised manuscript, we have marked the changes in blue. Below, we address the reviewer's concerns in detail.

W1/Q1. Absence of comparison with MOF-specific deep learning models. Add necessary (general CSP) baselines for a fair comparison, which include but are not limited to CDVAE, MOFDiff, FlowMM, DiffCSP++, etc.

First, we note that there exist no MOF-specific deep learning models that are directly comparable to ours. We already compare with the self-assembly algorithm of MOFDiff in Section 5.4. Note that MOFDiff solves a problem different to MOF structure prediction, so the algorithms are not directly comparable.

To alleviate your concern, we have incorporated your suggestion and are currently adding the most recent CSP algorithms (FlowMM[1] and EquiCSP[2]) as new baselines. These experiments are ongoing, and results will be updated during the rebuttal period.

In Table A, we report the performance of FlowMM evaluated after 8 days of training. We observe that FlowMM performs even worse than DiffCSP, which might be attributed to its high resource demands; training FlowMM for 70 epochs took 8 days on an 80GB A100 GPU, while DiffCSP completed 200 epochs in just 5 days on a 3090 GPU. Note that MOFFlow takes 5 days to train for the TimestepBatch version and 1~2 days for the Batch version (discussed extensively in W3/Q3). We are continuing to train FlowMM to ensure fair comparisons in the final paper.

If the reviewer is aware of any more recent and relevant MOF-specific models appropriate for this task, we would be happy to evaluate them in our framework.

Table A. Structure prediction accuracy of Random Search (RS), Evolutionary Algorithm (EA), DiffCSP, FlowMM, and MOFFlow. MR denotes match rate and RMSE denotes root mean squared error.

W2/Q2. Lack of ablations to understand contributions on each module. Providing insights into the relative contributions of different components (e.g., block-level embedding, MOFAttention module) would add clarity.

Thank you for the suggestion. Based on your feedback, we performed ablation studies to evaluate the relative contributions of the block-level embedding module and the MOFAttention module. Specifically, we varied the number of layers and reduced the hidden dimension from 64 to 32, similar to the ablation studies conducted by DiffDock [3]. The results, summarized in Table B, indicate that a single MOFAttention layer has a greater impact on performance than a layer of the building block embedding module.

Table B. Ablation study on MOFAttention and building block embedding layers. We report the match rate (%) and RMSE of the base MOFFlow model and its variations with one layer removed. The results highlight that the MOFAttention layer has a greater impact on performance than a building block embedder layer.

While it would be desirable to evaluate the relative contributions by completely removing each module, this is not possible since there exists no alternatives to our proposed pipeline.

W3/Q3. MOFFlow is stated to be computationally efficient compared to baseline methods. Including a detailed comparison of computational resources (e.g., GPU hours, memory usage, inference costs) against other models would be valuable.

Thank you for the insightful comment. We would like to clarify that MOFFlow’s computational efficiency is highlighted in terms of inference costs, as shown in Table 1. To address your concerns on training costs, we added a comparison of computational resources required to train each model in Appendix C. This covers batch size, GPU type and count, GPU memory usage, and total training time. We present a summary of this in Table C for your convenience.

Table C. Comparison of computational resources required for DiffCSP, FlowMM, and MOFFlow. We report the batch size, GPU type, number of GPUs, GPU memory utilization in GB, and training time in days/hours.

In the process of addressing this suggestion, we identified that the TimestepBatch implementation [4] used in our main experiments takes longer time in terms of GPU hours. Motivated by this insight, we refactored our code, leading to the development of the Batch version (link), which requires significantly less training time to achieve similar performance (1~2 days vs. 5 days). Detailed discussions and results related to the Batch implementation have been added to Appendix D. Both versions will be released upon acceptance.

We deeply appreciate the reviewer’s suggestion, which not only improved the completeness of our work but also guided us toward a meaningful improvement in our methodology.

Q1. For the ablation studies, the gap seems small. Could you provide error bars for multiple runs to show the statistical significance?

Yes, we are currently running additional experiments to address your concerns and will keep our results updated.

Please note that we initially aimed to produce statistically significant results by halving the number of layers in each module. However, this caused training to diverge, highlighting that both modules have important contributions.

Q2. Can you elaborate a bit more on the differences between those two "batched" implementations? What makes the improved version more efficient?

TimestepBatch [1]: This is the approach used in our original manuscript. To create a batch of size $B$, we sample a single datapoint from the dataset and generate $B$ noised instances of this single datapoint. The advantage of this method is that each batch contains the same number of building blocks and the same number of atoms within those building blocks, simplifying implementation and facilitating parallelization.

Batch: This is how a usual batch is created for flow matching models. To create a batch of size $B$, we sample $B$ datapoints from the dataset and generate a single noised instance for each datapoint. While this method stabilizes training, it is much more difficult to implement, as it requires handling varying numbers of building blocks and differing numbers of atoms within those building blocks.

The slow convergence of the TimestepBatch implementation arises because a batch effectively contains variations of a single data instance. This results in high gradient variability between each step, causing unstable training and longer convergence time.

Q3. For the training costs, I noticed that only MOFFlow used multi-gpu training. Do the other baselines support this?

Yes, the other baselines support multi-GPU training. DiffCSP was run on a single GPU since only one additional GPU was available at the time. FlowMM, on the other hand, requires an 80GB A100 GPU due to memory constraints. If you think this is a critical issue we will be happy to run DiffCSP again with 8 GPUs.

We hope that our responses have addressed all your concerns thoroughly. If you have any further questions or require additional clarifications, we would be happy to provide them!

Q. I'm a bit confused about why halving the number of layers would make the training diverge. You still have the necessary components/modules in your model but it's just the model is not that deeper as before. In deep learning, usually a deeper model would be harder to train to converge. That is why we have other tricks like residual connections. Can you provide more insights?

To be honest, it is challenging to pinpoint the exact reason, as understanding the behavior of deep learning models is still an open research question. One of our explanations is that halving the number of layers reduces the model’s capacity to a point where it becomes barely sufficient to fit the training data. Studies suggest that models in such cases are highly sensitive to even minor noise in the data, which can disrupt the model's global structure (refer to Discussion of Section 5 in [1]).

Please note that we are using training stabilization techniques including gradient clipping, layer normalization, residual connections, and more.

Q. I understood that "TimestepBatch" is easier to implement since everything is aligned and consistent. But for "batch", I think there's just one more step to do padding for building blocks and atoms. Is that right?

We do not use padding for building blocks and atoms, as it would be highly memory-inefficient, especially given the large size and variability of MOF structures. In other words, the size of zero-padding is prohibitively large when a pair of small and large MOF structures are inside the same mini-batch.

Instead, similar to existing GNN implementations [2], we represent each batch as a disconnected graph, with edges constructed between atoms within each building block and between building blocks within each data point. This inconsistent data structure posed some challenges during implementation.

Q. One thing I would still want to ask is, since you have an attention module, etc. involved, can you provide the big O complexity to help me understand the time complexity better in terms of both practice and theory?

We appreciate your question. We here share our analysis of the big O complexity.

Consider a MOF containing $N$ atoms and $M$ building blocks, where the number of atoms in $m$th building block is denoted $N_m$ such that $\sum_{m=1}^M N_m = N$. Let $L$ denote the number of layers and $D$ the feature dimension.

The big-O complexity of GNNs that operate on atomic coordinates is $O(L N^2 D)$. In contrast, our building block embedder has a time complexity of $O(L (N_1^2 + \dots+ N_M^2) D)$, as it operates on the building blocks. Note that $N_1^2 + \dots + N_M^2 \leq N^2$ since $N^2 = (\sum_{i=1}^M N_m)^2 = \sum_{i=1}^N N_m^2 + 2 \sum_{m < n}N_m N_n \geq \sum_{i=1}^N N_m^2.$ The equality holds only when there is one building block, which does not occur in our dataset. Therefore, the time complexity of our building block approach is lower.

We note that the time complexity of our attention module is $O(LM^2D)$, which is negligible since the number of building blocks is less than the number of atoms.

For simplicity, this discussion assumes a fully connected graph for the upper-bound complexity. While we construct edges with radius cutoff in practice, analysis with this is difficult since the number of edges depends on the atomic positions.

[1] Nakkiran, Preetum, et al. "Deep double descent: Where bigger models and more data hurt." Journal of Statistical Mechanics: Theory and Experiment 2021.12 (2021): 124003.
[2] Fey, M., & Lenssen, J. E. (2019). Fast graph representation learning with PyTorch Geometric. arXiv preprint arXiv:1903.02428.

Thank you for your reply. I think most of my concerns have been addressed. I'll raise the rating to 6. But I'm still looking forward to the ablation study.

Dear reviewer 6goZ,

We appreciate your patience. We present ablation results over three runs, as shown below. In terms of match rate, the performance of removing a MOFAttention layer and a building block embedder layer is comparable. This indicates that both layers contribute similarly to overall performance.

We attribute this to the complementary roles of these modules: the building block embedder captures atom-level resolution, while the MOFAttention layer focuses on block-level predictions. Together, they are crucial for achieving high performance.

We truly appreciate your helpful feedback and the effort you have put into improving our work.

Dear reviewer 6goZ,

We sincerely appreciate the time and effort you have dedicated to reviewing our work. We would like to address a possible misunderstanding in your comment: without the proposed modules, the pipeline is ill-defined and fails to define a valid mapping from MOF building blocks to the MOF structure.

As noted earlier, removing additional layers caused the model to diverge, highlighting the essential role of these modules in achieving meaningful performance. To ensure stability, we limited our ablation to the removal of a single layer. While this resulted in relatively small changes, the differences remain meaningful; specifically, the difference between the base model and the removal of a MOFAttention layer does not overlap within the standard deviation. We hope this clarifies the importance of our proposed modules.

Dear reviewer hN1H,

We deeply appreciate your thorough review and constructive comments on our manuscript. In the revised manuscript, we have marked the changes in blue. Below, we address the reviewer's concerns in detail.

W. Limited applicability. Crystal structures that cannot be decomposed into secondary structures cannot readily use these techniques.

We agree with this point. Nevertheless, we still believe that our work is significant since decomposable crystals like MOFs, covalent organic frameworks, and molecular crystals, have important real-world applications such as gas separation[1], catalysis[2], drug delivery[3], and sensing[4].

Q1. Awkward notation in equation 1. $x_n + Lk$ is not proper.

We apologize for the oversight. We have revised our manuscript to correct this as $x_n' =x_n + kL^{\top}$ where $k =(k_1, k_2, k_3) \in \mathbb{Z}^{1\times3}$ and $L = (l_1, l_2, l_3) \in \mathbb{R}^{3 \times 3}$. We appreciate your careful attention to detail.

Q2. What is the reference frame for local building block coordinates in line 156?

The reference frame is the PCA axes, with their orientation fixed by a reference equivariant vector as outlined in [5]. This is detailed in Section 4.1, and we have added a cross-reference at line 156 to improve readability.

Q3. In Section 5.2, are the property evaluations done on the predicted structures and compared with the property evaluations on ground-truth structures?

Yes, we compare the property evaluations on predicted structures to that of ground-truth structures. We apologize for the ambiguity and have revised Section 5.2 (lines 461-467) to clarify this.

References

[1] Li et al., Recent advances in gas storage and separation using metal-organic frameworks, Materials Today 2018.
[2] Lee et al., Metal-organic framework materials as catalysts, Chemical Society Reviews 2009.
[3] Horcajada et al., Metal-organic frameworks in biomedicine, Chemical Reviews 2012.
[4] Kreno et al., Metal-organic framework materials as chemical sensors, Chemical Reviews 2012.
[5] Gao & Günneman, Ab-initio potential energy surfaces by pairing gnns with neural wave functions, ICLR 2022.

Q1. How did the authors obtain the initial structure of the individual building blocks? Do they use the conformation from the data? Or predict those with a specific algorithm?

As explained in Section 5.2 (line 367-369), we use the metal-oxo decomposition algorithm of MOFid [2] to obtain the initial building block structures. This algorithm takes a complete MOF structure (in CIF format) as input, identifies the metal-oxo clusters (metal nodes), and classifies the remaining fragments as organic building blocks.

Q2. Have the authors studied the number of steps required for the reverse process in flow matching?

Thank you for pointing this out. Yes, we studied the number of steps to balance performance (match rate, RMSE) and sampling time, concluding that 50 integration steps offered the best trade-off, and used it in our main experiments. These results are now included in Appendix G (and cross-referenced in line 425) of the revised manuscript.

Q3. Have the authors studied the scalability of MOFFlow and the self-assembly algorithm used in MOFDiff, similar to Fig. 6?

We appreciate your suggestion. To address your concern, we ran additional analysis and added Appendix H in our revised manuscript, which studies the scalability and sampling efficiency of MOFFlow and the self-assembly algorithm. In summary, while MOFFlow demonstrates superior overall performance and significantly faster inference, the self-assembly algorithm exhibits slightly better scalability for structures with many building blocks.

Q4. I would like to see more details on how Fig. 5 is produced. Since MOFDiff and MOFFlow models different types of problems, how do the authors perform a fair comparison in Fig. 5?

We note that Figure 5 compares MOFFlow to DiffCSP, not MOFDiff. Indeed, as the reviewer mentions, we do not compare with MOFDiff since it does not address the structure prediction problem.

Reference

[1] Corso et al., DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking, ICLR 2023.
[2] Bucior et al., Identification schemes for metal-organic frameworks to enable rapid sesarch and cheminformatic analysis, 2019.

Dear reviewer zVZp,

We deeply appreciate your thorough review and constructive comments on our manuscript. In the revised manuscript, we have marked the changes in blue. Below, we address the reviewer's concerns in detail.

W1. The individual pieces of the paper are not particularly novel.

Thank you for your feedback. We respectfully highlight that our work introduces both novel pieces and a unique integration as meaningful contributions to the field. Specifically:

1. MOFAttention module: As noted by the reviewer, we introduce the MOFAttention module, a novel block-level update mechanism for MOFs.
2. End-to-end framework: We present the first end-to-end framework that jointly generates lattice parameters, block-level translations, and rotations, offering an comprehensive solution for MOF structure prediction.
3. Novel integration: By uniquely combining existing components with the MOFAttention module, we achieve a substantial performance improvement (31.69% vs. ~0%).

W2. Lack of ablation study that illustrates key component of the model. From DiffCSP to MOFFlow, there are many changes: 1) the adoption of a hierarchical modeling approach; 2) change from diffusion to flow matching; 3) change of the architecture of the score network.

Thank you for highlighting these points. Below, we provide ablation studies to evaluate the contributions of the building block embedder and MOFAttention, two core components of MOFFlow. Specifically, we varied the number of layers and reduced the hidden dimension from 64 to 32, similar to the ablation studies conducted by DiffDock [1]. The results, summarized in Table A, indicate that a single MOFAttention layer has a greater impact on performance than a layer of the building block embedding module.

Table A. Ablation study on MOFAttention and building block embedding layers. We report the match rate (%) and RMSE of the base MOFFlow model and its variations with one layer removed. The results highlight that the a MOFAttention layer has a greater impact on performance than a building block embedder layer.

Regarding the specific ablations suggested by the reviewer:

1. Hierarchical vs. non-hierarchical. We are currently training FlowMM as an additional baseline, which can be considered the non-hierarchical counterpart of MOFFlow.
2. Diffusion vs. flow matching. Implementing a diffusion-based version of MOFFlow is challenging due to our specific choice of priors, which requires deriving an appropriate closed-form kernel for diffusion that converges to this prior. While this task is infeasible within the current timeline, we will explore its implementation and aim to share results in the final version of our paper.
3. Other network architectures. Since our work is the first to introduce a block-level architecture for MOFs, direct comparisons with other architectures are difficult. Instead, we focus on evaluating relative contributions of key model components, as shown in Table A.

Thank you for your support! Regarding the follow-up questions:

Q1. I understand that MOFid is used to decompose the MOF structure to building blocks. I am asking how the authors obtain the initial 3D structure of the building blocks. Do the authors take average over 3D structures in training data? Or do they use RDKit?

We apologize for misunderstanding your question. With ~2 million building blocks in the training data, categorizing them to compute an average 3D structure is computationally infeasible. Instead, we follow MOFDiff and treat each building block as a unique entity, allowing the building block embedder to account for discrepancies in structure by learning consistent representations during training. During generation, we assume the structures are known; this is a valid assumption given the inherent rigidity of building blocks [1] and the availability of extensive libraries of MOF building blocks with detailed structural information [2, 3].

We also find your idea of using RDKit compelling, as it could improve robustness and enable text-based MOF structure prediction. Thank you for sharing this valuable idea.

Q2. Why does the model achieve the best performance at 50 steps, instead of monotonically increasing performance with more integration steps?

This is an interesting observation and we share your curiosity. FlowMM, which exhibits a similar trend, also lacks a clear explanation for this behavior. One of our thoughts is that increasing the number of integration steps may lead to a greater accumulation of numerical errors, which could counteract the benefits of finer integration. However, further investigation is needed to determine the exact reason for this behavior.

We hope that our responses have addressed all your concerns thoroughly. If you have any further questions or require additional clarifications, we would be happy to provide them!

[1] Yusuf et al., Review on Metal–Organic Framework Classification, Synthetic Approaches, and Influencing Factors: Applications in Energy, Drug Delivery, and Wastewater Treatment, ACS Omega 2022.
[2] Boyd, Peter G., et al. "Data-driven design of metal–organic frameworks for wet flue gas CO2 capture." Nature 576.7786 (2019): 253-256.
[3] Gibaldi, Marco, et al. "The HEALED SBU library of chemically realistic building blocks for construction of hypothetical metal–organic frameworks." ACS Applied Materials & Interfaces 14.38 (2022): 43372-43386.