An End-to-End Framework for Molecular Conformation Generation via Bilevel Programming

TL;DR

ConfVAE addresses the challenge of generating realistic molecular conformations from graphs by modeling the conditional distribution $p({\bm{R}}|{\mathcal{G}})$ with an end-to-end bilevel CVAE. It employs a CNF-based distance-decoder to predict distance geometry and backpropagates through an inner distance-geometry solver to optimize alignments via RMSD, while preserving invariances. The approach achieves state-of-the-art results on GEOM benchmarks for both conformation generation and distance distribution modeling, with ablations confirming the necessity of the bilevel objective. This work provides a scalable and principled framework for 3D molecular generation and suggests potential extensions to larger biomolecules such as proteins.

Abstract

Predicting molecular conformations (or 3D structures) from molecular graphs is a fundamental problem in many applications. Most existing approaches are usually divided into two steps by first predicting the distances between atoms and then generating a 3D structure through optimizing a distance geometry problem. However, the distances predicted with such two-stage approaches may not be able to consistently preserve the geometry of local atomic neighborhoods, making the generated structures unsatisfying. In this paper, we propose an end-to-end solution for molecular conformation prediction called ConfVAE based on the conditional variational autoencoder framework. Specifically, the molecular graph is first encoded in a latent space, and then the 3D structures are generated by solving a principled bilevel optimization program. Extensive experiments on several benchmark data sets prove the effectiveness of our proposed approach over existing state-of-the-art approaches. Code is available at https://github.com/MinkaiXu/ConfVAE-ICML21

Figures (5)

• Figure 1: The overall framework of ConfVAE. At training time, given the graph ${\mathcal{G}}$ and conformation ${\bm{R}}$, we: 1) compute the distributions of $q_\phi(z|{\mathcal{G}}, {\bm{R}})$ and $p_\psi(z|{\mathcal{G}})$, and calculate ${\mathcal{L}}_{prior}$; 2) sample $z$ from $q_\phi$ by reparameterization, and feed it into the decoder (generator) $p_\theta$ to generate inter-atomic distances ${\bm{d}}$, after which we can obtain an auxiliary objective ${\mathcal{L}}_{aux}$ from the true distances $\hat{{\bm{d}}}$ derived from ${\bm{R}}$; 3) run the inner loop (distance geometry) to recover the 3D structure from ${\bm{d}}$, and compute the reconstruction RMSD loss ${\mathcal{L}}_{recon}$. The model is trained end-to-end by optimizing the sum of three object components ${\mathcal{L}}_{prior}$, ${\mathcal{L}}_{aux}$ and ${\mathcal{L}}_{recon}$.
• Figure 2: Schematic illustration of the forward and backward computational graph through the inner loop (distance geometry optimization). We repeatedly update ${\bm{R}}$ with the gradient $\nabla_{\bm{R}} H$ during the forward computation, and accumulate hypergradients $\nabla_{\theta,\phi}{\bm{R}}$ to update parameters $\theta$ and $\phi$ from backward computation.
• Figure 3: Visualization of generated conformations from state-of-the-art baselines, our method and the reference set, where four molecular graphs are randomly taken from the test set of GEOM-Drugs. C, O, H, S and Cl are colored gray, red, white, yellow and green respectively.
• Figure 4: Curves of the Coverage score with different thresholds $\delta$ on GEOM-QM9 (left two) and GEOM-Drugs (right two) datasets. The first and third curves evaluates the generated conformations from different generative models, while the other two are further optimized with the empirical force field.
• Figure 5: Marginal distributions $p(d_{uv}|{\mathcal{G}})$ of ground-truth and generated conformations from generative models. We study the edges between C and O atoms, and omit the H atoms for clarity. In each subplot, the annotation ($u-v$) denotes the corresponding atoms connected by the chemical bond $d_{uv}$.