LMDM:Latent Molecular Diffusion Model For 3D Molecule Generation

TL;DR

This work addresses 3D molecular generation with diffusion models by introducing LMDM, a latent diffusion framework operating in a $SE(3)$-equivariant latent space. It combines a molecular variational autoencoder (MVAE) with an EGNN-based encoder/decoder and a dual equivariant diffusion network to model short-range covalent bonds and long-range van der Waals interactions, aided by variational noise to boost diversity. Empirically, LMDM outperforms state-of-the-art methods (EDM, GeoLDM) on QM9 and GEOM-Drug in validity, uniqueness, novelty, and stability, especially for large molecules, and supports conditional generation across multiple properties with smooth property interpolation. The approach promises efficient, controllable 3D molecule generation with potential impact on drug design and materials science.

Abstract

n this work, we propose a latent molecular diffusion model that can make the generated 3D molecules rich in diversity and maintain rich geometric features. The model captures the information of the forces and local constraints between atoms so that the generated molecules can maintain Euclidean transformation and high level of effectiveness and diversity. We also use the lowerrank manifold advantage of the latent variables of the latent model to fuse the information of the forces between atoms to better maintain the geometric equivariant properties of the molecules. Because there is no need to perform information fusion encoding in stages like traditional encoders and decoders, this reduces the amount of calculation in the back-propagation process. The model keeps the forces and local constraints of particle bonds in the latent variable space, reducing the impact of underfitting on the surface of the network on the large position drift of the particle geometry, so that our model can converge earlier. We introduce a distribution control variable in each backward step to strengthen exploration and improve the diversity of generation. In the experiment, the quality of the samples we generated and the convergence speed of the model have been significantly improved.

Figures (7)

• Figure 1: Illustration of LMDM. We outline the training process of the proposed LMDM model. The encoder $\mathcal{E}_{\phi}$ coordinates $\mathrm{x}$ and molecular features $\mathrm{h}$ are encoded into equivariant latent variables $\mathrm{R}$,$\mathrm{A}$, and the time step encoding is used to incorporate the sequential information into the molecular information. We gradually add noise through the latent diffusion transformation $q(\mathcal{G}_{t} \mid \mathcal{G}_{t-1})$ until the latent variable distribution converges to a Gaussian distribution. Similarly, for the reverse generation process, the initial state $\mathcal{G}_{T} \sim \mathcal{N}(0, I)$ is gradually denoised by using the Markov kernel $\mathrm{p}_{\theta}(\mathcal{G}_{t-1} \mid \mathcal{G}_{t})$ and gradually refined by the equivariant denoising dynamics $\mathbf{\epsilon_\theta}(\mathcal{G}_{t}, \mathbf{t})$. The final latent variables $\mathrm{R}$, $\mathrm{A}$ are further decoded by the decoder $\mathcal{D}_{\epsilon}$ to generate the molecular point cloud.
• Figure 2: An overview of the one-stage molecular variational autoencoder. We encode the molecular structure through the encoder $\mathcal{E}_\phi$. Due to the unique properties of EGNN, the encoded latent variables still maintain equivariance on the SE(3) group action. The latent variables are sampled by reparameterization and then decoded by the decoder $\mathcal{D}_\epsilon$ to restore the latent variables to the original molecular structure. As in the first term of \ref{['MVAE:lossfunction']}, we use $d(\mathcal{G}, \mathcal{\hat{G}})$ to achieve the reconstruction loss, and in order to make the distribution of the latent variable $z_{x,h}$ closer to the prior distribution, making the distribution more regular and smooth, we added the $KL$ regularization term.
• Figure 3: Illustration of the specific implementation of the Markov kernel (double equivariant denoising score network). In fact, depending on the local and global edges of the input, we can use it as a local or global equivariant encoder to capture the molecular internal forces in the model and output the expected target score. The implementation of Schnet comes from schutt2017schnet.
• Figure 4: Molecules generated by LMDM trained on QM9 (left four) and DRUG (right two).
• Figure 5: Molecules generated by conditional LMDM. We conduct controllable generation with interpolation among different Polariz-ability $\alpha$ values with the same reparametrization noise $\epsilon$. The given $\alpha$ values are provided at the bottom.

Theorems & Definitions (1)

• Proposition 1