MOFDiff: Coarse-grained Diffusion for Metal-Organic Framework Design

The first paragraph in section 3 is difficult to read - would you consider including a "concept figure" that outlines the overall training + generative process?

Thank you for the great suggestion. We have added Figure 11 to illustrate the training and sampling processes. It is currently placed in Appendix B (page 18) due to the page limitation. We welcome further feedback on the new figure.

We look forward to further discussions if you have additional questions or suggestions.

Reference:

[1] Yim, Jason, et al. "SE (3) diffusion model with application to protein backbone generation." arXiv preprint arXiv:2302.02277 (2023).

We thank reviewer 1BzN for helpful feedback and comments. We address each of the reviewer’s concerns below.

There are many "moving parts" in the modeling pipeline. It would be enlightening with more ablation studies or comparisons of different modeling choices to yield some insight into the sensitivity to the different modeling choices.

We appreciate your feedback regarding the complexity of our modeling pipeline. To provide a more comprehensive comparison, we elaborate on some early experiments and decision-making:

• We experimented with an autoencoder-based approach for building block embedding. However, this method significantly underperformed compared to the contrastive learning embedding technique we ultimately employed.

• We also explored incorporating orientation into the building block representation and experimented with learning orientation diffusion. This approach, however, yielded significantly inferior results compared to our current assembly algorithm. More details on this will be provided in our subsequent response.

We are now exploring alternative building block identity other than ECFP4. We will update the results as they finish. We are happy to further discuss what other comparison/analysis might be interesting to include.

Some technical details are not described in detail. The paper builds on several existing techniques, which are referenced but not described technically. This is just a minor point, but in my view the paper could perhaps be more self contained.

Thank you for the concrete feedback. We have included a description of the metal-oxo/MOFid algorithms, the ECFP4 fingerprint, and the GemNet-OC architecture in the paper in Appendix B.3 (page 20). We welcome further feedback on this section.

In the coarse grained structure, are building blocks only represented using their identity and position? Did you consider including other features such as their orientation?

Thank you for the insightful question. Yes, in the coarse grained structure, the building blocks are only represented using their identities and positions. In early experiments, we have also attempted to include the orientation as part of the building block representation and diffusing the orientation of the building blocks based on SO(3)-diffusion [1]. However, we could not get good performance. The trained model is unable to recover accurate orientation through the reverse diffusion process. We believe this is due to the orientation of building blocks being ill-defined. In our failed attempts, we use principal component analysis to obtain the canonical orientation of the building blocks. Unlike amino acids, which have a natural definition of orientation based on the $N - C_\alpha - C$ atoms [1], the geometry of the building blocks is much more diverse. There exist many almost-2D or near-isotropic building blocks which makes the orientation ambiguous. Further, the layout of the building blocks is delicate and requires accurate alignment between the building blocks to render a valid MOF.

How sensitive is the representation learning to the choice of fingerprint? Is ECFP4 the only fingerprint you have considered, and why did you choose that?

In the current paper, we have only considered the ECFP4 fingerprint because of its simplicity, popularity, and satisfying results. We agree with the reviewer the choice for building block identity is worth further consideration. As also suggested by reviewer yqBG, we plan to experiment with learned identity for building block representation. We will update the results as they finish.

Abstract: What does "predicting scores in E(3)" mean? At this point it is not clear which scores this refers to.

We are sorry for the unclarity. We have revised the abstract in the updated version.

"small geometric variations in 3D orientation": I assume "orientation" here means something other than rotation/tranlation?

Thanks for pointing out this unclarity. Yes, since the GemNet-OC encoder is SE(3) invariant, the building block embedding is invariant to global translation/rotation. We have revised this sentence in the updated version.

Would there be room to include a short technical description of the metal-oxo / MOFid algorithms that are used? Similarly, a sentence that technically describes ECFP4?

We have included a short description of the metal-oxo / MOFid algorithms and ECFP4 in Appendix B.3.

Could you include a very brief technical description of the gemnet-oc architecture?

We have included a brief description of the GemNet-OC architecture in Appendix B.3.

We thank reviewer yqBG for helpful feedback and comments. We address each of the reviewer’s concerns below.

As acknowledged in the paper, the validity of the MOFs decreases as number of building blocks, this could limit applicability to larger and more complex MOFs.

We recognize that as the complexity of MOFs increases, the validity of our approach may diminish. In future work, we aim to refine our methodology to enhance its applicability to more complex MOFs. This could involve incorporating the lattice parameters into the diffusion process, using known templates for guidance, and enhancing building block embeddings.

The authors use a MOF dataset with less than <20 building blocks. But the paper doesn't discuss about size extensivity of the GNN encoder and denoising module in detail.

It is straightforward to apply the current model to MOFs with a larger number of building blocks. In this paper, we limit the size of MOFs under the hypothesis that MOFs with extremely large primitive cells may be difficult to synthesize.

There is limited discussion on the synthetic accessibility of the designed MOFs. Finally, the building blocks are constrained to a certain set to allow coarse graining, this may led to less diverse MOFs.

Determining the synthetic accessibility of MOFs is very challenging, and there is unfortunately no widely accepted and generally applicable method. To the best of our capability, we define chemically informed validity criteria and use molecular simulations to filter for promising MOF candidates. We look forward to attempting the synthesis of top candidates in future endeavors.

The building blocks used in this work are extracted from the training dataset, BW-DB. This building block space contains 242,000 distinctive building blocks that enable a broad generative scope and satisfying performance in our tasks. While we focus on the BW-DB dataset in this paper due to the availability of gas adsorption labels, the space of possible building blocks can be expanded by incorporating new datasets. Our fully automated pipeline allows us to extract new building blocks without the need to curate the building blocks for template compatibility.

Why was the UFF chosen as force field? Would a MOF specific force field result in better structures?

Both UFF [1] and UFF4MOF [2] are widely used in existing MOF literature. It has been shown in previous work [3] that they have similar performance. We chose UFF over UFF4MOF to be consistent with the well-established GCMC simulation protocol for gas adsorption, which uses the UFF force field. In addition, through manual inspections, we find the relaxed structures from UFF reasonable.

ECFP4 was chosen as similarity measure for contrastive learning. Would a better similarity measure e.g., learned embeddings (ChemBERTa) result in better GNN encoder?

Thanks for the great suggestion. We agree that a more chemically informed building block identity has the potential to enhance our model. We are now exploring the ChemBERTa [4] embedding instead of ECFP4 for building block representation and will update further results when they finish.

On page 19, there is a mismatch between Py-G and pytorch references.

Thank you for pointing this out. We have corrected this in the revised manuscript.

We look forward to further discussions if you have additional questions or suggestions.

Reference:

[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] Nandy, Aditya, et al. "A database of ultrastable MOFs reassembled from stable fragments with machine learning models." Matter 6.5 (2023): 1585-1603.

[3] Boyd, Peter G., et al. "Force-field prediction of materials properties in metal-organic frameworks." The journal of physical chemistry letters 8.2 (2017): 357-363.

[4] Ahmad, Walid, et al. "Chemberta-2: Towards chemical foundation models." arXiv preprint arXiv:2209.01712 (2022).

In general, while the authors include relevant work in the paper, the authors should have a related work section to place their work in the context of the field. The following papers could potentially be related to the present work [2], [3].

Thank you for the thoughtful suggestion of adding a related work section, and for pointing out additional related works. We agree more context for the present work will enhance the manuscript. We have added a related work section (Appendix A) to discuss the relevant works in more detail, as well as include additional related works pointed out by the reviewer. Due to the page limitation, the related work section is in the appendix. We welcome further feedback on this related work section.

In Figure 5, are only novel MOFs (i.e. not in the reference dataset) used to generate the

Yes, only MOFs that are valid, novel, and unique are included in Figure 5.

In the representation learning of building blocks, the authors mention small geometric variations of the building. Are these variations in the coordinate space? In other words, what transformations are used to provide positive samples in the contrastive loss?

The geometric variation naturally exists in the dataset because the same building block appears in the same or different MOFs multiple times with geometric differences (reflected in Figure 3). Therefore, we don’t need to define transformations to provide positive samples. Instead, we use any two 3D building block structures as a positive pair if they share the same ECFP4 fingerprint and the same number of connection points.

How are the coordinates of the building blocks assigned?

The coordinates of the building blocks are generated through the reverse diffusion process. The reverse diffusion process starts by randomly sampling the building block coordinates and identities. From this noisy structure, the learned score networks iteratively denoise the noisy structure to a final coarse-grained MOF structure. The all-atom MOF structure is then recovered from the CG structure through the assembly algorithm and force field relaxation.

We look forward to further discussions if you have additional questions or suggestions.

Reference:

[1] Yao, Z., Sánchez-Lengeling, B., Bobbitt, N. S., Bucior, B. J., Kumar, S. G. H., Collins, S. P., ... & Aspuru-Guzik, A. (2021). Inverse design of nanoporous crystalline reticular materials with deep generative models. Nature Machine Intelligence, 3(1), 76-86.

[2] Jablonka, K. M. (2023). mofchecker (Version 1.0.0) [Computer software]. https://doi.org/10.5281/zenodo.1234

[3] Bucior, Benjamin J., et al. "Identification schemes for metal–organic frameworks to enable rapid search and cheminformatics analysis." Crystal Growth & Design 19.11 (2019): 6682-6697.

[4] Boyd, Peter G., et al. "Data-driven design of metal–organic frameworks for wet flue gas CO2 capture." Nature 576.7786 (2019): 253-256.

We thank reviewer nQJ6 for helpful feedback and comments. We address each of the reviewer’s concerns below.

In the generation process, it is unclear why the latent vector Z can be sampled from the standard normal distribution. The latent vector is generated using the OrbNet encoder trained with the contrastive learning loss. Without further regularization, like in the case of variational autoencoders using the KL divergence, the latent vector should not conform to the distribution.

We are sorry for the confusion. We hope to clarify that the MOFDiff model is indeed a variational autoencoder model trained with the KL regularization loss (Equation 12). This allows sampling MOF structures by sampling latent vectors from a Gaussian prior. The contrastive loss is not used for CG diffusion of MOFs. Instead, it is only used for learning the building block embedding. Only the building block encoder is trained with the contrastive loss. The building block encoder is then frozen and used to obtain learned embedding for all building blocks in the dataset. These building block embeddings are then used to represent CG MOF structures and diffusion-based generative modeling.

We hope the response above clarifies the confusion. We are more than happy to provide additional clarification if needed.

It would be useful to compare with previous work [1] on the inverse design of MOFs. The authors mention and cite a source on the low diversity of the generated MOFs with template-based systems such as SmVAE, but a comparison isn’t present in the experiments section.

We have attempted to conduct more experiments on the SmVAE model proposed in [1]. However, although the VAE source code of [1] is released, the MOF deconstructor and reconstructor used in [1] are not clearly stated or publicly available. Without the code for these steps, we are unable to run SmVAE for a new dataset, or recover the all-atom structures for the dataset or generated samples of SmVAE.

While a direct lateral comparison is not feasible, we were able to download the top 9 candidate MOFs released by [1] (GMOF-1 to GMOF-9) for CO2/N2 separation. We run our molecular simulation workflow on these MOFs and report their gas adsorption properties in Table 1. The best GMOF attains a working capacity of 2.20 mol/kg, a CO2 uptake at the adsorption stage of 2.53 mol/kg, and a CO2/N2 Selectivity of 11.46. In comparison, the MOFs generated by MOFDiff exhibit better carbon capture performance. However, it is important to acknowledge that this comparison is not entirely rigorous due to the difference in training, dataset, simulation, and optimization settings. All of our code will be open-sourced.

In addition, we hope to highlight the methodological difference between our method and SmVAE. Our method generates 3D MOF structures without templates and thus offers a distinct perspective compared to existing template-based methods such as SmVAE in addressing the diversity of computational MOF design.

The validity and novelty of generated structures are low compared to template-based models as in [1].

In the previous response, we highlighted the significant obstacles in making an apple-to-apple comparison to [1]. We wish to clarify that the validity and the novelty statistics reported in [1] and our manuscript are not directly comparable for the following reasons:

1. The criterion for validity/novelty is not clearly defined in [1]. Our validity criterion is defined as the simultaneous satisfaction of (1) matched connection, (2) successful force field relaxation, and (3) passing all criteria defined in MOFChecker [2]. Our novelty is defined through MOFid [3]. Our code will be open-sourced so the validity/novelty criterion can be reproduced in future works.

2. The training data set is different. The SmVAE model [1] was trained on a customized database of “around 2 million MOFs”. Our model is trained on BW-DB (~300k MOFs) with the original labels for gas adsorption properties [4].

3. The generation scheme is different. The SmVAE model generates MOFs based on predefined templates. Our method generates 3D structures without templates. The SmVAE model and our method also use different sampling protocols.

Reference:

[1] Yao, Zhenpeng, et al. "Inverse design of nanoporous crystalline reticular materials with deep generative models." Nature Machine Intelligence 3.1 (2021): 76-86.

[2] Park, Hyunsoo, et al. "Inverse design of metal-organic frameworks for direct air capture of CO2 via deep reinforcement learning." (2023).

[3] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.

[4] Gruver, Nate, et al. "Protein Design with Guided Discrete Diffusion." arXiv preprint arXiv:2305.20009 (2023).

[5] Xie, Tian, et al. "Crystal diffusion variational autoencoder for periodic material generation." arXiv preprint arXiv:2110.06197 (2021).

[6] Yim, Jason, et al. "SE (3) diffusion model with application to protein backbone generation." arXiv preprint arXiv:2302.02277 (2023).

[7] Jiao, Rui, et al. "Crystal Structure Prediction by Joint Equivariant Diffusion on Lattices and Fractional Coordinates." NeurIPS (2023).

We deeply appreciate your suggestions on further improvements over the current model, which we respond to below.

One potentially problematic aspect is that the number of blocks K and the lattice structure L are computed based on the sample z which also conditions the diffusion of the attributes and the locations of the blocks. This can also lead to invalid MOF structures as the block's location and orientation may not match the lattice geometry. You identify this as well in your last sentence. Why not incorporate these in the diffusion process? Or alternatively, why not compute the lattice structure based on the atoms and the bond in the unit cell? Wouldn't the atoms (block) and the bonds uniquely identify the lattice structure of the crystal?

Thank you for the great suggestion. Indeed, as we conclude in our paper, we agree incorporating the lattice parameters into the diffusion process is a promising future direction to resolve the mismatch between the blocks and the lattices, and significantly improve the current model.

However, we find incorporating lattice diffusion to MOFDiff presents significant challenges. Existing work that focuses on inorganic crystals [7] has explored lattice diffusion. Their formulation is to diffuse the lattice parameters to N(0, I) and represent the coordinates of atoms in the fractional coordinates because the cartesian coordinates will become numerically unstable for small lattice sizes and the cartesian scale of the system changes across the diffusion process.

Applying a similar strategy to MOFs is more complex. First, using fractional coordinates to represent the coarse-grained coordinates may not be ideal, as fractional coordinates cannot accurately represent the essential size information of these blocks, while the sizes of the building blocks vary drastically. Moreover, as the building block identities and their respective sizes evolve during the reverse diffusion process, maintaining accurate and meaningful representations of the lattice and building blocks may be challenging. Given these challenges, we believe it requires careful consideration and novel approaches to address the unique characteristics and demands of MOFs when designing a lattice diffusion process. We value your input and recognize the importance of this issue. We are committed to addressing it in our future research endeavors.

Why generate from an uninformed prior N(0, I)? It seems unreasonable to expect that sampling from such a distribution would give good coverage of the vast space of possible MOF configurations. Would it not be more effective to condition on a number of building blocks or present partial coarse structures? Possibly many other well-understood properties of the MOFs?

Our first experiment samples MOFs from the Gaussian prior N(0, I) of the VAE model. For a well-trained VAE, sampling from the prior should result in the generation of new data instances that reflect the general characteristics of the training data. In this unconditional generation experiment, we aim to evaluate our model’s capability to recover the training distribution in a general sense. In other words, we aim to validate that MOFDiff can generate valid MOFs (Figure 6) with a wide spectrum of structural properties that resemble the training distribution (Figure 5). In contrast, a badly trained model may have “mode collapse”, and consequently be unable to capture the training distribution.

Your generation process is limited to using the building blocks present in the training data. How broad of a coverage does this give the generating process? Are there many other MOFs possible with building blocks not present in the training data?

Our method operates over 242,000 distinctive building blocks from the BW-DB dataset. Our experimental results demonstrate the current space of building blocks can already enable structurally diverse generation, as well as effective inverse design for carbon capture applications. While we focus on the BW-DB dataset in this paper due to the availability of gas adsorption labels, it is straightforward to expand the space of possible building blocks (and the scope of MOF generation) by incorporating new datasets.

We look forward to further discussions if you have additional questions or suggestions.

We thank reviewer e39L for helpful feedback and comments. We address each of the reviewer’s concerns below.

While a strong contribution to Material Science, from the ML perspective the work is a good combination of existing methods but does not go much further than that. The paper tells a story about many of the design decisions that were taken, but many other choices could have been made with even more recent approaches. In that sense, the main weakness is the alignment to the ICLR conference.

Thank you for appreciating the materials science and ML application aspects of our paper. We believe our paper also makes methodological contributions and is aligned with the ICLR audience for the following reasons:

1. Our paper addresses the important and difficult task of MOF design, for which ML methods have great potential. We believe our model design makes novel and non-trivial contributions from an ML perspective:

• We propose contrastive learning for embedding the enormous space of building blocks to be used for generative modeling. In particular, in early experiments, we found autoencoding-based embedding has a much inferior performance compared to contrastive learning.
• We combine learning and algorithmic methods to solve a task that is hard to solve with a pure learning approach. In early experiments, we attempted a pure learning approach of SO(3) diffusion for building block orientation, but couldn’t get good results.
• Coarse-grained diffusion. Our diffusion model operates over CG coordinates and a vast space of learned building block identities. This is novel compared to existing diffusion models for small molecules, proteins, or crystals [3,4,5], where the atoms/amino acids have discrete types.

2. Our paper provides scientifically meaningful tasks for the ML community to make rapid progress on MOF design.

3. We will open-source our code for the entire pipeline of MOF decomposition, MOF coarse-graining, diffusion model, all-atom MOF reconstruction, relaxation, and molecular simulation. We believe our code will make MOF modeling and design much more accessible to the broader ML community.

In summary, we believe the methods, tasks, and codebase proposed in our paper will be interesting to the ICLR audience, especially the AI4Science community. In particular, our method offers new perspectives for the ML modeling of multi-scale molecular and materials systems that are ubiquitous in biological and physical science.

More details on the orientation diffusion attempt:

In early experiments, we also explored diffusing the orientation of the building blocks based on SO(3)-diffusion, which has shown promising results in protein diffusion models [6]. However, we could not get good performance. We believe this is due to the orientation of building blocks being ill-defined. In our failed attempts, we use principal component analysis to get canonical orientations for the building blocks. Unlike amino acids which have a natural definition of orientation based on the $N - C_\alpha - C$ atoms, the geometry of the building blocks is much more diverse. There exist many almost-2D or near-isotropic building blocks which makes the orientation ambiguous. Further, the layout of the building blocks is delicate and requires accurate alignment between the building blocks to render a valid MOF.