MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models¶

会议: NeurIPS 2026
arXiv: 2602.17602
OpenReview: https://openreview.net/forum?id=QkdUSTr5Md
论文: https://arxiv.org/abs/2602.17602
领域: 计算生物
关键词: 分子图生成, 离散扩散模型, 分层扩散, 解耦原子编码, 药物发现

一句话总结¶

提出 MolHIT 分子图分层离散扩散框架，通过分层状态空间融入化学先验，配合解耦原子编码消除表示歧义，在图扩散模型中首次实现接近 100% 的化学有效性并超越 1D 序列基线。

研究背景与动机¶

基于生成式 AI 的从头分子设计是加速药物研发与材料创新的关键路径。分子在化学本质上是由原子与化学键构成的离散二维图结构，但长久以来分子生成领域一直面临“1D 序列模型”与“2D 图生成模型”之间的根本权衡。以 SMILES 为代表的 1D 序列模型能够直接套用成熟的自回归或掩码语言模型架构，具备极高的化学有效性（Validity）；然而，序列表示极易记住训练集的局部子串与文本模态模式，导致模型陷入结构探索瓶颈，难以跳出已有骨架生成真正多样且新颖的分子。

为了突破探索边界，基于图结构的模型（尤其是与离散属性天然契合的离散图扩散模型）成为研究热点。然而现存图扩散模型在生成化学有效且稳定的分子方面表现欠佳，化学有效性普遍徘徊在 80%~90% 甚至更低。这种性能缺陷根源于两个长期被忽视的设计瓶颈：第一，传统离散图扩散采用均匀（Uniform）或吸收（Absorbing/Masked）转移核，将所有原子类别视为相互独立的孤立状态，完全抛弃了元素周期律与官能团相似性等化学先验，迫使模型在稀缺的高质量分子数据上从零硬学高维组合空间；第二，现有图扩散模型普遍沿用极为粗糙的原子编码方案（仅按原子序数编号），忽略了同一原子在不同化学环境下的形式电荷与芳香性差异。例如对含有吡咯氮（[nH]）的杂环，粗粒度编码在重构阶段出现严重的信息损失与病态多对一映射，导致生成结果存在不可逆的化学缺陷。

面对图扩散在“探索自由度”与“化学有效性”之间的割裂，本文旨在探索能否在保留图建模结构探索优势的同时，达到乃至超越序列模型的化学有效性。核心 idea：构建分层离散扩散模型（HDDM）以粗到细的方式注入化学语义分组先验，并提出解耦原子编码（DAE）显式建模芳香性与形式电荷，配合投影加噪采样器实现高有效性与高质量外推。

方法详解¶

整体框架¶

MolHIT 框架将分子的二维拓扑拆解为原子节点矩阵 $\mathbf{X}$ 与键邻接张量 $\mathbf{E}$。在扩散前向过程中，节点与边采用解耦动力学演化：对原子类别引入中间语义状态集合 $\mathcal{S}_1$ 构建分层马尔可夫转移链，边类别则保持均匀转移；在逆向去噪阶段，图 Transformer 仅需预测原始干净数据分布，并通过投影加噪采样策略（PN-sampler）在离散流形上引导分子从掩码高斯噪声逐步粗粒化聚集再细化为精确化学结构。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入分子图<br/>节点特征 X + 边张量 E"] --> B["解耦原子编码<br/>显式解耦元素/芳香性/形式电荷"]
    B --> C["分层离散扩散前向过程<br/>S0原子状态 → S1化学语义组 → S2掩码吸收态"]
    C --> D["图Transformer去噪网络<br/>节点自适应LayerNorm注入时间与属性"]
    D --> E["投影加噪采样器<br/>类别投影到干净流形 + 重加噪探索"]
    E --> F["有效药物分子图输出"]

关键设计¶

1. 解耦原子编码：消除节点化学角色的病态多对一映射
以往图扩散模型（如 DiGress、DeFoG）仅用原子序数表示节点，迫使模型仅依赖 $O(n^2)$ 的键邻接矩阵去隐式推断原子的电荷状态与环共轭性。在包含吡咯环或带电官能团的药物骨架中，仅靠原子序数无法判定氢原子的连接位置与饱和度，导致分子重构成功率大幅下跌。解耦原子编码（Decoupled Atom Encoding, DAE）将节点的表示空间显式解耦为元素类型、芳香态（Aromaticity）与形式电荷（Formal Charge）。例如在 MOSES 数据集上将 7 种重原子细分为 12 种化学标记（区分脂肪族 C/N/O/S 与芳香族 c/n/nH/s/o）；在 GuacaMol 数据集上扩展至 56 种解耦 token。这种设计彻底避免了将显式氢作为独立图节点引起的拓扑复杂度和计算量暴增，在保持图稀疏性的同时赋予每个原子清晰的化学约束。

2. 分层离散扩散模型：构建语义群到精确原子的由粗到细扩散链
现有离散扩散直接在 $K$ 个原子类别与掩码吸收态间跳跃，忽略了元素之间的近亲替换规律。HDDM 将离散状态空间扩充为三级子空间 $\mathcal{T} = \mathcal{S}_0 \cup \mathcal{S}_1 \cup \mathcal{S}_2$：$\mathcal{S}_0$ 为 $K$ 类干净原子状态，$\mathcal{S}_1$ 为 $G$ 类中间化学语义群（如卤素群、杂原子群、芳香环群），$\mathcal{S}_2 = \{\mathbf{m}\}$ 为唯一的全掩码吸收态。通过行随机矩阵 $\mathbf{\Phi} \in [0, 1]^{K \times G}$ 定义投影转移核 $Q^{(1)}$ 与全掩码转移核 $Q^{(2)}$，构建多阶段前向扰动矩阵： $$Q_{X,t} = \alpha_t \mathbf{I} + (\beta_t - \alpha_t) Q^{(1)} + (1 - \beta_t) Q^{(2)}$$ 其中单调递减调度函数满足 $0 \le \alpha_t \le \beta_t \le 1$。该转移核在数学上严格满足 Chapman-Kolmogorov 一致性。逆向去噪时，模型网络仅需估计干净原子概率 $p_\theta(\mathbf{X}_0 \mid G_t, t)$，HDDM 对应的负变分下界（NELBO）被证明可自然退化并以标准交叉熵损失训练，驱动生成过程先决定高层化学骨架基团再推断精确原子。

3. 投影加噪采样器：打破离散后验约束以提升骨架多样性与有效性
基于理论后验分布 $q(G_{t-\Delta t} \mid G_t, \hat{G}_0)$ 逐步回推的传统采样器极易受到当前带噪图 $G_t$ 的局部状态拖累，陷入次优局部解。本文设计了投影加噪采样器（Project-and-Noise Sampler, PN-sampler）：在每一个反向步 $t$，模型预测出干净图概率后，通过 Categorical 采样将其直接“投影”到离散真实流形空间获得单热离散候选 $\hat{G}_0 = (\hat{\mathbf{X}}_0, \hat{\mathbf{E}}_0)$，随后直接利用累积前向转移矩阵 $q(G_{t-\Delta t} \mid \hat{G}_0)$ 将其重加噪至前序时刻 $t-\Delta t$。此外，PN-sampler 针对原子预测采用温度系数与 Top-$p$（Nucleus）截断采样，只控制原子类型的候选集质量，在极大拓宽分子骨架探索自由度的同时保持了极低的无效结构率。

损失函数 / 训练策略¶

模型网络主干采用 12 层 Graph Transformer，参数量约为 16.2M。在训练阶段，原子与边分别施加掩码交叉熵损失，总损失函数为： $$\mathcal{L}_{\theta} = \mathbb{E}_{t, G_t \sim q(\cdot \mid G_0)} \left[ \sum_{i=1}^n -\log p_\theta^X(X_{0,i} \mid G_t, t) + \lambda \sum_{1 \le i < j \le n} -\log p_\theta^E(E_{0,ij} \mid G_t, t) \right]$$ 其中 $\lambda > 0$ 平衡节点与边损失。调度策略上，原子与边分别采用 $\alpha_{X,t} = \alpha_{E,t} = 1 - t$ 及 $\beta_{X,t} = 1 - t^2$。在 MOSES 上仅训练 100 轮，GuacaMol 上训练 50 轮，配合梯度裁剪（阈值 2.0）与指数移动平均（EMA 0.999）。条件生成任务中，利用自适应层归一化（AdaLN）注入由属性 MLP 与正弦时间编码构成的控制向量 $\mathbf{C}$，并辅以无分类器引导（Classifier-Free Guidance, CFG）。

实验关键数据¶

主实验¶

在 MOSES 基准数据集上的无条件分子生成测试中，MolHIT 全面超越既有 2D 图扩散模型，并在多项核心指标上击败强有力的 1D 序列基线。

数据集 / 模型类别	模型	Quality (↑)	Validity (%) (↑)	Unique (%) (↑)	Novelty (%) (↑)	Scaf-Novel (↑)	FCD (↓)
MOSES (训练集基准)	Ground Truth	95.4	100.0	100.0	—	—	0.48
1D Sequence	VAE	92.8	97.7	99.7	69.5	0.22	0.57
1D Sequence	CharRNN	92.6	97.5	99.9	84.2	0.29	0.52
1D Sequence	SAFE-GPT	92.8	99.8	98.9	43.7	0.12	0.72
1D Sequence	GenMol	62.1	99.7	64.0	68.9	0.05	16.36
2D Graph	DiGress	82.5	87.1	100.0	94.2	0.26	1.25
2D Graph	Cometh	82.1	87.2	100.0	96.4	0.36	1.44
2D Graph	DeFoG	88.5	92.8	99.9	92.1	0.26	1.95
2D Graph (本文)	MolHIT	94.2	99.1	99.8	91.4	0.39	1.03

在未经过滤的完整 GuacaMol 基准上，MolHIT 达到 87.1% 的 Validity 和 96.7% 的 KL 散度，相较于全量数据训练下的 DiGress（74.7% Validity，92.4% KL）实现显著代差跨越。

消融实验¶

在 MOSES 数据集上逐项剥离各核心组件的递增消融分析：

配置	Quality (↑)	FCD (↓)	Validity (%) (↑)	说明
DiGress 基线	82.5	1.25	87.1	传统边缘与节点离散扩散
+ DAE (解耦原子编码)	87.6	0.89	96.2	消除电荷与芳香歧义，Validity 激增 9.1%
+ PN Sampler (投影加噪)	92.9	1.65	99.4	消除后验牵制，Validity 达到近乎完美的 99.4%
+ HDDM (完整 MolHIT)	94.2	1.03	99.1	注入化学层级先验，综合药效质量 Quality 达到最优

关键发现¶

解耦原子编码 DAE 是化学有效性跃升的第一功臣：仅引入 DAE 就使 DiGress 的有效性从 87.1% 跃升至 96.2%，重构失败率归零，彻底消除了含吡咯氮 [nH] 分子的系统性生成缺失。
PN 采样器是达成“序列级有效性”的关键：将去噪估计投影至流形再重加噪，绕过了连续概率后验在复杂图上的扰动阻碍，促使有效性突破 99%。
HDDM 实现了最佳的分布保真度与创新性权衡：单纯使用 PN 采样器会使 FCD 略微退化（1.65），而引入 HDDM 分层状态后 FCD 回降至 1.03，且 Quality 达到 94.2%，骨架创新率 Scaf-Novel 达 0.39，处于全局帕累托前沿。
在四属性联合引导条件生成中，MolHIT 的宏平均 MAE 仅为 0.058，比 Marginal+DAE 降低了 52.4%，logP 与 QED 的皮尔逊相关系数分别达到 0.950 和 0.804。

亮点与洞察¶

解耦化学状态而非暴增节点：将形式电荷和芳香度作为原子特征解耦建模，既避免了显式加氢造成的图节点规模膨胀，又从根源上治愈了图生成中普遍存在的化学不合法问题。
数学严谨的离散状态分层扩散：通过构造分块行随机矩阵，在无需复杂连续时间马尔可夫链数值求解的前提下，推导出闭式分层转移核与 NELBO，可推广至任意树状分层离散词表。
重构离散扩散采样范式：提出的 PN-sampler 打破了严格后验转移的传统束缚，用“流形投影+前向重注噪”的思想极大增强了组合空间的跳出逃逸能力。

局限与展望¶

模型体量与骨干网络未作针对性扩展：实验仍受限于 DiGress 提出的早期 16M 参数量小型 Graph Transformer，未在类似现代大语言模型的数十亿级超大网络架构上探究 Scaling Law。
全量大集训练步数不足：在化学空间更广阔的 GuacaMol 全集上仅训练了 50 轮，模型收敛尚未饱和，FCD 指标仍有进一步优化空间。
尚未扩展至 3D 几何构象：目前框架聚焦于 2D 分子图拓扑生成，未来可将分层离散状态与等变 3D 坐标扩散结合，实现端到端的立体药物分子构象设计。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [数学闭式解的分层离散扩散与解耦原子编码极具洞察力]
实验充分度: ⭐⭐⭐⭐⭐ [MOSES/GuacaMol双基准全面评测，覆盖无条件、多属性引导与骨架扩展]
写作质量: ⭐⭐⭐⭐⭐ [论证逻辑严密，图表详实，理论推导与经验结论高度自洽]
价值: ⭐⭐⭐⭐⭐ [攻克了图扩散模型多年来弱于序列模型的顽疾，为 AI 辅助药物分子设计树立了新标杆]