Score-based 3D molecule generation with neural fields

We thank the reviewer for the comments and questions. A general rebuttal is posted above. Below we address the reviewer's individual questions.

1. MolDiff Baseline. MolDiff shows that incorporating bond information in point cloud-based approaches improves the quality of the generated samples. This contribution is orthogonal to ours and can potentially be incorporated into our generative model, e.g. via additional channels. This, however, is not the focus of our work, nor that of the baselines we considered. Here, we aim at proposing for the first time neural fields as a new representation for 3D molecules (a non-trivial task).

For completeness (despite different training assumptions), we compare FuncMol to MolDiff. Since the weights for MolDiff with hydrogens were unavailable, we compared FuncMol using MolDiff’s metrics and the MolDiff performance reported in their Appendix D.1, Table 8. We observe that FuncMol achieves competitive results in most metrics despite not using bond information.

2. RINGER Baseline. To our knowledge, we are the first to report results on CREMP in the unconditional all-atom 3D molecule generation setting. This experiment is more for qualitative purposes as CREMP is very recent and not a standard benchmark for this task. RINGER [1] also considered CREMP but tackled conformer generation, a different problem from ours: it assumes knowledge of the molecular sequence/graph during training and sampling. This simplifies generation, since the model knows a priori the number of atoms, their types, the bonds between them and the approximate atom locations. RINGER only parametrizes angles and torsions, while FuncMol and its baselines perform all-atom generation. We tried to extend our baselines to CREMP but did not succeed, mainly due to the high memory consumption (e.g. VoxMol took 40 GPU-hours per epoch, while FuncMol took 2.7 GPU-hours). To encourage comparisons to FuncMol, we include some quantitative metrics used in [1] that measure the distance between test and generated distributions using KL divergence. The KL divergence for bond angles are 0.1615 ($\theta_1$), 0.1345 ($\theta_2$) and 0.2197 ($\theta_3$); the ones for the dihedral angles are 0.1127 ($\phi$), 0.1178 ($\psi$) and 0.1813 ($\omega$). The percentage of valid generated peptides (for which we can extract the sequence of amino acids from their SMILES) is 82.7%.

3. Focus on unconditional generation, which is not as good as conditional generation. We agree that conditional generation presents more practical value than unconditional generation. We note, however, that our approach is the first to introduce neural fields as a representation of 3D molecules. The focus of our submission is to show that (i) it is a feasible approach (this is non-trivial), (ii) it scales well with data and molecule size, and (iii) it achieves competitive results with fast sampling time. We will include a discussion about this topic on the appendix and we will leave extensions to conditional generation for future work (similar to how TargetDiff [1]/DiffSBDD [2] extend EDM or VoxBind [3] extends VoxMol).

4. The model is not SE(3) equivariant. How to leverage rotation/translation? We do not leverage any SE(3) equivariance on the approach described in the original manuscript. We acknowledge that this is a good point and address it in the rebuttal by performing data augmentation during training (see main Rebuttal comment "Issues with memorization"). Data augmentation improves the quality of the representations as it helps learn a "more semantic" latent space z. This is reflected in the empirical results of Rebuttal Tables 1 and 2, where the uniqueness score significantly improves.

5. Ablation study missing: Impact of different discretization steps. Please, see the main Rebuttal, where we address this comment. Appendix D of the manuscript contains four other thorough ablation studies that justify other design choices in FuncMol.

[1] Guan et al. "3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction". ICLR23

[2] Schneuing et al. "Structure-based Drug Design with Equivariant Diffusion Models". Arxiv22

[3] Pinheiro et al. "Structure-based drug design by denoising voxel grids" ICML24

I would like to thank the authors for their answers on my inquiries, especially concerning the ablation study and the comparative analysis with the MolDiff model. I hope, that all provided metrics and the new baselines will be included in the final version. Additionally, it would be beneficial to include results from the MolDiff model trained on a hydrogen-depleted graph. I raise my score.

We thank the reviewer for the comments and positive feedback. A general rebuttal is posted above. Below we address the specific issues raised by the reviewer.

“Why does FuncMol not outperform GeoLDM and VoxMol?" It is challenging to explain why these models perform differently, as they make different choices w.r.t. data representation, NN architectures and generative modeling. FuncMol outperforms GeoLDM in most of the metrics of Rebuttal Table 2 and some metrics of Rebuttal Table 1 (c.f. main rebuttal). VoxMol outperforms FuncMol on GEOM-drugs (the performance gap is smaller when adding data augmentation to FuncMol). VoxMol uses 3D UNets, an expressive and well-established NN architecture tailored to discrete data, while FuncMol uses relatively new neural field architectures; neural fields tend to underperform against discrete-data architectures in many computer vision applications [1]. FuncMol, however, generates 3D molecules an order of magnitude faster than VoxMol and GeoLDM. Moreover, unlike them, FuncMol easily scales to larger molecules (such as macro-cyclic peptides).

Scaling laws for molecules. We agree that coming up with scaling laws for 3D molecules is an interesting research direction. However, this lies beyond the scope of this work, especially when taking into account the large amount of computation resources required for this type of study. That being said, we observed that at fixed dataset, increasing the parameter count of our models (both the neural field and our denoiser) significantly improved our results. Moreover, increasing the dataset e.g. with data augmentation leads to better performance.

Downstream applications: "Molecule generation may not be a perfect scenario for this representation". We thank the reviewer for this comment that helped improve the paper. Generative modeling for molecules is an important task with many practical applications. It also serves as a testbed to demonstrate that neural fields provide a very expressive/efficient/scalable decoder for 3D molecules, unlike point clouds or voxel grids.

At the reviewer’s request, we show that neural field representations can also be used in other downstream tasks, e.g. discriminative tasks related to property prediction. We consider some properties from QM9 ($\mu$ dipole moment, $u_0$ internal energy, $\alpha$ isotropic polarizability, $C_v$ heat capacity) and report the Spearman correlation performance (ranging from -1 to 1) in the linear probing setting (i.e. training a linear regression on top of frozen codes). We observe high ranking scores and correlations between predictions and ground truth labels as shown on Rebuttal Fig 4 (c.f. main rebuttal pdf) and the table below.

Why use walk-jump sampling? We explain why we use walk-jump sampling/neural empirical Bayes on L40-41: "[it] enjoys many properties such as fast-mixing, simplicity for training and fast sampling speed" as it relies on a single noise level. It has been successfully applied in 3D molecule generation (e.g., VoxMol) and antibody sequence generation [2]. Despite choosing walk-jump sampling for our main results, we stress that our framework is agnostic to the generative model used; we reported results with diffusion models on Appendix C. Compared to diffusion, we found that walk-jump sampling on this task is simpler/faster to train, faster to sample from (less sampling steps are required) and achieves slightly better empirical results.

How to guarantee the generalization of learned "codes"? Random initialization and Langevin MCMC may produce invalid molecules. We agree with the reviewer that the manifold of valid molecules is very complicated (low-dimensional in a large ambient space and potentially non-smooth). However, the manifold of codes after gaussian-smoothing is much simpler, especially with large noise. This allows us to perform Langevin MCMC from random initialization. This is precisely the intuition of neural Empirical Bayes [3], the generative framework we used in this work. We observed empirically that Langevin MCMC on smoothed codes mixes well. For example, Rebuttal Fig. 5 (c.f. main rebuttal pdf) shows two single MCMC chains initialized randomly with different seeds, where molecules are generated after each 200 "walk" steps. This phenomenon has also been observed on different data modalities, e.g., images [4], biological sequences [2] and voxelized molecules (VoxMol).

[1] Dupont et al. "From data to functa: Your data point is a function and you can treat it like one." ICML2022.

[2] Frey et al. "Protein Discovery with Discrete Walk-Jump Sampling". ICLR24

[3] Saremi et al.. "Neural empirical Bayes". JMLR19

[4] Saremi et al.. “Multi-measurement Generative Models”. ICLR22

We thank the reviewer for the valuable feedback. A general rebuttal is posted above. Below we address the reviewer's specific questions.

Issues with overfitting/memorization. By further investigating this issue, we realized that memorization is likely caused by a degenerate latent space—as alluded to by the reviewer. Data augmentation helped mitigate this issue and learn a more "semantically meaningful" latent space. This is reflected in the results of Rebuttal Table 1 and 2 (see main Rebuttal "Issues with memorization"), where uniqueness is significantly improved and we no longer observe memorization. This is a consequence of the model overfitting less, which also allowed us to train denoisers with higher capacity.

Worse performance on QM9. We agree that QM9 is a good task to test generative models. FuncMol and VoxMol perform comparably on QM9 and both are outperformed by point cloud approaches. FuncMol, however, has an order of magnitude faster sampling time. QM9 has limited relevance for drug discovery: the number of stable and unique molecules is limited as QM9 is an enumeration of all molecules up to 9 heavy atoms satisfying some constraint. Achieving a high performance on QM9 does not mean the model can handle more complex distributions. In fact, the best models we reported on QM9 (GeoLDM, EDM) perform the worst on GEOM-drugs.

Quantitative results on CREMP. To our knowledge, we are the first to report results on CREMP in the unconditional all-atom 3D molecule generation setting. This experiment is more for qualitative purposes as CREMP is very recent and not a standard benchmark for this task. RINGER [1] also considered CREMP but tackled conformer generation, a different problem from ours: it assumes knowledge of the molecular sequence/graph during training and sampling. This simplifies generation, since the model knows a priori the number of atoms, their types, the bonds between them and the approximate atom locations. RINGER only parametrizes angles and torsions, while FuncMol and the baselines perform all-atom generation. We tried to extend our baselines to CREMP but did not succeed, mainly due to the high memory consumption (e.g. VoxMol took 40 GPU-hours per epoch, while FuncMol took 2.7 GPU-hours). To encourage comparisons to FuncMol, we include some quantitative metrics used in [1] that measure the distance between test and generated distributions using KL divergence. The KL divergence for bond angles are 0.1615 ($\theta_1$), 0.1345 ($\theta_2$) and 0.2197 ($\theta_3$); the ones for the dihedral angles are 0.1127 ($\phi$), 0.1178 ($\psi$) and 0.1813 ($\omega$). The percentage of valid generated peptides (for which we can extract the sequence of amino acids from their SMILES) is 82.7%.

Sampling time on CREMP. We reported the sampling time on L327: "our model takes around 1.4s to generate a molecule. [..] should VoxMol be trained successfully, it would take over a minute [per] molecule".

Why switch from the meta-learning approach in [6]? Joint optimization, a.k.a. auto-decoding has been applied in several works [2, 5] to learn signed-distance functions. Auto-decoding is attractive as it does not require the memory-expensive double loop optimizations in meta-learning and is simpler to optimize. We attribute the novelty/uniqueness issue to the lack of data augmentation during training (see "Issues with memorization" in the main Rebuttal).

Filtering samples due to memorization. As written on L261: "we only consider the novel molecules of FuncMol$_{drop=0}$ and exclude those pathological chains when benchmarking this model"; the total number of molecules that are evaluated is 10,000. The other reported FuncMol model (with dropout) considers both novel and non novel molecules. Our initial novelty scores are 56.3% for FuncMol$_{drop=0}$ and 77.3% for FuncMol (with dropout). With data augmentation, the novelty score is over 99%.

How can the percentage of valid molecules be higher than the percentage of stable molecules? We use the same metrics as previous work for an apples to apples comparison to other work. In fact, all baselines that report these two metrics have higher validity than mol stability (EDM, GeoLDM, VoxMol). Validity is measured as the success rate of RDKit sanitization and is computed only on the largest fragment of the generated molecules (when they consist of disjoint fragments). Atom and molecule stability are defined as the reviewer described and were introduced in [4]. As explained by [3]: "[mol stability] is similar to validity but, in contrast to RDKit sanitization, they do not allow for adding implicit hydrogens to satisfy the valency constraints."

[1] Grambow et al. "RINGER" Arxiv23

[2] Park et al. “DeepSDF” CVPR19

[3] Vignac et al. "MiDi" ECML23

[4] Satorras et al. "E(n) equivariant normalizing flows" NeurIPS21

[5] Chou et al. "Diffusion-SDF" ICCV23

[6] Dupont et al. "From data to functa" ICML22

We thank the reviewer for the positive and valuable feedback. A general rebuttal is posted above. Below we address the reviewer's concerns.

1. Meaningful representation of latent space. See "Meaningful representation of latent space [yg8b]" on the main Rebuttal. Overall, the additional experiments show that the latent space is well structured.

2. Peak finding + iterative refinement can be time consuming. To clarify, the average sampling time reported on Table 2 is the time required to fully generate the molecules ("from noise to set of atom types/coordinates"). It includes four steps: (i) the walk-jump sampling steps to sample latent codes, (ii) rendering (going from the code to the voxel grid at resolution .25A), (iii) the peak finding and (iv) the iterative refinement. The sampling time for all steps together is an order of magnitude faster than alternative baselines. The bottleneck in sampling time is the rendering phase (ii). The sampling time can be reduced at the cost of performance by rendering codes at a lower resolution (see "Ablation study" on the main Rebuttal).

3. Explicitly modeling bonds between atoms. Some recent works e.g., MolDiff, MiDi, show that incorporating extra information such as bonds and formal charges in point cloud-based approaches (e.g. EDM) improves the quality of the generated samples. These contributions are orthogonal to ours and can potentially be incorporated into our generative model, e.g. via additional channels as suggested by the reviewer. This is, however, not the focus of our work (nor that of the baselines we considered). Here, we aim at proposing for the first time neural fields as a new representation for 3D molecules (a non-trivial task).

For completeness (despite different training assumptions), we compare FuncMol to MolDiff. MolDiff only incorporates bond information into the diffusion process, making it a simple representative baseline for this class of model. Since the weights for MolDiff with hydrogens were unavailable, we compared FuncMol using MolDiff’s metrics and the MolDiff performance reported in their Appendix D.1, Table 8. We observe that FuncMol achieves competitive results in most metrics despite not leveraging bond information during training.