Equivariant Neural Diffusion for Molecule Generation

We thank the reviewer for their constructive feedback.
We address the reviewer’s concerns/questions here below:

Limited technical novelty
As mentioned in our general rebuttal, END is indeed a combination of existing ideas. We however want to emphasize that (1) our work is the first to seek improvement by adding a learnable forward process, orthogonally to most previous work that has mostly focused on the reverse process, (2) designing an appropriate transformation was not straightforward, it required some care as to preserve the desired invariance of the learned distribution.

Relevance of learnable forward process
Most improvements to diffusion models for molecules have focused on designing more expressive denoisers, better noise schedules, or embedding discrete features in a continuous latent space. In END, we seek improvement in an orthogonal direction. Adding a learnable forward allows the hierarchy of latent variables to be learnt, instead of fixed as in classical diffusion models, thereby offering greater flexibility and enhanced generative modeling. We note that END can be combined with previous developments for even better performance.

This enhanced generative modelling is greatly showcased on the more challenging and realistic GEOM-Drugs, where the learnable forward greatly improves connectivity (~20% across all sampling schemes), and leads to 3D structures of better quality as evidenced by the significantly lower strain energies.

Finally, the addition of a learnable forward process comes with two practically relevant by-products: (1) END achieves better performance with fewer integration steps, and (2) the learnable forward yields an improved conditional generative process – as demonstrated on two conditional tasks where cEND even outperforms by a significant margin a model leveraging an auxiliary model through classifier-guidance (EEGSDE) (Tables 3 and 4 in the submitted manuscript).

Previous work
We first note that both papers were already included as baselines / previous work in the submitted manuscript, but, as suggested by the reviewer, we will add a more thorough discussion to the final version. We summarize it here:

• EEGSDE is a continuous-time formulation of EDM, where conditional generation is performed by combining (1) a conditional model score model, (2) a method similar to classifier-guidance (requiring the training of an auxiliary model). In cEND, we instead only learn a conditional model, akin to classifier-free guidance.
• Similar to END, EDM-BRIDGE builds on the observation that there exists an infinity of processes mapping from prior to target distributions. EDM-Bridge constructs one such process that incorporates some prior knowledge, i.e. part of the drift term is a physically-inspired force term. END can be seen as a generalization of EDM-Bridge, where the forward drift term is now learned instead of pre-specified. Through experiments, we show that a learnable forward performs better than a fixed one, even when the latter is physics-inspired.

Additional metrics
We thank the reviewer for their suggestion. We added a citation to HierDiff $1$ and computed additional metrics taken from the paper, namely: SAScore, QED, logP and MW (see rebuttal PDF). To further assess the quality of the generated 3D structures, we also added a metric that measures the strain energy – expressed as an energy difference between the generated geometry and its relaxation (obtained through force-field geometry optimization).

These additional metrics were computed on the valid x connected samples generated by each method. We converted the generated samples to SDF using OpenBabel, and read them in using RDKit. For the SAScore, we normalized the values between 0 and 1, with 0 being “difficult to synthesize” and 1 “easy to synthesize”.
All the additional metrics are provided in the rebuttal PDF. We summarize here the main results.

• On QM9, END shows better agreement with the data distribution (except for QED which is captured perfectly by all methods). In particular, the reduction in strain energy demonstrates that END yields better geometries than the baselines (Table 1 in rebuttal PDF).
• On GEOM-Drugs, in addition to greatly improved connectivity, the SAScore, QED and logP are shown to be in better agreement with the data distribution (Table 3 in rebuttal PDF). The geometries are also shown to be of better quality as per the significantly reduced strain energy (Table 2 in rebuttal PDF).

Comparison to GeoLDM
We downloaded the GEOM-Drugs checkpoint available with the official implementation of GeoLDM, and evaluated the generated samples. We added one line with the obtained results in Table 2 in the rebuttal PDF, and summarize the findings here. While all samples generated by GeoLDM were effectively deemed “valid”, only around 46% were connected, against nearly 83% for END. The strain energy of the samples generated by END was also significantly lower, 55 kcal/mol vs 133 kcal/mol for GeoLDM, underscoring the advantages of a learnable forward process.

$1$ Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3D. ICML 2023.

We thank the reviewer for their constructive feedback. We are happy that our work was positively received by the reviewer.
We address the reviewer’s concerns/questions here below:

Ablation
The key component in our method is the learnable forward process. Hence, the logical ablation is whether to include a learnable forward (=END), or not (=EDM).

As stated in our global rebuttal, we made sure that all baselines shared the same amount of parameters as the models featuring a learnable forward to ensure that the difference in performance did not come from an increase in parameters. We did so by having 10 rounds of message-passing in EDM, against 5 in the forward and 5 in the reverse for END – all other things equal.

As pointed out by the reviewer, the ablated version where only the mean is learned (the standard deviation of the conditional marginal is pre-specified and derived from the same noise schedule as EDM) is shown to perform on par with the full version. We believe this is due to the relative simplicity of the task. We are currently running the same ablation on the more challenging GEOM-Drugs, and hope to be able to report the results before the end of the discussion period.

As nicely suggested by the reviewer, we conducted the same ablation in the conditional setting (composition-conditioned generation). The corresponding results are provided in the table here below – we will also add it to the final version. The ablated model is presented in the last row.

Regarding additional ablations, disabling equivariance in $F_\varphi$ could also be a possible option, previous work $1$ has however clearly shown that an architecture obeying the relevant symmetries is a useful inductive bias. We would be happy to hear the additional ablations that the reviewer would like to see in the final version of the paper.

Timing
In terms of timing, a training step on QM9 (common batch size of 64) takes on average ~0.37s (END) vs ~0.16s (EDM), this corresponds to a ~2.3x relative increase.
We observe the same trend on GEOM-Drugs (with an effective batch size of 64), a training step takes on average ~0.40s for END vs. ~0.15s for EDM (corresponding to a ~2.7x relative increase).
We do not observe that END requires a larger number of epochs to converge, compared to EDM. All models were trained for the same number of epochs / steps.

With our current implementation, sampling with END requires slightly less than 3x the time required by EDM on average. We report sampling times on QM9 (1024 samples) as an example, for varying numbers of integration steps.

| time $s$

However, we want to stress that END can achieve comparable (or better) accuracy with less integration steps. A concrete example of this can be seen in Table 2 in the rebuttal PDF (GEOM-Drugs), where with only 100 steps END yields better samples (connectivity and strain energy) than EDM with 1000 integration steps, i.e. in ~⅓ of the time.
As mentioned in the global rebuttal, we note that alternative parameterizations of the reverse process are possible. In particular, the drift of the reverse process $\\hat{f}\_{\\theta, \\varphi}$ could be learned without direct dependence on $f\_{\\varphi}$, thereby leading to very limited overhead with respect to vanilla diffusion models for sampling.

$1$ Equivariant Diffusion for Molecule Generation in 3D. ICML 2022.

We thank the reviewer for their constructive feedback. We are happy that our work was positively received by the reviewer.
We address the reviewer’s concerns/questions here below.

Performance Consistency
Regarding validity, we first note that (1) cheminformatics software implicitly adds hydrogens if an atom has a valence smaller than expected, (2) validity is usually computed on the largest connected fragment (thereby ignoring the remaining atoms).
On QM9, mol. stability gives a better picture, as it ensures that each atom exactly has the desired variance. Models including a learnable forward pass are shown to perform significantly better than baselines without. Improvements are also observed across other metrics (Table 1 in submission).
On the more challenging GEOM-Drugs, we believe that evaluating the model in terms of (validity x connectivity) is more realistic. According to that metric, END is clearly better than the baseline and samples far more connected molecules (Table 2 in rebuttal PDF).

GeoDiff
To the best of our knowledge, GeoDiff $1$ is a diffusion model for conformer generation, i.e. generation of coordinates given a molecular graph. Perhaps the reviewer had another model in mind?

Ablation
The key component in our method is the learnable forward process. Hence, the logical ablation is whether to include a learnable forward (=END), or not (=EDM).

To ensure that the difference in performance does not come from an increase in parameters, we made sure that all EDM baselines shared the same amount of parameters as the models featuring a learnable forward. We did so by having 10 rounds of message-passing in EDM, against 5 in the forward and 5 in the reverse for END – all other things equal.

We additionally provided another ablated version of END on QM9, where only the mean is learned, whereas the standard deviation of the conditional marginal is pre-specified and derived from the same noise schedule as EDM.
As suggested by another reviewer ** $jNQ2$ **, we conducted the same experiment in the conditional setting, where fixing the standard deviation is shown to lead to a small decrease in performance.
We are currently running the same ablation on GEOM-Drugs, and hope to be able to report the results before the end of the discussion period.

Complexity and Scalability
Complexity:
In terms of training, we did not observe that END required more epochs than the baseline to reach convergence. However, it is correct that each training step takes longer, and that the total training time is increased. All other things equal, we evaluated that increase to ~2.5x relative to EDM. In terms of resources, END could be trained on a single GPU for both datasets.

Regarding sampling, our current implementation incurs a ~3x increase relative to EDM.
However, we emphasize that END can be competitive with fewer integration steps. As a concrete example, on GEOM-Drugs, END with only 100 steps yields better samples than EDM with 1000 integration steps, i.e. in ~⅓ of the time.
Finally, alternative efficient parameterizations of the reverse process are possible. In particular, the drift of the reverse process $\\hat{f}\_{\\theta, \\varphi}(\\boldsymbol{z}\_t, t)$ could be learned without direct dependence on $f\_{\\varphi}$, thereby leading to a very limited overhead with respect to vanilla diffusion models for sampling.

Scalability: As all concurrent approaches, scaling to large systems is currently limited by the full connectivity of the message passing scheme.
An element specific to END, is that, due to the learnable forward, the prior is no longer forced to be $\\mathcal{N}(0, I)$ as in conventional diffusion models and can instead e.g. become size-specific. For large molecules, one could imagine scaling the prior with the number of atoms present in the system, i.e. $\\mathcal{N}(0, sI)$ where $s$ is a function of the number of atoms. Combined with a distance-based cutoff function to determine the neighborhood, this would allow for a more graceful scaling – as early steps of the denoising process would not result in fully-connected message passing as it is effectively the case with a simple $\\mathcal{N}(0, I)$ prior.

Another possibility to scale to systems such as large coordination complexes, is to design priors that “encode shapes”, e.g. square planar geometry could be built as a center and 4 groups of atoms representing the coordinated ligands, and thereby enabling more local message passing schemes.

$1$ GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation, ICLR 2022

We thank the reviewer for their constructive feedback. We address the reviewer’s concerns/questions here below:

Limited technical novelty
As mentioned in our general rebuttal, END is indeed a combination of existing ideas. We however want to emphasize that (1) our work is the first to seek improvement by adding a learnable forward process, orthogonally to most previous work that has mostly focused on the reverse process, (2) designing an appropriate transformation was not straightforward, it required some care as to preserve the desired invariance of the learned distribution.

Marginal performance gain
While we agree with the reviewer that the performance gains are limited on QM9, we have shown that END works considerably better than the considered baselines on GEOM-Drugs -- as summarized in the general rebuttal. Most notably, END significantly improves connectivity and geometry quality. It also shows better agreement with the training data on (newly added) drug-related metrics (SAScore, QED and logP).

Missing baseline GeoBFN
Unfortunately, since GeoBFN was published on arXiv ~2 months prior to the submission deadline a thorough comparison to this concurrent work is unfeasible. However, we extracted the relevant numbers presented in that paper and a direct comparison with GeoBFN in Tables 1 and 2 will be part of the final version. We summarize the comparison below.

On QM9, END generally performs on par with GeoBFN on the stability / validity metrics, across sampling steps. We could not evaluate the additional metrics as we do not have access to samples / pretrained models for GeoBFN.

On the more realistic GEOM-Drugs, END is shown to outperform GeoBFN in terms of atom stability (first row in the table below) – especially for small numbers of steps. As we can not check whether GeoBFN generates disconnected fragments, solely based on validity GeoBFN seems to perform better (second row in the table below).

It would be interesting to evaluate the connectivity / strain energy of the samples produced by GeoBFN for a full comparison.

Application to large systems
There is no restriction to the systems the END could be applied to.
As all concurrent approaches, scaling to large systems is currently limited by the full connectivity of the message passing scheme.
As for other methods treating categorical features as a continuous relaxation, we also note that for systems with a large number of atom types, alternatives to one-hot encoding (linear scaling wrt cardinality of space) might be beneficial. Methods such as Analog Bits $1$ (logarithmic scaling), or continuous lower-dimensional embedding such as that of GeoLDM $2$ are interesting approaches to deal with discrete features within END.

$1$ Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning, ICLR 2023.

$2$ Geometric Latent Diffusion Models for 3D Molecule Generation, ICML 2023

Thank you for the response. However I do believe the paper would be in better shape by additional efforts in some modifications e.g., adding comparison with advanced baselines (e.g., GeoBFN). I understand that GeoBFN was published ~2months prior to the deadline but the whole experiment would just take around several days and a fair comparison is expected, since you both work on the same dataset and benchmark and claim the advantage of fewer step sampling. Moreover, the work still seems technically incremental. I will slightly increase the score but still hold an opinion that the paper could be further strengthened by either showcasing strong performance against sota methods or exploring novel applications/use cases that constitutes a unique contribution.