cs.LGOct 11, 2025

Transport-Coupled Bayesian Flows for Molecular Graph Generation

Authors: Yida Xiong, Jiameng Chen, Kun Li, Hongzhi Zhang, Xiantao Cai, Lei Lei, Wenbin Hu

Abstract

Molecular graph generation (MGG) is essentially a multi-class generative task, aimed at predicting categories of atoms and bonds under strict chemical and structural constraints. However, many prevailing diffusion paradigms learn to regress numerical embeddings and rely on a hard discretization rule during sampling to recover discrete labels. This introduces a fundamental discrepancy between training and sampling. While models are trained for point-wise numerical fidelity, the sampling process fundamentally relies on crossing categorical decision boundaries. This discrepancy forces the model to expend efforts on intra-class variations that become irrelevant after discretization, ultimately compromising diversity, structural statistics, and generalization performance. Therefore, we propose TopBF, a unified framework that (i) performs MGG directly in continuous parameter distributions, (ii) learns graph-topological understanding through a Quasi-Wasserstein optimal-transport coupling under geodesic costs, and (iii) supports controllable, property-conditioned generation during sampling without retraining the base model. TopBF innovatively employs cumulative distribution function (CDF) to compute category probabilities induced by the Gaussian channel, thereby unifying the training objective with the sampling discretization operation. Experiments on QM9 and ZINC250k demonstrate superior structural fidelity and efficient generation with improved performance.

Explore similar work

CardsList
  1. Variational Bayesian Flow Network for Graph Generation

    Jan 30, 2026Yida Xiong, Jiameng Chen, Xiuwen Gong +3Generative Flow Networks