Gradual Optimization Learning for Conformational Energy Minimization

TL;DR

This work presents the Gradual Optimization Learning Framework (GOLF) for energy minimization with neural networks that significantly reduces the required additional data and demonstrates that the neural network trained with GOLF performs on par with the oracle on a benchmark of diverse drug-like molecules.

Abstract

Molecular conformation optimization is crucial to computer-aided drug discovery and materials design. Traditional energy minimization techniques rely on iterative optimization methods that use molecular forces calculated by a physical simulator (oracle) as anti-gradients. However, this is a computationally expensive approach that requires many interactions with a physical simulator. One way to accelerate this procedure is to replace the physical simulator with a neural network. Despite recent progress in neural networks for molecular conformation energy prediction, such models are prone to distribution shift, leading to inaccurate energy minimization. We find that the quality of energy minimization with neural networks can be improved by providing optimization trajectories as additional training data. Still, it takes around $5 \times 10^5$ additional conformations to match the physical simulator's optimization quality. In this work, we present the Gradual Optimization Learning Framework (GOLF) for energy minimization with neural networks that significantly reduces the required additional data. The framework consists of an efficient data-collecting scheme and an external optimizer. The external optimizer utilizes gradients from the energy prediction model to generate optimization trajectories, and the data-collecting scheme selects additional training data to be processed by the physical simulator. Our results demonstrate that the neural network trained with GOLF performs on par with the oracle on a benchmark of diverse drug-like molecules using $50$x less additional data.

Figures (5)

• Figure 1: Mean squared error (MSE) of energy and forces prediction for NNPs trained on $\mathcal{D}_0, \mathcal{D}_{\text{traj-10k}}, \mathcal{D}_{\text{traj-100k}}, \mathcal{D}_{\text{traj-500k}}$. To compute the MSE, we collect NNP-optimization trajectories of length $T = 100$ and calculate the ground truth energies and forces on steps $t = 1, 2, 3, 5, 8, 13, 21, 30, 50, 75, 100$. Solid lines indicate the median MSE, and the shaded regions indicate the 10th and the 90th percentiles. Both the x-axis and y-axis are log scaled
• Figure 2: Violin plots of the percentage of optimized energy $\operatorname{pct}(s_T)$ calculated for various NNPs on $\mathcal{D}_{\text{test}}$ and $\mathcal{D}_{\text{test}}^{\text{SPICE}}$. Blue marks denote the mean percentage of optimized energy $\overline{\operatorname{pct}}_T$, the 10th, and the 90th quantile.
• Figure 3: $\overline{\operatorname{pct}}_t$ and $\operatorname{pct}^t_{\text{div}}$, $t = 1, 2, 3, 5, 8, 13, 21, 30, 50, 75, 100$. Shaded regions indicate the 10th and the 90th percentiles of the $\operatorname{pct}(s_t), s \in \mathcal{D}_{\text{test}}$ distribution. The x-axis is log-scaled.
• Figure 4: Mean squared error (MSE) of energy and forces prediction for NNPs trained on $\mathcal{D}_0$ and $\mathcal{D}_{\text{traj-500k}}$, and NNP trained with GOLF. To compute the MSE, we collect NNP-optimization trajectories of length $T = 100$ and calculate the ground truth energies and forces in steps $t = 1, 2, 3, 5, 8, 13, 21, 30, 50, 75, 100$. Solid lines indicate the median MSE, and the shaded regions indicate the 10th and the 90th percentiles. Both the x-axis and y-axis are log scaled
• Figure 5: Visualization of final conformations obtained by various models, the 2D view of the molecule, and the reference optimal conformation obtained with the $\mathcal{O}_G \ $.