Unified Generative Modeling of 3D Molecules via Bayesian Flow Networks

TL;DR

This work tackles the challenge of generating accurate 3D molecular geometries with diffusion-based methods by introducing GeoBFN, a Bayesian flow framework that operates in the differentiable parameter space to model SE(3) invariant densities over atom coordinates and features. By enforcing equivariance through an EGNN-based updater and a continuous-time loss, GeoBFN unifies multiple molecular modalities (coordinates, discretized charges, and atom types) and achieves state-of-the-art generation quality on QM9 and GEOM-DRUG benchmarks. A key contribution is the ability to sample at any step count, enabling substantial speedups (around $20\times$) without sacrificing performance, and to improve controllable generation via a probabilistic, symmetry-aware update mechanism. The framework is generalizable to other molecular tasks and offers a scalable, symmetry-preserving approach to 3D molecular generation with strong empirical performance and efficiency benefits.

Abstract

Advanced generative model (e.g., diffusion model) derived from simplified continuity assumptions of data distribution, though showing promising progress, has been difficult to apply directly to geometry generation applications due to the multi-modality and noise-sensitive nature of molecule geometry. This work introduces Geometric Bayesian Flow Networks (GeoBFN), which naturally fits molecule geometry by modeling diverse modalities in the differentiable parameter space of distributions. GeoBFN maintains the SE-(3) invariant density modeling property by incorporating equivariant inter-dependency modeling on parameters of distributions and unifying the probabilistic modeling of different modalities. Through optimized training and sampling techniques, we demonstrate that GeoBFN achieves state-of-the-art performance on multiple 3D molecule generation benchmarks in terms of generation quality (90.87% molecule stability in QM9 and 85.6% atom stability in GEOM-DRUG. GeoBFN can also conduct sampling with any number of steps to reach an optimal trade-off between efficiency and quality (e.g., 20-times speedup without sacrificing performance).

Figures (3)

• Figure 1: The framework of GeoBFN
• Figure 2: Graphical View of Comparison between BFN and Diffusion
• Figure 3: The Bayesian Flow and Diffusion Process of GeoBFN and EDM.

Theorems & Definitions (3)

• Theorem 3.1
• Proposition 3.2
• Remark 3.3