GeoBFN

Adding Gaussian noise to an image degrades it gracefully; adding it to a molecule's atom coordinates does not. Displacing one atom changes its distance to every other atom, so a perturbation small enough to look harmless can push a carbon-carbon distance outside any plausible bond length, erasing the bonded structure the point cloud encodes. Diffusion models for 3D molecules inherit that fragility: their intermediate states carry little usable chemical information, and most of the molecule materialises only in the final few denoising steps. Molecular geometry is also multi-modal, pairing continuous coordinates with discrete element types that one Gaussian noise process cannot describe.

GeoBFN answers both by moving generation out of sample space entirely. Built on Bayesian flow networks, it maintains a collection of independent distributions over the clean molecule — Gaussian means and precisions for coordinates, categorical probabilities for atom features — and generates by applying Bayesian updates to those parameters as an equivariant network sharpens its reconstruction. Every modality becomes a point in the same differentiable parameter space, whose input variance is far lower than the sample space's, so intermediate states stay chemically meaningful throughout. Making the inter-dependency modeling over parameters equivariant preserves SE(3)-invariant density modeling, leaving likelihood unchanged by rotating or translating the molecule.

The work comes from the Institute for AI Industry Research (AIR) at Tsinghua University with the University of Illinois Urbana-Champaign and the Shanghai Institute of Materia Medica, presented as an ICLR 2024 Oral. It is the parent formulation for MolCRAFT, which carries the same parameter-space generation into pocket-conditioned structure-based drug design, and for the MolJO and MolPilot models built on it.

#Key Features

• Generation in parameter space: Bayesian update, sender and receiver distributions replace the forward and reverse noising of diffusion, giving a smoother transformation of molecular geometry and a lower-variance input to the network.
• Equivariant inter-dependency modeling: An SE(3)-equivariant graph neural network models the dependency among distribution parameters, so the implied density is roto-translation invariant and the generative process equivariant.
• Any-step sampling: A continuous-time training loss lets the number of sampling steps be chosen at inference, with no retraining and no extra training overhead.
• Redundancy-free atom representation: Element identity is carried by the discretised atomic charge alone rather than a one-hot type variable alongside it, leaving two similar modalities instead of three heterogeneous ones.
• Mode-seeking discretised sampling: Clamping the cumulative probability at the bucket boundaries piles density into the first and last bins; snapping the output distribution to the nearest bin center removes that mode redundancy.

#Technical Details

The network module is an EGNN operating on distribution parameters, trained on the two standard 3D molecule generation corpora, QM9 and the larger, drug-like GEOM-DRUG. On QM9 with 2,000 sampling steps GeoBFN reaches 99.31% atom stability, 93.32% molecule stability and 96.88% validity, against 99.0% and 95.2% for the training data itself and 82.0% molecule stability for EDM, 89.4% for GeoLDM; 4,000 steps raise it to 94.25%. On GEOM-DRUG it reports 86.17% atom stability where the data upper bound is 86.5%. The any-step property is where the efficiency argument lands: at 50 steps GeoBFN already gives 85.11% molecule stability on QM9, above EDM (82.0%) and EDM-Bridge (84.6%) at 1,000 steps though below GeoLDM's 89.4%, a roughly 20-fold reduction in sampling cost. A conditional variant trained on six QM9 properties lowers mean absolute error to 2.34 Bohr³ for polarizability, 328 meV for the HOMO energy and 0.998 D for dipole moment, from 2.76, 356 and 1.111 for EDM. An ablation confirms the redundancy argument: the discretised-charge-only representation reaches 90.87% molecule stability against 88.93% for the discrete one-hot form.

#Applications

The direct use is unconditional generation of chemically valid 3D molecular geometries — filling a virtual library with novel conformers, or supplying starting geometries for downstream docking — with the conditional variant instead sampling molecules that hit a requested electronic or thermodynamic property, and the any-step sampler making library-scale generation a tunable cost. Running it means training from the repository's Docker environment on a GPU, which places it with computational chemistry and method-development groups rather than at the bench.

#Impact

GeoBFN brought Bayesian flow networks into geometric generative modeling and made the case that the difficulty of 3D molecule generation is a property of the space generation happens in, not of the network architecture — a claim the group's pocket-conditioned successors then carried into drug design, where it is now a standard comparison. Its limits are concrete: no trained checkpoint has been released, so every user starts from scratch; the public code implements QM9 only and raises a not-implemented error for any other dataset, leaving the GEOM-DRUG results outside what the release reproduces; the repository carries no license file; and all reported results are in-silico stability, validity and property-prediction metrics rather than synthesised and measured molecules.