2025-11-11T14:46:09.738382

Hierarchical Bayesian Flow Networks for Molecular Graph Generation

Xiong, Chen, Li et al.

Molecular graph generation is essentially a classification generation problem, aimed at predicting categories of atoms and bonds. Currently, prevailing paradigms such as continuous diffusion models are trained to predict continuous numerical values, treating the training process as a regression task. However, the final generation necessitates a rounding step to convert these predictions back into discrete classification categories, which is intrinsically a classification operation. Given that the rounding operation is not incorporated during training, there exists a significant discrepancy between the model's training objective and its inference procedure. As a consequence, an excessive emphasis on point-wise precision can lead to overfitting and inefficient learning. This occurs because considerable efforts are devoted to capturing intra-bin variations that are ultimately irrelevant to the discrete nature of the task at hand. Such a flaw results in diminished molecular diversity and constrains the model's generalization capabilities. To address this fundamental limitation, we propose GraphBFN, a novel hierarchical coarse-to-fine framework based on Bayesian Flow Networks that operates on the parameters of distributions. By innovatively introducing Cumulative Distribution Function, GraphBFN is capable of calculating the probability of selecting the correct category, thereby unifying the training objective with the sampling rounding operation. We demonstrate that our method achieves superior performance and faster generation, setting new state-of-the-art results on the QM9 and ZINC250k molecular graph generation benchmarks.

academic

Hierarchical Bayesian Flow Networks for Molecular Graph Generation

基本信息

论文ID: 2510.10211
标题: Hierarchical Bayesian Flow Networks for Molecular Graph Generation
作者: Yida Xiong, Jiameng Chen, Kun Li, Hongzhi Zhang, Xiantao Cai, Wenbin Hu (武汉大学计算机学院)
分类: cs.LG (机器学习)
发表时间: 2025年10月11日 (arXiv预印本)
论文链接: https://arxiv.org/abs/2510.10211

摘要

分子图生成本质上是一个分类生成问题，旨在预测原子和化学键的类别。目前主流的连续扩散模型将训练过程视为回归任务，预测连续数值，但在最终生成时需要通过舍入操作转换为离散分类类别。由于训练过程中未包含舍入操作，模型的训练目标与推理过程存在显著差异，导致过拟合、学习效率低下、分子多样性降低等问题。为解决这一根本性限制，作者提出GraphBFN，一个基于贝叶斯流网络的分层粗到细框架，通过创新性地引入累积分布函数来计算选择正确类别的概率，从而统一训练目标与采样舍入操作。

研究背景与动机

核心问题

分子图生成中存在一个根本性的训练-推理不一致问题：

训练阶段：连续扩散模型将离散的原子/键类别映射到连续空间，通过回归损失优化连续数值预测
推理阶段：需要通过硬舍入将连续预测值转换回离散类别
不一致性：训练时未考虑舍入规则，导致模型过度关注类内变化而忽略离散本质

问题重要性

分子图生成是药物发现的关键技术，影响分子优化、药物-靶点结合亲和力预测等下游任务
现有方法的不一致性导致分子多样性降低、泛化能力受限
即使微小的回归偏差也可能导致完全错误的分类结果

现有方法局限性

离散扩散模型：虽然适合离散图结构，但牺牲了连续表示的平滑性和动态生成特性
连续扩散模型：训练目标与推理过程脱节，容易过拟合到无关的类内变化
传统贝叶斯流网络：假设所有类别在概率单纯形中等距，导致收敛缓慢且噪声较大

核心贡献

首次将贝叶斯流网络应用于分子图生成，通过分层分子表示监督增强生成效果
创新性引入累积分布函数(CDF)，计算各类别概率而非拟合特定数值，统一训练目标与采样舍入操作
提出分层粗到细框架，通过多尺度图表示同时捕获局部原子连接性和全局分子拓扑
实现更快的训练和采样，在QM9和ZINC250k基准上达到新的最先进性能，采样步数显著减少

方法详解

任务定义

给定分子图 $G = (X, A)$ ，其中：

$X \in \{0, \ldots, K_X - 1\}^D$ ： $D$ 个原子特征矩阵，来自 $K_X$ 个类别
$A \in \{0, \ldots, K_A - 1\}^{D \times D}$ ：邻接矩阵，包含 $K_A$ 个键类别的特征

目标是学习生成符合真实分子分布的新分子图。

模型架构

1. 分层粗到细框架

多尺度表示：使用DiffPool构建 $L$ 层粗化层，生成分子图的金字塔表示
自底向上生成：从最粗层开始无条件生成，逐层细化到完整原子图
条件传递：每层的上采样模块 $\phi_1^{(l)}$ 将粗层输出转换为细层条件 $c^{(l)}$

2. 图表示映射

将离散类别 $k \in \{0, \ldots, K-1\}$ 映射到连续空间 $[-1, 1]$ ：

k_c = (2k + 1)/K - 1  # 中心点
k_l = k_c - 1/K       # 左边界  
k_r = k_c + 1/K       # 右边界

3. 贝叶斯流网络组件

输入分布：使用高斯分布建模

p_I(G|θ) = N(G|μ, ρ^{-1}I)

发送分布：添加高斯噪声

p_S(Y|G; α) = N(Y|G, α^{-1}I)

输出分布：通过CDF计算离散概率

p_O^{(d)}(k|θ; t) = F(k_r|μ_x^{(d)}, σ_x^{(d)}) - F(k_l|μ_x^{(d)}, σ_x^{(d)})

接收分布：

p_R(Y|θ; t, α) = ∏_{d=1}^D ∑_{k=0}^{K-1} p_O^{(d)}(k|θ; t)N(Y^{(d)}|k_c, α^{-1})

4. 关键创新：CDF机制

使用截断的累积分布函数连接连续分布与离散类别：

F(x|μ_x^{(d)}, σ_x^{(d)}) = {
  0,                    if x ≤ -1
  1,                    if x ≥ 1  
  1/2[1 + erf((x-μ_x^{(d)})/(√2σ_x^{(d)}))], otherwise
}

技术创新点

训练-推理一致性：CDF直接计算离散概率，避免连续预测与离散舍入的不匹配
非等距类别映射：不同于传统BFN假设类别等距，允许更快更平滑的收敛
多尺度监督：分层框架在不同粒度提供结构信息，增强生成质量
端到端优化：统一的损失函数同时优化BFN生成损失和池化损失

实验设置

数据集

QM9：包含134k个小分子的量子化学数据集
ZINC250k：包含250k个相对较大分子的药物样分子数据集

评价指标

Validity w/o correction：未经校正的有效分子比例
Uniqueness：生成分子的唯一性比例
FCD (Fréchet ChemNet Distance)：使用ChemNet特征计算的训练集与生成集距离
NSPDK MMD：考虑原子和键特征的邻域子图成对距离核最大均值差异

对比方法

包括多种最先进基线：

流模型：MoFlow
扩散模型：EDP-GNN, GDSS, DiGress, GSDM
流匹配：Dirichlet FM, CatFlow
能量模型：GraphEBM

实现细节

采样步数：GraphBFN使用100×L步（L为层数），显著少于基线的400-1000步
多尺度损失平衡参数：λ₁, λ₂
最小时间阈值：t_min = 10⁻⁵

实验结果

主要结果

方法	QM9 Val.↑	QM9 Unique↑	QM9 FCD↓	QM9 NSPDK↓	ZINC250k Val.↑	ZINC250k Unique↑	ZINC250k FCD↓	ZINC250k NSPDK↓	采样步数
GDSS	95.72	98.46	2.565	0.0033	97.12	99.64	14.032	0.0192	1000
CatFlow	99.81	99.95	0.441	0.0029	99.21	100.00	13.211	0.0207	-
GraphBFN	99.60	99.97	0.214	0.0008	96.00	100.00	5.743	0.0069	100×L