Flow Matching on Neural Fields for 3D Molecule Generation and Property-Based Guidance

Generative models can quickly design new molecules to assist in drug discovery. Point clouds and graphs are common representations for molecular generative models, but more informative representations can enable higher quality molecule generation. Neural fields represent 3D spaces and objects in a memory-efficient form. Applications of neural fields in molecule design have shown improvements in molecule quality over other representations. The choice of sampling method heavily contributes to model performance and molecule quality. Flow matching has emerged as a state-of-the-art sampling method for generating molecules. Neural fields have shown promise as an effective molecular representation but have not been directly measured against flow matching-based methods. In this work, we apply flow matching to a neural field code space to generate small molecules. We train neural fields based on voxelized versions of molecules in the QM9 dataset and then train a flow matching model to generate neural field codes which map to molecules. We compare our model performance against neural field-based methods with walk-jump sampling and graph-based methods with flow matching to assess molecule quality. We also enable property-guided generation for molecular neural fields with a two-part guiding and refining sampling procedure that optimizes for desired properties while promoting high validity in resulting molecules.