[论文解读] Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3D

Bo Qiang, Yuxuan Song+6 more|arXiv (Cornell University)|May 5, 2023

Innovative Microfluidic and Catalytic Techniques Innovation被引用 5

一句话总结

该论文提出 HierDiff，一种用于 3D 分子生成的层次扩散模型，采用自粗到细的策略以提升结构有效性与质量。首先通过等变扩散过程生成片段级表征，随后通过消息传递与迭代采样对表征进行精炼，生成精确的原子 3D 结构，在药物类脂性、构象准确性与稳定性方面优于现有方法，尤其在大分子上表现更优。

ABSTRACT

Generating desirable molecular structures in 3D is a fundamental problem for drug discovery. Despite the considerable progress we have achieved, existing methods usually generate molecules in atom resolution and ignore intrinsic local structures such as rings, which leads to poor quality in generated structures, especially when generating large molecules. Fragment-based molecule generation is a promising strategy, however, it is nontrivial to be adapted for 3D non-autoregressive generations because of the combinational optimization problems. In this paper, we utilize a coarse-to-fine strategy to tackle this problem, in which a Hierarchical Diffusion-based model (i.e.~HierDiff) is proposed to preserve the validity of local segments without relying on autoregressive modeling. Specifically, HierDiff first generates coarse-grained molecule geometries via an equivariant diffusion process, where each coarse-grained node reflects a fragment in a molecule. Then the coarse-grained nodes are decoded into fine-grained fragments by a message-passing process and a newly designed iterative refined sampling module. Lastly, the fine-grained fragments are then assembled to derive a complete atomic molecular structure. Extensive experiments demonstrate that HierDiff consistently improves the quality of molecule generation over existing methods

研究动机与目标

• 解决原子级 3D 分子生成的局限性，特别是大分子中结构有效性差与误差累积问题。
• 克服基于片段的框架中非自回归 3D 生成的挑战，确保片段连接性与几何结构得以保留。
• 通过片段级建模而非原子级扩散编码化学归纳偏置，以提升生成质量。
• 实现稳定且高质量的 3D 构象生成，同时尊重真实分子的几何与拓扑约束。
• 在类药物性质与构象准确性方面，相比最先进扩散模型与自回归模型，实现更优性能。

提出的方法

• 提出两阶段层次扩散框架：先生成粗粒度片段表征，再解码为细粒度原子结构。
• 使用等变扩散模型以非自回归方式生成片段级坐标与特征，保留几何与化学约束。
• 引入两类粗粒度特征：基于元素的直方图与基于性质的描述符，以实现化学可解释的片段表征。
• 应用等变消息传递网络，确保片段解码过程中拓扑一致性与连接性。
• 设计迭代精炼采样模块，以纠正几何偏差，并在片段与原子层级提升构象准确性。
• 通过组合精炼后的片段构建最终的原子 3D 结构，同时保持键价与空间约束。

实验结果

研究问题

• RQ1 与原子级扩散模型相比，自粗到细的层次扩散框架是否能提升 3D 分子生成的结构有效性与质量？
• RQ2 基于片段的建模在多大程度上可缓解大分子生成中的误差累积并提升泛化能力？
• RQ3 化学可解释的粗粒度特征（基于元素 vs. 基于性质）在生成类药物与稳定构象方面的影响程度如何？
• RQ4 在片段级抽象下，非自回归 3D 生成能否实现优于自回归基线的构象准确性？
• RQ5 所提方法的计算效率如何？是否可在减少扩散步数的同时保持质量，实现加速？

主要发现

• 在 QM9 与 GEOM_DRUG 数据集上，HierDiff 均优于最先进模型，在非自回归设置下，对含最多 8 个重原子的分子实现 100% 的结构有效性。
• 在 GEOM_DRUG 数据集上，HierDiff 的片段级 RMSD 为 0.609，原子级 RMSD 为 0.632，显著优于基线模型。
• HierDiff-P（基于性质的特征）在 GEOM_DRUG 上获得更高的类药物性评分，而 HierDiff-E（基于元素的特征）在 MCF 与化学安全性指标上表现更优。
• 该模型显著减少了大分子中的误差累积，在重原子数超过 20 的分子上仍保持高有效性，而自回归模型在超过 20 个原子后即失效。
• 采样效率具有竞争力：HierDiff 使用 250 次扩散步长时，每样本仅需 3.02 秒，且由于潜在空间维度降低，支持更大的批量处理。
• 可视化结果表明，与原子级基线相比，HierDiff 生成的亚结构（如芳香环）更稳定，几何畸变极小。