Scalable Non-Equivariant 3D Molecule Generation via Rotational Alignment

We thank Reviewer bfzT for taking the time to review our paper and provide valuable feedback. We respond to the concerns as follows.

Not SOTA on QM9: Our paper mainly aims to improve non-equivariant diffusion models. On QM9, compared with the best non-equivariant baseline GraphLDM-aug, all three variants of our model improve the performance drastically, which validates the main claims of our paper. In fact, our best model ALDM$\_{\text{DiT-B}}$ achieves the highest Validity (94.2%) on QM9 across all (equivariant and non-equivariant) baselines and is only negligibly worse than GeoLDM on other metrics. We kindly request the reviewer to also consider the efficiency improvement shown in Table 3. Overall, our model achieves the best performance-efficiency trade-off on QM9. Additionally, on GEOM-DRUGS, which is much larger and more challenging than QM9, our model achieves the best performance.

Difference between Eq. 17 and regular GNN: Eq. 17 is indeed a regular GNN. We use it as a basic non-equivariant architecture (for the decoder and the noise prediction network of the diffusion model). Unlike EGNN, it simply concatenates atom coordinates and types and is thus non-equivariant. In Table 1, both ALDM$\_{\text{GNN}}$ and the baseline GraphLDM (from the GeoLDM paper) use Eq.17 as the noise prediction network, except that learned rotations are applied to the latent space of ALDM$\_{\text{GNN}}$. The performance gain validates that the aligned latent space can improve the training of non-equivariant diffusion models.

Alignment Quality Uncertainty: The lack of supervision is an inherent challenge for generative modeling. Our approach manages to learn rotations in such an unsupervised setting, which we believe is part of our contribution rather than a weakness of our model. Furthermore, the theme of our paper is to show that “alignment” is better than “no alignment” for non-equivariant diffusion models, which is already supported by our experiments. We agree with the reviewer that investigating the “optimality” of the aligned representations would be interesting, but we think it is currently beyond the scope of this paper. We’d love to explore this in future work.

Limited Theoretical Justification on Performance Trade-offs: In theory [1], diffusion models are able to learn the dataset distribution as long as the score function is learned well, without imposing any constraints (e.g., equivariance) on the target distribution. Therefore, non-equivariant diffusion models, in principle, have the full potential to learn the dataset distribution. On the contrary, previous equivariant diffusion models took equivariance for granted and didn’t justify why equivariance is necessary. Additionally, a few recent works (e.g., [2] (AlphaFold 3), [3]) also demonstrate the superior performance of non-equivariant diffusion models, though the specific problems they address are different from ours.

Code and detailed experimental settings: We already detailed the experimental setup in the main body of the paper (as mentioned by Reviewer t7B8: “The authors provided sufficient details about their experiment settings”). Specifically, we describe the dataset information in Section 5.1. We follow exactly the dataset preprocessing and splits of baselines EDM and GeoLDM and point to their public repositories. In Section 5.2, we introduce hyper-parameters (learning rate, number of layers, hidden dimension, batch size, etc.) we use before we present the experimental results. We will organize these information in a better format (e.g., using tables) in the next version. Our code will be released after the paper is accepted.

[1] https://arxiv.org/abs/2011.13456
[2] https://www.nature.com/articles/s41586-024-07487-w
[3] https://proceedings.mlr.press/v235/wang24q.html

We kindly request Reviewer bfzT to consider raising their score if they think our response can address the concerns. We are more than happy to answer any further questions!

We thank Reviewer Biv1 for taking the time to review our paper and provide valuable feedback. We are glad to see that they found our approach interesting and our results good enough to support our claims. We respond to the concerns as follows:

Lack of stronger baselines: We thank the reviewer for pointing this out! We will add more recent baselines (e.g., [1]) in an updated version of our paper. The performance of our non-equivariant models is still highly competitive compared with stronger equivariant baselines, and as the reviewer said, we believe the existing experimental results are sufficient to validate the effectiveness of our approach.

[1] https://arxiv.org/abs/2403.15441

Lack of analysis of learned alignment and latent representation. We provide more analysis below. First, for a single molecule, its canonical form can be obtained (e.g., by applying PCA to atom coordinates). However, the challenge is that such canonicalization may not have consistent effects for different molecules. For example, the directions (or signs) of the principal components obtained by PCA are ambiguous. Moreover, applying PCA to atom coordinates ignores the dependencies between coordinates and atom types, which should be generated together. Therefore, we believe it makes more sense to learn canonical forms across the training set by maximizing the overall reconstruction quality. As an ablation, we report the results based on PCA: PCA$\_{\text{GNN}}$ achieves 82.8% molecule stability on QM9 and 81.1% molecule stability on DRUGS. Its performance is better than the best non-equivariant baseline GraphLDM but still significantly worse than ALDM$\_{\text{GNN}}$ (87.4% on QM9 and 83.0% on DRUGS).

To analyze the rotations learned by the autoencoder, we randomly sample a mini-batch of molecules from the training set and compare their raw positions and the positions after rotations. We provide some visualization in the following anonymous link:

https://docs.google.com/document/d/e/2PACX-1vQeBYa4zljEoIM_h-4cvlOH-unSQFJhDgwVuq0FxtISIXZJ-T9T8q53OoWOKtoMR0d4O1SA7Vp12xq3/pub

We indeed found that common structural semantics (e.g., rings) of some molecules are easier to recognize after the alignment. Since our encoder is equivariant, such rotations are applied to the latent representations in the same way and thus make them more applied

We hope our response can address the reviewer's concerns. We are happy to answer any further questions.

We thank Reviewer t7B8 for taking the time to review our paper and provide valuable feedback. We respond to the concerns as follows.

Relation to canonicalization: We thank the reviewer for pointing this out! We admit that we overlooked the literature on canonicalization, and we will add them to the related work. Furthermore, we want to highlight some differences between our paper and existing works on canonicalization:

1. Existing canonicalization works mostly consider the supervised learning setting, while our paper studies a generative modeling problem.
2. The main computational cost of equivariant baselines is the training of diffusion model (e.g., ~3000 epochs on QM9). Our approach uses an autoencoder to learn canonicalization in a few epochs before training the diffusion model. The additional overhead from learning canonicalization is marginal compared with the great speedup of training diffusion models (Table 3).
3. While theoretical works have analyzed the sample complexity gain of equivariance under the supervised learning setting, their conclusions may not be trivially extendable to generative modeling. Empirically, recent works (e.g., notable AlphaFold 3) also show the good performance of non-equivariant generative models. We believe this problem is worth further investigation, and our work is one step in this direction.
4. Regarding the smoothness of canonicalization, the rotation estimation by Eq (16) is smooth except when det(M)=0, which doesn’t notably impede training in practice [1]. Besides, we conjecture that the discontinuity issue of canonicalization is less of a concern in our setting, because we don't test canonicalization on unseen data and only use it to reshape the training data of which we have full control.

Ablations on architectural choices: Our experiments already contain ablations. Specifically, our baseline GeoLDM uses equivariant networks for both encoder and decoder. We use the same equivariant encoder as GeoLDM to keep the architectural factors affecting the generation of latent representations the same, except that our model applies the learned rotations. Next, to learn rotations, our decoder has to be non-equivariant; otherwise, the reconstruction loss for atom coordinates (L2 loss) would be invariant to rotations and prevent the rotation network from learning anything. Note that although we have an equivariant encoder, it is only used to generate latent representations, and the probability distribution is learned by a non-equivariant diffusion model. This doesn't violate the claim of the paper.

For the ablation results, we kindly refer the reviewer to the performance of GraphLDM and ALDM$\_{\text{GNN}}$ in Table 1. GraphLDM is a non-equivariant baseline from the GeoLDM paper that replaces the original equivariant denoising network of the diffusion model with a non-equivariant network (Eq 17). ALDM$\_{\text{GNN}}$ also uses Eq (17) as the denoising network for diffusion and the same hyper-parameters (noise schedule, diffusion steps, etc.). As a result, the significant performance gain of ALDM$\_{\text{GNN}}$ comes from the rotations applied to the latent space, which validates the effectiveness of our proposal. We also tried to use DiT for GraphLDM, but the gap was still significant.

Use of a canonicalization algorithm (PCA): We considered PCA in our preliminary experiments. There are two issues with PCA. First, the directions (or signs) of the principal axes are ambiguous. For prediction tasks, frame averaging can address this ambiguity, but it's not applicable to generative modeling. Second, applying PCA to atom coordinates ignores the dependencies between coordinates and atom types, which should be generated together. To further address the reviewer’s concern, we report here the results based on PCA: PCA$\_{\text{GNN}}$ achieves 82.8% molecule stability on QM9 and 81.1% molecule stability on DRUGS. Its performance is better than the best non-equivariant baseline GraphLDM but still significantly worse than ALDM$\_{\text{GNN}}$ (87.4% on QM9 and 83.0% on DRUGS)

Use of an equivariant canonicalization network: We agree that using an equivariant canonicalization network (the rotation network in the paper) makes the canonical form of a molecule rotation invariant. However, to learn the specific canonical forms for different molecules that facilitate the reconstruction of the training set, the canonicalization network still needs to be trained jointly with the VAE (and probably less efficient). We'll try it in the updated version. Since our paper mainly aims to improve non-equivariant diffusion models, we believe the existing results are sufficient to support our claim.

[1] https://arxiv.org/abs/2006.14616

We hope our response can provide a clearer explanation of the key points of our paper and kindly request Reviewer t7B8 to consider revising their score if they think our response can address their concerns. We are happy to answer any further questions.

We thank Reviewer uMCt for taking the time to review our paper and provide valuable feedback. We respond to the concerns as follows:

Scalability of non-equivariant diffusion models: In fact, we did experiments to show the scalability of our model. We kindly refer the reviewer to Table 3 on page 8, where we show the model sizes. We tested non-equivariant denoising networks of three different sizes for the diffusion model: GNN is from Eq (17), DiT-S is the small version model from the diffusion transformer paper, and DiT-B is the base version with a larger hidden dimension (4 times #params as DiT-S and more than 10 times #params as GeoLDM). Our largest model achieves the best performance on the DRUGS dataset and still maintains higher training/sampling efficiency than existing equivariant diffusion models.

Why is non-equivariant decoder necessary to learn alignment?: Let $x$ denote the original atom coordinates and $x' = \mathcal{D}(\mathcal{E}(x))$ denote the reconstructed coordinates. The reconstruction error reduces to L2 loss $\Vert x' - x \Vert^2$ for continuous coordinates. Now suppose $x$ is rotated by $R$, i.e., $Rx$. If both the encoder and decoder are equivariant, the reconstruction would become $Rx'$, and the L2 loss between $Rx$ and $Rx'$ would be the same as $\Vert x' - x\Vert^2$ (because $R$ is orthogonal), making the rotation network unable to receive any supervision signals. Therefore, encoder and decoder should not both be equivariant.

We chose to make the decoder non-equivariant and keep the encoder equivariant because we wanted to compare directly with the baseline GeoLDM. GeoLDM uses equivariant networks for both encoder and decoder. We use the same equivariant encoder as GeoLDM to keep the architectural factors affecting the generation of latent representations the same, except that our model applies the learned rotations.

Ablations on the architecture and latent space. Our experiments already contain ablations. We kindly refer the reviewer to the performance of GraphLDM and ALDM$\_{\text{GNN}}$ in Table 1. GraphLDM is a non-equivariant diffusion model baseline from the GeoLDM paper that replaces the original equivariant denoising network of the diffusion model with a non-equivariant network (Eq 17). ALDM$\_{\text{GNN}}$ also uses Eq (17) as the denoising network for diffusion and the same hyper-parameters (noise schedule, diffusion steps, etc.). Therefore, the difference between ALDM$\_{\text{GNN}}$ and GraphLDM lies in that the latent space of ALDM$\_{\text{GNN}}$ is aligned by rotations. The experimental results show a significant performance gain of ALDM$\_{\text{GNN}}$ compared with GraphLDM, which validates that the aligned representations improve the learning of non-equivariant diffusion models. We also tried to use DiT for GraphLDM, but the gap was still significant.

Additionally, we randomly sample a mini-batch of molecules from the training set and compare their original and rotated positions. The visualization is contained in the following anonymous link: https://docs.google.com/document/d/e/2PACX-1vQeBYa4zljEoIM_h-4cvlOH-unSQFJhDgwVuq0FxtISIXZJ-T9T8q53OoWOKtoMR0d4O1SA7Vp12xq3/pub We indeed found that common structural semantics (e.g., rings) are easier to recognize after rotation. Since our encoder is equivariant, these rotations are applied to the latent representations in the same way and thus make them more aligned.

Novelty of the proposed method. While the evaluation of novelty depends on one's own perspective, we want to emphasize that our paper is not simply “a change of architecture”. In fact, in this paper, we switch from equivariant models to non-equivariant models and propose an approach to close the gap between them. As we show in the experiments (explained above), the improvement brought by our approach is significant even if we use the same non-equivariant architecture (i.e., a basic GNN) for the diffusion model.

Besides, our model not only improves the sampling speed, but also significantly saves training cost. On QM9, training the baseline equivariant diffusion model (e.g., GeoLDM) to convergence takes nearly 3000 epochs, which needs 4 days to finish on a single 4090 GPU (in Table 3 on page 8 we show the average training time per epoch). As a comparison, the training of our largest model ALDM$\_{\text{DiT-B}}$ takes 3 days on a 4090 GPU. We use a larger batch size and thus more epochs for ALDM$\_{\text{DiT-B}}$, but the total training cost is still reduced significantly. We believe this provides more support for our model.

We hope our response can provide a clearer explanation of the key points of our work and kindly request Reviewer uMCt to consider revising their score if they think our response can address their concerns. We are happy to answer any further questions.