small systems - simple setting…include experiments on more challenging tasks

---

In our response to Reviewer WqZ3, we presented results for GEOM based on your suggestion. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.

---

We agree that systems from QM9 and MP20 are small, but they are complex and challenging. Doing well on them requires models to learn something about the underlying physics of atomic interactions.

Let’s take MP20 as an example:

• MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is complex due to periodicity.
• The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has improved S.U.N. rate from 4-5% of previous works up to 6.5% (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.

Ultimately, the goal of this paper was to introduce the first unified architecture for generative modelling of 3D atomic systems, and to demonstrate transfer learning across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.

Also note that our best models are based on pure Transformers - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.

extremely small systems - up to 10 atoms

Just to clarify the details here:

• MP20: 20 atoms in a periodically repeating unit cell. This means that each atom can interact with a far greater number of atoms than the size of the unit cell due to periodicity.
• QM9: 9 heavy atoms, which doesn’t include hydrogen. When including hydrogen (which is how our models are trained), systems go up to 20-30 atoms on average.

Which is why we wrote “tens of atoms”, not 10 atoms.

extend metrics for small molecules, following MiDi

Our evaluation followed Symphony, the most recently accepted small molecule generation paper at ICLR 2024. We believe that Symphony's evaluation metrics are an improvement upon MiDi.

Here are our justifications:

• We have included the validity and uniqueness metrics from MiDi/EDM. We found ADiT to outperform EDM and GeoLDM, both of which do better than MiDi.
• Instead of atom and molecule stability metrics from MiDi/EDM which based on simple valency heuristics, Symphony used the PoseBusters metrics suite which provides information about the physical realism of generated molecules. PoseBusters provides a more holistic set of metrics using force field relaxations and physics-based checks instead of relying on heuristics.
• PoseBusters also includes metrics for bond distances, angles, rings, and geometries, which supersede MiDi's histogram metrics and are more interpretable (in our opinion).
• PoseBusters is now widely known and used in academia and industry, e.g. AlphaFold3 used PoseBusters.

"the paper proposes another approach to generate both material and molecules; as such the idea of learning a model across modalities is interesting and fairly new"

The goal of this paper was to introduce the first unified architecture for generative modelling of 3D atomic systems, and to demonstrate transfer learning across modalities.

We disagree with your characterisation that our work is just “another” approach as we are not aware of any other generative architecture that jointly generates periodic and non-periodic systems. We believe this will make a novel and valuable contribution to the ICML community, and would kindly request you to reconsider your position on our novelty.

citing Flam-Shepherd et al. 2023 and Pinheiro et al 2023

We will definitely cite and discuss both papers in the Related Work - thanks for the pointers. We will also cite and discuss the MiDi paper and metrics in our revision.

(Note that this year, ICML does not allow us to upload an updated PDF.)

relatively small structures - benchmark against GEOM

---

In our response to Reviewer WqZ3, we presented results for GEOM based on your suggestion. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.

---

We agree that systems from QM9 and MP20 are small, but they are complex and challenging. Doing well on them requires models to learn something about the underlying physics of atomic interactions.

Let’s take MP20 as an example:

• MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is complex due to periodicity.
• The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has improved S.U.N. rate from 4-5% of previous works up to 6.5% (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.

Ultimately, the goal of this paper was to introduce the first unified architecture for generative modelling of 3D atomic systems, and to demonstrate transfer learning across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.

Also note that our best models are based on pure Transformers - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.

MPTS-52

Not applicable as MPTS is meant for evaluating structure prediction, where TS = temporal split, which is different from our de novo generation task.

unjustified claims of primacy - Flam-Shephard

To the best of our knowledge, ADiT is the first unified, jointly trained generative model for periodic and non-periodic systems to demonstrate transfer learning.

Flam-Shepherd trains 3 independent models on 3 different datasets with a different tokenisation strategy for each. They don't discuss how/whether their method can be applied for one unified model. Also, our results on MP20 are better than FlowLLM, another language model which outperforms Flam-Shepherd.

We will discuss their paper in Related Work, as well as all other references you shared.

latent diffusion for generating chemical structures has been explored prior

Our novelty is about how we used latent diffusion: For unification of system types (periodic and non-periodic systems together), as well as unification of mulit-modal data (categorical, numerical, floating point).

We think this is new - nobody has used latent diffusion in this unified manner - and will be of interest to the ICML community.

• [a] - Not directly related to 3D structure
• [b] - They technically do latent diffusion, but the latents are still multi-modal (four latent scalars & one latent vector) - thus GeoLDM still uses equivariant diffusion, which is slow, and GeoLDM is not applicable beyond small molecules. Additionally, there are several discrepancies between the implementation reported in the paper and released on GitHub (most concerning: https://github.com/MinkaiXu/GeoLDM/issues/6). ADiT’s latent diffusion formulation is unified and faster, as well as outperforms GeoLDM.
• [c] - Highly specific to proteins as latent representations are created by aggregating along the sequence

(We will cite [a] and [c])

high number of lambda coefficients, lack of intuition

The coefficients are needed for each of the different data types that constitute a 3D atomic system. The coefficients are not dataset dependent. The intuition for choosing them is simple: “balance the relative magnitudes of the various losses”. This is stated in the paper and is standard practice for supervised learning tasks like the VAE reconstruction used here.

conditional downstream tasks

The focus of this paper on introducing the architecture and demonstrating transfer learning. Efforts for extensions to conditional tasks such as protein-pocket conditioned molecule generation (ie. SBDD) as well as property conditioned crystal generation (similar to MatterGen) are underway as independent papers. From a methods standpoint, adding conditioning is straightforward in Diffusion Transformers via classifier-free guidance tokens.

We don’t feel that property conditioned small molecule generation, as benchmarked in papers like EDM/GeoLDM, is very practically relevant. We feel that it is unlikely that generating molecules with specific HOMO-LUMO gaps or dipole moments are relevant for drug discovery.

Suggestions

We will implement both.

trained on relatively small datasets - scalability on larger datasets

---

In our response to Reviewer WqZ3, we presented results for the larger GEOM dataset of small molecules with up to 180 atoms. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.

---

We agree that systems from QM9 and MP20 are small, but they are complex and challenging. Doing well on them requires models to learn something about the underlying physics of atomic interactions, and cannot be solved by simple approaches.

Let’s take MP20 as an example:

• MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is complex due to periodicity.
• The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has improved S.U.N. rate from 4-5% of previous works up to 6.5% (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.

Ultimately, the goal of this paper was to introduce the first unified architecture for generative modelling of 3D atomic systems, and to demonstrate transfer learning across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.

Also note that our best models are based on pure Transformers - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.

lack of theoretical justification for why a shared latent space works

There are strong theoretical justifications of our choices:

• The underlying physics of atomic interactions is the same across all 3D atomic systems. Interatomic distance between a carbon atom double bonded to an oxygen atom will be the same, whether they are part of a molecule or a crystal structure. Thus, a shared/unified latent space enables the model to learn shared principles of interatomic interactions.
• Using a shared latent space unifies system modalities (both periodic and non-periodic atomic systems embedded together), as well as unifies multi-modality data (categorical, numerical, floating point embedded together). This makes it very easy to train a simple Gaussian diffusion model on the latent representations, instead of more complex equivariant diffusion formulations.
• The advantage of jointly embedding periodic & non-periodic systems for ML force fields, aka. universal interatomic potentials was shown by recent works like JMP and MACE. More broadly, all of AI research is moving towards joint training of large models on multiple datasets, ie. learning unified latent representations.

trade-offs of equivariance? ability to respect physical symmetries?

When developing ADiT, we ablated the impact of enforcing equivariance on the model -- see Appendix D. We tried to be pragmatic about whether or not to use equivariance.

The non-equivariant version improved generative performance and physical realism of the crystals/molecules, as judged by evaluation metrics explicitly focussed on physics-based tests s.a. PoseBusters suite for molecules, and DFT-based rate for crystals.

Overall, we and others have found that equivariance is not strictly necessary for training a good generative model. Perhaps it can even be an advantage: if you start with two rotated versions of a noisy Gaussian and non-equivariant denoising leads to two different (but equally valid) molecules/crystals, that gives you a more diverse generative model without sacrificing on performance!

joint training on MOFs reduces validity?

Simple answer: training did not converge during the time period that we were allocated compute resources on a large enough GPU cluster for this experiment.

Note that joint training lead to very strong validity for both crystals (91%) and molecules (95%), while reducing validity for MOFs. When looking at training dynamics (right side plot in Table 3), the model first learns crystals, then molecules, and finally starts learning MOFs later in training. We believe training for longer is likely to close the gap and potentially improve beyond the MOF-only variant of ADiT.

Ultimately, our goal with including preliminary MOF results was to encourage the community to work further on MOFs - hybrid organic-inorganic materials for carbon capture and sustainability. We will release open source code and datasets which make it easy for others to build upon our initial experiments.

Thank you for your positive comments and excellent summary of the work! Thanks for appreciating that the latent diffusion idea can be further applicable to multi-modality data in other scientific domains, too.

does not necessarily indicate that the model has learned the underlying physics

ADiT obtains very strong results on both MP20 and QM9, exceeding the current state-of-the-art. These systems are small but they are complex and challenging. Doing well on them requires models to learn something about the underlying physics of atomic interactions, and cannot be solved by simple co-occurence/statistical approaches.

Let’s take MP20 as an example:

• MP20 is restricted to up to 20 atoms in a unit cell because this is practically most relevant and largely covers the Materials Project distribution of known materials. Each atom's local environment is highly complex due to infinite periodic tiling of a unit cell.
• The S.U.N. rate metric for crystals/MP20 is a very challenging and practically relevant metric for materials discovery. ADiT has improved S.U.N. rate from 4-5% of previous works up to 6.5% (see Table 4). ADiT significantly outperformed the recently published MatterGen (Nature, 2024). Previous works showing smaller improvements have been published at top conferences (most recent example: FlowLLM at NeurIPS, 2024). There is a large community of researchers who care about S.U.N. rate on MP20.

or that it can generalize well to very large-scale data

We agree, but the goal of this paper was to introduce the first unified architecture for generative modelling of 3D atomic systems, and to demonstrate transfer learning across periodic crystals and non-periodic molecular systems. We believe this will make a novel and valuable contribution to the ICML community.

Below, we presented results for the larger GEOM dataset of small molecules with up to 180 atoms. ADiT continues to show strong results, generating physically realistic molecules and outperforming baselines.

Also note that our best models are based on pure Transformers - which have shown to be extremely scalable to large-scale inputs - so there is nothing inherent to our methodology that would prevent scaling up to larger systems or more samples.

---

TO ALL REVIEWERS - PLEASE READ BELOW - RESULTS ON GEOM

---

To demonstrate the scalability of ADiT to larger systems, we have run experiments on GEOM, as suggested by Reviewers KRy3 and 7DhL. GEOM includes 430,000 unique molecules and consists of larger systems than QM9, up to 180 atoms. Our experiments and evaluation follows the EDM/MiDi paper, with additional PoseBusters physics-based tests. Model setup and hyperparameters used are exactly as described in the paper.

ADiT-S (32M) outperforms or matches EDM and MiDi across all metrics. ADiT generates physically realistic molecules based on all PoseBusters criteria. Notably, molecules generated by ADiT are significantly more likely than EDM/MiDi to be connected as measured by the 'Atoms connected' score (measures % of molecules where there exists a path along bonds between any two atoms).

We have not trained the larger ADiT-B and ADiT-L models due to resource constraints in the short rebuttal period. We expect larger models to further improve performance based on the scaling trends we have seen on QM9 and MP20.