人工智能技术正在深刻重塑药物研发的进程，其中分子性质预测与分子生成作为两大核心任务，长期以来却沿着独立的技术路径发展。分子性质预测的目的是给定分子结构信息，预测分子多样的化学、生物性质，加速药物筛选。分子生成旨在估计分子数据分布，潜在地学习原子间相互作用与构象信息，能够从头生成化学合理的新分子，开拓药物设计的可能性边界。尽管近年来在这些领域已经进行了大量研究，但它们基本上是独立发展的，始终未能有效打通这两个关键环节的协同通道。
针对于此，清华大学联合中科院团队提出 UniGEM 模型，首次基于扩散模型实现两任务协同增强。研究团队指出，生成与性质预测依赖于有效的分子表示而高度相关。团队创新地提出两阶段生成过程，克服了传统联合训练中的不一致性，为分子生成与性质预测领域开辟新路径。该成果以「UniGEM: A Unified Approach to Generation and Property Prediction for Molecules」为题，入选 ICLR 2025 。
论文地址：
https://openreview.net/pdf?id=Lb91pXwZMR
QM9 量子化学数据集：
GEOM-Drugs 3D 分子构象数据集：
开源项目「awesome-ai4s」汇集了 200 余篇 AI4S 论文解读，并提供海量数据集与工具：
https://github.com/hyperai/awesome-ai4s
生成与预测任务统一的动机
研究团队认为，生成与预测任务的本质均在于学习分子表示。一方面，各种分子预训练方法的有效性表明，分子性质预测依赖于稳健的分子表示作为基础。另一方面，分子生成需要对分子结构有深入的理解，从而能够在生成过程中创建良好的表示。
最近的研究结果为这一观点提供了支持。例如，计算机视觉领域的工作表明，扩散模型本身具有学习有效图像表征的能力。在分子领域，研究表明生成式预训练可以增强分子性质预测任务，不过这些方法通常需要额外的微调才能实现最佳预测性能。此外，虽然预测器可以通过 classifier guidance 方法指导分子生成，但预测器的训练是否可以直接提高生成性能仍不清楚。
因此，现有研究尚未充分阐明生成任务和预测任务之间的关系，这提出了一个关键问题：我们能否构建统一模型，实现生成与预测任务的协同增强？
传统方法失败的原因分析
将这两项任务结合的一个直接方法是使用传统的多任务学习框架，其中模型同时优化生成损失与预测损失。然而，该研究团队的实验表明，这种方法显著降低了生成任务和性质预测任务的性能 (生成稳定性骤降 6%，预测误差同步增大 1 倍以上) 。即使冻结生成模型的权重，并为性质预测任务添加一个单独的 head 以保持生成性能，研究人员观察到与从头开始训练相比，性质预测的性能并没有改善。
研究人员将传统方法结果不理想归因于生成任务和预测任务之间固有的不一致性。在扩散生成过程中，分子结构需要历经从无序噪声到精细结构的渐进重构。然而，在预测任务中，有意义的分子性质只有在分子结构基本建立后才能定义。因此，仅仅采用简单的多任务优化方法会导致在早期扩散阶段，高度无序的分子构象会与性质标签形成错误关联，从而对分子生成和属性预测产生负面影响。
为了进一步说明这一观点，研究人员对扩散训练过程中去噪网络内的中间表示与目标分子之间的互信息进行了理论分析。进而从理论上证明，扩散模型隐式地最大化了中间表示与目标分子互信息的下界，表明具有扩散模型表示学习的能力。然而，中间表示与目标分子间的互信息呈现单调递减趋势，在较大时间步时趋近于零，这意味着无序阶段的中间表示无法支持有效预测。因此，直觉和理论都表明，只有在较小的时间步，即当分子仍然保持较为有序时，生成任务和预测任务才能对齐。
两阶段扩散生成机制
基于上述分析，研究团队提出了一种新颖的两阶段生成方法，旨在统一分子的性质预测和生成，如下图所示。
图注：UniGEM 的两阶段生成过程
研究人员将分子生成过程分为两个阶段，即「分子成核阶段」和「分子生长阶段」，这一划分灵感来源于物理学中的晶体形成过程。
在分子成核阶段，分子从完全无序状态形成其骨架，随后基于该骨架生长出完整的分子。这两个阶段通过「成核时间」来分隔。研究人员引入了一种新的分子生成方式来表述这两个阶段。其中，在「成核时间」之前，扩散模型逐步生成分子坐标；在成核之后，模型继续调整分子坐标，同时优化性质与原子类型预测损失。
不同于传统生成模型通常进行原子类型和坐标的联合扩散，这一创新方法仅专注于坐标的扩散，将原子类型作为单独的预测任务。因为研究人员观察到，原子类型通常可以根据成型的分子坐标推断出来。具体来说，在成核之前，扩散过程的目标是重建坐标；成核之后，其将原子类型和性质的预测损失整合进统一的学习框架中。
UniGEM 的训练策略
为了便于与传统的联合扩散方法进行比较，研究人员采用与 E(3) 等变扩散模型 (EDM) 使用 EGNN 作为网络结构骨架。其中，生长阶段只占整个训练过程的约 1% 。如果按照标准的扩散训练程序，并均匀地采样时间步长，预测任务的迭代次数仅占总训练过程的 1%，这将显著降低模型在该任务上的表现。因此，为了确保对预测任务的充分训练，研究人员在生长阶段对时间步长进行过采样。
然而，研究人员观察到，过采样可能导致整个时间步长范围内的训练不平衡，进而影响生成过程的质量。为了解决这个问题，其提出了一种多分支网络架构。该网络在浅层共享参数，但在更深的层次分裂成两个分支，每个分支拥有一套独立的参数。这些分支在训练的不同阶段激活：一个分支专注于成核阶段，另一个分支处理生长阶段，如下图所示。该设计确保了预测任务和生成任务能够有效训练，而不会相互影响。
图注：UniGEM 的训练过程
UniGEM 的推理过程
图注：UniGEM 的分子生成过程
在 UniGEM 中，分子生成是通过反向扩散过程来重建原子坐标，然后基于生成的坐标预测原子类型，如图。对于性质预测，固定网络输入的时间步为零，并使用性质预测 head 。值得注意的是，这种方法对于生成任务和预测任务都不会产生额外的计算开销，总推理时间与基线相同。
对于分子生成任务，研究人员还分析了 UniGEM 与传统联合生成方法在生成误差上的差异。首先，其观察到在 UniGEM 中，原子类型预测损失的误差小于联合生成中的原子类型去噪生成损失。其次，在联合生成过程中，坐标生成会受到原子类型预测结果震荡的影响，进而导致误差增大。最后，联合生成方法还会引入更大的初始分布误差和离散化误差。这些因素共同解释了 UniGEM 如何实现更优的生成效果。
实验结果：在分子生成与性质预测任务上均优于基线模型
分子生成：UniGEM 优于基准模型
研究人员首先在 QM9 和 GEOM-Drugs 数据集上比较了基于 EDM 改造的 UniGEM 与 EDM 变体，UniGEM 几乎在所有评估指标上都优于基准模型，如下图所示。值得注意的是，与其他 EDM 变体相比，UniGEM 显著更简单，因为它既不依赖于先验知识，也不需要额外的自动编码器训练，但它的性能优于 EDM-Bridge 和 GeoLDM，突显了 UniGEM 的优势。
图注：UniGEM 的生成效果
为了展示 UniGEM 在适应各种生成算法方面的灵活性，研究人员将 UniGEM 应用于贝叶斯流网络（BFN），在 QM9 数据集上超过了联合生成坐标与原子类型的 GeoBFN，从而取得了 SOTA 结果。
此外，研究人员还测试了 UniGEM 在条件生成任务中的表现，通过在采样过程中使用模型自带的性质预测模块作为指导，避免了重新训练条件生成模型的需求。
分子性质预测：UniGEM 超越了大多数预训练方法
图注：UniGEM 的性质预测效果
研究人员在 QM9 数据集上评估了 UniGEM 性质预测的效果，采用测试集上的平均绝对误差 (MAE) 作为评估指标。如图，UniGEM 显著优于从头训练的 EGNN，证明了统一建模的有效性。令人惊讶的是，尽管这些先进的预训练方法利用了额外的大规模预训练数据集，UniGEM 仍然超越了大多数预训练方法。这突显了其生成与预测统一模型的优势，能够在生成过程中有效利用分子表示学习的能力，而无需额外的数据和预训练步骤。
结语
UniGEM 模型统一了分子生成和性质预测任务，并显著提高了分子生成和性质预测的性能。 UniGEM 的增强性能得到了坚实的理论分析和全面的实验研究支持。我们相信，创新的两阶段生成过程及其相应的模型为分子生成框架的发展提供了新的范式，并可能激发更多高级分子生成框架的开发，进而惠及更多特定应用领域的分子生成。
本研究由 ATOM Lab 主导，在分子预训练、分子生成、蛋白质结构预测、虚拟筛选等领域，该团队还有更多研究成果，欢迎关注！
欢迎访问 ATOM Lab 主页：
https://atomlab.yanyanlan.com/
论文作者简介：
* 兰艳艳，清华大学智能产业研究院（AIR）教授，研究方向为 AI4Science 、机器学习、自然语言处理。
* 冯世坤，清华大学智能产业研究院 (AIR) 在读博士生，研究方向包括表示学习，生成模型和 AI4Science 。
* 倪雨嫣，中国科学院数学与系统科学研究院 (AMSS) 在读博士生，研究兴趣包括生成模型、表示学习、 AI4Science 和深度学习理论。
本论文的主要作者冯世坤和倪雨嫣博士目前正在寻找工作机会，感兴趣的朋友可以去联系他们哦～
* 冯世坤邮箱：mailto:fsk21@mails.tsinghua.edu.cn
* 倪雨嫣邮箱：mailto:niyuyan17@mails.ucas.ac.cn
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用。你还可以使用@来通知其他用户。