MolHIT：利用层级离散扩散模型推进分子图生成

MolHIT通过HDDM和DAE，首次在图扩散中实现99.1%有效率。

前置知识

离散扩散模型 (Discrete Diffusion Models)

离散扩散模型是一类直接在离散状态空间（如词表、原子类型集合）上定义噪声过程的生成模型。其前向过程通过转移矩阵逐步将干净数据替换为其他类别或掩码状态，反向过程则学习从噪声逐步恢复原始数据。与连续扩散模型（如 DDPM）不同，离散扩散模型天然适合处理分子中的原子类型、键类型等分类变量。典型方法包括 D3PM（Austin et al., 2021）定义的均匀转移和吸收转移，以及 DiGress（Vignac et al., 2022）将离散扩散应用于图结构生成。转移核的设计直接决定了模型的学习难度和生成质量。

本文的核心贡献是在离散扩散框架中引入层级结构（HDDM），理解标准离散扩散的前向加噪和反向去噪机制是理解 HDDM 创新点的必要前提。

分子图生成 (Molecular Graph Generation)

分子图生成是将分子表示为图结构（节点为原子、边为化学键）并用生成模型采样新分子的任务。核心挑战在于生成的图不仅要满足数学上的合法性，更要通过化学规则验证（如价态合理性、芳香性一致性）。与基于 SMILES 字符串的 1D 序列生成相比，图结构能更自然地表达分子拓扑，但搜索空间更大、约束更复杂，因此图扩散模型长期面临化学有效率低于 1D 模型的困境。

本文的目标正是解决图生成模型在化学有效率上长期落后于 1D 模型的问题，理解这一背景有助于把握论文的定位和贡献。

解耦原子编码 (Decoupled Atom Encoding, DAE)

传统分子图生成模型通常仅根据原子序数（如 C、N、O）对节点进行编码，忽略了同一元素在不同化学环境下的差异。例如，脂肪族氮（N）和芳香族氮（n）、吡咯氮（nH）在化学性质上截然不同，但传统编码将它们视为同一个 Token。DAE 通过扩展原子词表，根据芳香性、形式电荷和氢饱和度对原子进行显式拆分。在 MOSES 数据集中，它将 7 种基础原子类型扩展为 12 种语义明确的 Token，从而解决了传统粗粒度编码导致的一对多映射歧义问题。

DAE 是本文方法的关键组成部分，直接解决了先前图扩散模型无法正确生成含吡咯氮骨架等复杂药物结构的根本原因。

Project-and-Noise (PN) 采样器

标准离散扩散模型的反向采样通常基于后验分布进行逐步更新。PN 采样器则采用不同的策略：在每一步去噪时，模型先预测干净数据的分布，通过 Top-p 和温度采样将其投影为一个确定性的离散候选样本，然后直接将该候选样本加噪回前一时间步。这种投影-再加噪机制绕过了标准后验约束，使得采样过程能够在化学合理的流形上进行更广泛的探索，从而提升生成分子的结构多样性。

PN 采样器是 MolHIT 实现高质量与高新颖性平衡的关键技术，配合温度和 Top-p 参数可以精细控制生成的多样性。

研究动机

现有分子生成模型面临一个长期存在的矛盾：基于 1D 序列（如 SMILES 表示）的模型虽然能达到较高的化学有效率（通常 97% 以上），但极易陷入对训练集文本模式的记忆，导致生成的分子结构缺乏新颖性，无法有效探索广阔的化学空间。例如，SAFE-GPT 在 MOSES 上的骨架新颖性仅为 0.12。另一方面，基于 2D 图结构的离散扩散模型虽然在结构探索上具有天然优势（骨架新颖性可达 0.36），但由于忽略了原子间明确的化学先验关系以及采用了粗粒度的原子编码，生成的分子化学有效率普遍较低（DiGress 仅 87.1%，Cometh 为 87.2%）。更严重的是，传统原子编码将所有原子类别视为独立且等价的，忽略了化学领域中已知的层级关系（如卤素原子之间更容易相互替换），使得学习过程变得不必要地困难。在 MOSES 数据集的重建实验中，传统编码对吡咯氮 [nH] 组的重建成功率仅为约 20%，几乎无法生成这类在药物骨架中极为常见的结构。

本文的目标是本文旨在打破图扩散模型在化学有效率和结构新颖性之间的性能瓶颈，提出一种既能像 1D 模型那样保证极高的化学有效率（接近 100%），又能保持 2D 图模型优越结构探索能力的分子生成框架。具体目标包括：在 MOSES 基准数据集上实现接近完美的有效率（目标 99% 以上），在质量（Quality）指标上超越所有现有 1D 和 2D 基线，在骨架新颖性（Scaffold Novelty）上达到帕累托最优前沿，并在条件生成和骨架扩展等下游任务中验证框架的实用价值。

与已有工作不同的是，本文的独特切入角度在于将层级化的思想引入离散扩散过程，并从原子表示的根本层面进行改进。作者发现了一个被已有工作忽视的关键问题：现有图扩散模型的原子编码存在不适定性（ill-posed），即同一个原子 Token 可以对应多种不同的物理状态（如中性与带电、脂肪族与芳香族），这使得模型在生成时无法精确区分这些化学角色，导致系统性的重建失败。此外，现有扩散框架将所有原子类别视为扁平的集合，忽略了化学领域中固有的层级结构（如元素周期表中的族划分、芳香族与脂肪族的分组）。MolHIT 通过引入层级离散扩散模型（HDDM）在扩散过程中显式建模这些化学先验，并通过解耦原子编码（DAE）从根源上解决表示信息丢失的问题，这是之前任何工作都没有尝试过的角度。

核心方法

MolHIT 的整体方法可以类比为一位化学家从模糊印象到精确结构的推理过程。想象你在雾中看到一个分子的轮廓，首先能判断出它大致含有哪些类型的化学基团（如卤素环、芳香环），然后逐步细化到具体的原子（如氯原子还是溴原子）。这就是 HDDM 的核心思想：在扩散过程中引入中间状态层，让模型先学习高层化学类别，再细化到具体原子类型。技术路线上，MolHIT 包含三个关键组件：（1）层级离散扩散模型（HDDM），通过在干净状态和掩码状态之间增加语义中间层，构建由粗到细的前向和反向过程；（2）解耦原子编码（DAE），根据芳香性和电荷等化学角色显式拆分原子类型，扩展原子词表以消除表示歧义；（3）Project-and-Noise 采样器，通过投影-再加噪机制增加生成多样性。这些组件共同作用，使模型能够生成既化学有效又结构新颖的分子。

MolHIT 的核心创新在于将层级结构引入离散扩散模型的转移核设计。与以往将所有原子类别视为扁平集合的均匀转移或吸收转移不同，HDDM 引入了中间状态空间 $S_1$，其中包含 $G$ 个语义组（例如在 MOSES 中将 12 种原子分为 4 组：碳族、氮氧硫族、卤素族、芳香族）。前向扩散过程通过两个转移矩阵 $Q^{(1)}$ 和 $Q^{(2)}$ 实现：$Q^{(1)}$ 将干净状态 $S_0$ 映射到中间状态 $S_1$（吸收过程），$Q^{(2)}$ 将 $S_0$ 和 $S_1$ 映射到掩码状态 $S_2$。这一设计的数学关键在于 Lemma 3.1 证明了存在满足 Chapman-Kolmogorov 方程的闭式转移核，从而保证了训练目标（NELBO）的理论正确性。与已有方法的本质区别是：DiGress 使用数据依赖的边际分布作为先验，DeFoG 使用连续时间马尔可夫链框架，而 MolHIT 通过显式的层级结构让模型在去噪时先确定原子的化学类别再确定具体类型，大幅降低了学习难度。

方法步骤详情

MolHIT 的完整方法流程如下。第一步是数据预处理与解耦原子编码：使用 DAE 将分子图中的原子节点从基础元素映射到包含芳香性和氢饱和度的扩展词表（如 MOSES 中从 7 个 Token 扩展到 12 个）。第二步是构建前向扩散过程：原子节点使用 HDDM 转移核，将干净状态 $S_0$ 经由中间状态 $S_1$ 转移到掩码状态 $S_2$；键（边）则使用标准的均匀转移过程。转移核的数学形式为 $Q_t = \alpha_t I + (\beta_t - \alpha_t) Q^{(1)} + (1 - \beta_t) Q^{(2)}$，其中 $\alpha_t$ 和 $\beta_t$ 是单调递减的扩散调度函数。第三步是模型训练：采用 DiGress 的图 Transformer 架构（12 层，约 16.2M 参数），训练网络预测去噪后的干净数据 $x_0$，损失函数为基于 HDDM 推导的负对数似然下界（NELBO）。第四步是反向生成（PN 采样）：在每一步去噪时，模型先预测 $x_0$ 的分布，通过 Top-p 和温度采样将其投影为离散候选样本 $\hat{x}_0$，然后直接加噪至前一时间步。第五步是条件生成：通过 AdaLN 注入条件信息（如 QED、SA 等属性），并使用 Classifier-Free Guidance 进行采样。

技术新颖性

MolHIT 的技术新颖性体现在多个层面。首先，它是首个将层级离散扩散（HDDM）应用于分子图生成的工作。与 Zhou et al. (2025) 在语言建模中提出的层级扩散方法相比，HDDM 从半群一致的闭式转移核出发，支持任意行随机投影矩阵 $\Phi$，具有更一般的数学形式和更强的灵活性。其次，DAE 机制首次在图扩散中系统性地解决了原子编码的不适定问题，使模型能够生成带形式电荷和显式氢的原子（如 [nH]），这在 DiGress、DisCo、Cometh、DeFoG 等所有先前图扩散模型中都无法实现。第三，PN 采样器是一种创新的离散扩散采样策略，通过投影-再加噪机制打破了标准后验更新的约束，有效地平衡了生成质量与新颖性。消融实验表明，DAE 贡献了从 82.5% 到 87.6% 的有效率提升，PN 采样器进一步将质量从 87.6% 提升到 92.9%，HDDM 最终将质量推至 94.2% 并优化了 FCD 指标。

实验结果

MolHIT 在多个基准测试中取得了突破性进展，全面超越现有 1D 和 2D 基线模型。在 MOSES 数据集的无条件生成任务中，MolHIT 实现了 99.1% 的高有效率，这是图扩散模型首次接近完美有效率，远超 DiGress（87.1%）、Cometh（87.2%）和 DeFoG（92.8%）。其质量得分（Quality）达到 94.2%，不仅创下图扩散模型的新纪录，甚至超越了所有 1D 序列基线（如 SAFE-GPT 92.8%、CharRNN 92.6%）。骨架新颖性为 0.39，较此前最佳的 Cometh（0.36）提升 8.3%。FCD 指标为 1.03，优于所有基线。在 GuacaMol 数据集上，MolHIT 在全量未筛选数据上训练，有效率达到 87.1%，远高于在筛选数据上训练的 DiGress（85.2%）和在全量数据上训练的 DiGress（74.7%），证明了 DAE 对复杂化学空间（含带电原子）的鲁棒性。在多属性条件生成任务中，MolHIT 的平均绝对误差仅为 0.058，较 Marginal+DAE 基线降低 52.4%，同时保持了 96.31% 的高有效率。在骨架扩展任务中，MolHIT 的有效率达到 83.9%（DiGress 仅 50.8%），Hit@5 为 9.79，较 DiGress（6.41）提升 52.7%。消融实验进一步验证了各组件的贡献：DAE 单独使用将有效率从 87.1% 提升至 96.2%，加入 PN 采样器后质量从 87.6% 提升至 92.9%，HDDM 进一步将质量推至 94.2% 并将 FCD 从 1.65 优化至 1.03。

**Table 1**: Comprehensive MOSES benchmark results.

**Table 3**: Multi-property guided generation on MOSES with four different conditions.

**Table 4**: Scaffold extension results on the MOSES dataset.

**Table 5**: Incremental performance gains on the MOSES dataset by integrating DAE, the PN Sampler, and HDDM into the DiGress.

**Figure 1**: MolHIT achieves SOTA result on MOSES dataset.

**Figure 4**: Effect of top-p sampling in MolHIT.

查看结构化数据

任务	指标	本文	基线	提升
MOSES 无条件生成	Quality (%)	94.2	88.5 (DeFoG), 82.5 (DiGress), 82.1 (Cometh), 92.8 (SAFE-GPT)	提升 5.7%（相对 DeFoG），首次超越 1D 模型
MOSES 无条件生成	Validity (%)	99.1	92.8 (DeFoG), 88.3 (DisCo), 87.2 (Cometh), 87.1 (DiGress)	提升 6.3%（相对 DeFoG），图扩散模型首次接近 100%
MOSES 无条件生成	Scaffold Novelty	0.39	0.36 (Cometh), 0.29 (CharRNN), 0.26 (DiGress/DeFoG)	提升 8.3%（相对 Cometh）
MOSES 无条件生成	FCD (lower is better)	1.03	1.25 (DiGress), 1.44 (Cometh/DisCo), 1.95 (DeFoG)	降低 17.6%（相对 DiGress）
GuacaMol 无条件生成	Validity (%)	87.1	85.2 (DiGress on filtered), 74.7 (DiGress on full)	提升 2.2%（相对 DiGress on filtered），在更难的全量数据上
MOSES 多属性条件生成	MAE Avg.	0.058	0.122 (Marginal+DAE), 0.143 (Marginal)	误差降低 52.4%（相对 Marginal+DAE）
MOSES 骨架扩展	Hit@5	9.79	6.41 (DiGress), 6.37 (Marginal+DAE)	提升 52.7%（相对 DiGress）

局限与改进

尽管 MolHIT 取得了显著成果，但仍存在一些值得关注的局限性。首先，作者明确指出尚未测试模型规模增加或架构改进带来的潜力，当前模型仅约 16.2M 参数，在更大化学空间中的表达能力可能受限。其次，在 GuacaMol 数据集上模型仅训练了 40 个 epoch（相比 DiGress 的 1000 epoch），FCD 指标（54.9）略低于原始 DiGress（68.0），表明在分布匹配的精细度上仍有提升空间，作者也承认进一步训练会改善指标。第三，虽然 DAE 扩展了原子词表解决了表示问题，但在 GuacaMol 上 DiGress+DAE 的 FCD（49.2）反而低于原始 DiGress（61.1），说明扩展词表本身增加了建模难度。此外，PN 采样器在极低 Top-p 值下会导致化学指标和多样性急剧下降，模型在极端多样性采样下的稳定性有待加强。条件生成中 CFG 引导权重超过 1.0 时效果不一致，限制了强条件约束下的生成能力。论文也未讨论计算开销的详细对比。

独立分析的弱点

从独立分析的角度来看，MolHIT 存在以下几个值得关注的弱点。第一，DAE 虽然解决了原子编码的不适定问题，但代价是词表规模大幅膨胀（MOSES 从 7 到 12，GuacaMol 从 12 到 56），这增加了模型的学习负担。GuacaMol 实验中 DiGress+DAE 的 FCD 下降表明，简单的词表扩展并非总是最优方案，未来可以探索基于子图或官能团的分层编码方式，在保留化学精度的同时控制词表规模。第二，PN 采样器引入了温度和 Top-p 两个超参数，需要根据具体任务进行调优，缺乏自适应机制。可以考虑设计基于去噪阶段的动态调度策略，在早期去噪步使用高温度促进探索，在后期步使用低温度保证质量。第三，HDDM 的分组策略是手动设计的（如 MOSES 中的 4 组），依赖领域知识，未来可以探索数据驱动的自动分组方法。第四，论文未与其他最新的分子生成方法（如基于 3D 几何的扩散模型）进行对比，难以评估在更现实场景中的竞争力。

未来方向

作者提出了几个有前景的未来研究方向。首先是将 HDDM 框架推广到语言模型和图像生成领域，验证其层级扩散思想的通用性。其次是结合更先进的分子标记化方法，如基于官能团或子结构的 Motif 分解（如 Jin et al., 2018 的 Junction Tree VAE），进一步提升生成的语义合理性。第三是将框架扩展到 3D 分子生成和蛋白质设计，这对于药物发现中的分子对接和蛋白质折叠任务至关重要。基于当前成果还可以延伸出更多方向：将 MolHIT 应用于逆向合成规划，利用其高有效率特性加速实际药物筛选流程；探索多目标帕累托优化，在多个属性约束下同时优化分子结构；以及将 HDDM 与大规模预训练结合，在更广泛的化学空间上训练基础模型。此外，结合强化学习进行基于反馈的分子优化也是一个值得探索的方向。

复现评估

该工作具有较高的可复现性。作者在论文中提供了详细的实验设置，包括模型架构（基于 DiGress 的图 Transformer，12 层，约 16.2M 参数）、训练超参数（梯度裁剪阈值 2.0、EMA 衰减率 0.999）、PN 采样器的完整伪代码（Algorithm 1）以及 HDDM 的数学推导和证明。数据集方面，MOSES 和 GuacaMol 都是公开可用的标准基准，DAE 的分组策略在附录中有详细表格（Table 8）。实验结果均报告了 3 次独立运行的均值和标准差。虽然论文未直接附带代码，但基于 DiGress 的开源代码库进行复现是可行的，DAE 和 HDDM 的实现复杂度适中。所需算力方面，模型规模较小（16.2M 参数），在单 GPU 上即可训练，适合学术实验室进行验证和改进。唯一需要注意的是 RDKit 版本差异可能影响分子属性计算的数值。

MolHIT：利用层级离散扩散模型推进分子图生成 MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models