Exploring Discrete Flow Matching for 3D De Novo Molecule Generation

TL;DR

This work benchmark the performance of existing discrete flow matching methods for 3D de novo small molecule generation and provides explanations of their differing behavior, and presents FlowMol-CTMC, an open-source model that achieves state of the art performance for 3D de novo design with fewer learnable parameters than existing methods.

Abstract

Deep generative models that produce novel molecular structures have the potential to facilitate chemical discovery. Flow matching is a recently proposed generative modeling framework that has achieved impressive performance on a variety of tasks including those on biomolecular structures. The seminal flow matching framework was developed only for continuous data. However, de novo molecular design tasks require generating discrete data such as atomic elements or sequences of amino acid residues. Several discrete flow matching methods have been proposed recently to address this gap. In this work we benchmark the performance of existing discrete flow matching methods for 3D de novo small molecule generation and provide explanations of their differing behavior. As a result we present FlowMol-CTMC, an open-source model that achieves state of the art performance for 3D de novo design with fewer learnable parameters than existing methods. Additionally, we propose the use of metrics that capture molecule quality beyond local chemical valency constraints and towards higher-order structural motifs. These metrics show that even though basic constraints are satisfied, the models tend to produce unusual and potentially problematic functional groups outside of the training data distribution. Code and trained models for reproducing this work are available at \url{https://github.com/dunni3/FlowMol}.

Figures (5)

• Figure 1: OverviewTop: We adapt the flow matching framework for unconditional 3D molecule generation and explore the use of different discrete flow matching methods. This CTMC trajectory shows masked atoms in gray. Middle: Trajectory of the atom type vector for a single atom under SimplexFlow, a variant of continuous flow matching developed for categorical variables. Atom type flows lie on the probability simplex. Bottom: Trajectory of an atom type vector for a CTMC flow. Atom types jump between the mask state and real atom types.
• Figure 2: Atom Type Assignment Times: Cumulative Density Functions (CDFs) of the time at which an atom is assigned its final atom type, for each DFM method tested. Green lines show the time of final atom type assignments in $g_t$. Gold lines show the times when the final atom type is assigned in $\hat{g}_1(g_t)$ (the predicted final molecule given the current molecule at time $t$).
• Figure 3: Prior distributions used with support on the probability simplex
• Figure 4: FlowMol ArchitectureTop left: An input molecular graph $g_t$ is transformed into a predicted final molecular graph $g_1$ by being passed through multiple molecule update blocks. Top right: A molecule update block uses NFU, NPU, and EFU sub-components to update all molecular features. Bottom: Update equations for graph features. $\phi$ and $\psi$ denote MLPs and GVPs, respectively.
• Figure 5: Atom type trajectories and associated times of final type assignment. Each plot contains the atom type trajectory for a single atom. Subplot columns correspond to the trajectory type: either a $\hat{g}_1(g_t)$ or $g_t$ trajectory. Subplot rows correspond to the DFM variant used by the model. Within a row, the $\hat{g}_1(g_t)$ and $g_t$ trajectories displayed are for the same atom. The time of final atom type assignment is overlaid in red on each trajectory.