Mixed Continuous and Categorical Flow Matching for 3D De Novo Molecule Generation

TL;DR

FlowMol advances 3D de novo molecule generation by extending flow matching to mixed continuous and categorical data, enabling joint sampling of atom positions, types, charges, and bond orders. It introduces SimplexFlow to handle categorical variables on the probability simplex, along with an endpoint-parameterized objective, optimal transport alignment, and a SE(3)-aware GVP-based architecture. Empirical results show Pure Gaussian priors with endpoint parameterization outperform simplex-constrained variants, FlowMol achieves competitive accuracy with diffusion models while offering markedly faster inference, and SimplexFlow does not consistently improve performance. The work raises important questions about prior design in flow matching for mixed-data tasks and provides code and models for reproducibility, positioning FlowMol as a practical, scalable alternative for 3D molecular generation.

Abstract

Deep generative models that produce novel molecular structures have the potential to facilitate chemical discovery. Diffusion models currently achieve state of the art performance for 3D molecule generation. In this work, we explore the use of flow matching, a recently proposed generative modeling framework that generalizes diffusion models, for the task of de novo molecule generation. Flow matching provides flexibility in model design; however, the framework is predicated on the assumption of continuously-valued data. 3D de novo molecule generation requires jointly sampling continuous and categorical variables such as atom position and atom type. We extend the flow matching framework to categorical data by constructing flows that are constrained to exist on a continuous representation of categorical data known as the probability simplex. We call this extension SimplexFlow. We explore the use of SimplexFlow for de novo molecule generation. However, we find that, in practice, a simpler approach that makes no accommodations for the categorical nature of the data yields equivalent or superior performance. As a result of these experiments, we present FlowMol, a flow matching model for 3D de novo generative model that achieves improved performance over prior flow matching methods, and we raise important questions about the design of prior distributions for achieving strong performance in flow matching models. Code and trained models for reproducing this work are available at https://github.com/dunni3/FlowMol

Figures (4)

• Figure 1: Overview of FlowMolTop: We adapt the flow matching framework for unconditional 3D molecule generation. An ordinary differential equation parameterized by a graph neural network transforms a prior distribution over atom positions, types, charges, and bond orders to the distribution of valid molecules. Black arrows show the instantaneous direction of the ODE on atom positions. Middle: Trajectory of the atom type vector for a single atom under SimplexFlow, a variant of flow matching developed for categorical variables. Atom type flows lie on the probability simplex. Bottom: Trajectory of an atom type vector starting from a Gaussian prior. This approach does not respect the categorical nature of the data; however, we find it yields superior performance to SimplexFlow.
• Figure 2: FlowMol ArchitectureTop left: An input molecular graph $g_t$ is transformed into a predicted final molecular graph $g_1$ by being passed through multiple molelcule 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$ is used to denote MLPs and GVPs, respectively.
• Figure 3: Interpolant Schedules for the QM9 Dataset
• Figure 4: Left: maximum integration step size to remain on the simplex for a cosine interpolant. Right: zoomed in view of the asymptotic decline of the maximum step size as $t \to 1$