Summary
The paper introduces a pretraining method for molecule joint auto-encoding (MoleculeJAE) for 2D molecular topology and 3D molecular geometry. Their approach adopts SE(3) symmetry and is trained by fitting the joint distribution of the trajectories from the forward process of the diffusion model.
The authors treat the 2D molecular structure and 3D geometry as continuous objects, which are perturbed by Gaussian noise in the diffusion process. Unlike traditional diffusion models that denoise data by marginal distribution modeling, their 'auto-encoding' approach focuses on learning the joint distribution of data pair $(x_0, x_t)$. The overall task is then optimized by two objective functions, including a reconstruction loss and a contrastive loss, as well as a surrogate model.
Experimentally, MoleculeJAE is pre-trained on the PCQM4Mv2 dataset and is used to perform downstream tasks on QM9 and MD17 datasets. The experimental results are competitive on property prediction and dynamics prediction tasks.
Overall, this work is well-motivated and novel, with a theoretical justification for their approach. This paper presents good quality work.
Strengths
First of all, this work is well-motivated, as jointly learning 2D and 3D molecular representations by diffusion models is still at the stage of development. MoleculeJAE can learn both 2D bond topology and 3D conformation geometry information and is designed to respect the SE(3) symmetry of molecule data and is trained by fitting the joint distribution of the data’s augmented trajectories extracted from the forward process of the diffusion model.
It is also novel that the joint 2D and 3D learning of diffusion models is trained in a self-supervised learning manner and optimized by a contrastive learning objective function. MoleculeJAE unifies both contrastive and generative learning approaches from a trajectory perspective, providing a versatile and powerful molecular representation that can be applied to various downstream applications.
Unlike traditional diffusion modeling, MoleculeJAE uses a contrastive learning paradigm to fit the trajectory of pairwise molecules, while also utilizing the reconstructive task to perform denoising for each individual molecule. By using contrastive learning, it allows the model to align the augmented views of the same data and simultaneously contrast the augmented views of different data. Also, ablation studies show the importance to have a contrastive loss in the task.
Weaknesses
For methodology, the authors treat the 2D graph topology as a continuous object that is perturbed by Gaussian noise. I understand it is easier and reasonable to put the three things, atom features, bonds, and coordinates, under a single SDE framework. However, following previous works [1][2], describing the 2D graph topology as a discrete object with a discrete diffusion process is more reasonable, and their models show superior results.
Experimentally, MoleculeJAE shows better results with only small margins. However, the authors do not compare the number of model parameters, the overall training time, and the error bars (standard deviation) with existing methods, so it is unclear to me how well MoleculeJAE surpasses those methods.
[1] Vignac, Clement, et al. "DiGress: Discrete Denoising diffusion for graph generation." arXiv preprint arXiv:2209.14734 (2022).
[2] Hua, Chenqing, et al. "MUDiff: Unified Diffusion for Complete Molecule Generation." arXiv preprint arXiv:2304.14621 (2023).
Questions
1. Jointly learning 2D and 3D molecular data is not completely underdeveloped. Previous works on diffusion models, [1][2], have shown that diffusion models can be used to learn 2D and 3D latent molecular representations effectively jointly. They perform better on molecular generation tasks than single-modal diffusion models, and it is proper to cite these works and mention the difference.
2. I wonder if the authors consider extending this auto-encoding framework to discrete graph topology, as a discrete representation and diffusion process is more reasonable for graph-like objects. If not, I would like to hear the reason.
3. I would like to see the error bars (standard deviation), the number of model parameters, and the total training time compared with different methods.
[1] Hua, Chenqing, et al. "MUDiff: Unified Diffusion for Complete Molecule Generation." arXiv preprint arXiv:2304.14621 (2023).
[2] Vignac, Clement, et al. "Midi: Mixed graph and 3d denoising diffusion for molecule generation." arXiv preprint arXiv:2302.09048 (2023).
Typo:
1. Line 257, the log-likelihood logp is also achieves $\rightarrow$ the log-likelihood logp also achieves
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.