Unified 2D and 3D Pre-Training of Molecular Representations

TL;DR

提出统一的2D和3D分子预训练模型,显著提升分子性质预测性能。

cs.LG 🔴 高级 2022-07-14 39 次浏览
Jinhua Zhu Yingce Xia Lijun Wu Shufang Xie Tao Qin Wengang Zhou Houqiang Li Tie-Yan Liu
分子表示 深度学习 多模态预训练 图神经网络 结构预测

核心发现

方法论

该方法结合图神经网络(GNN)对分子2D图和3D构象进行编码,通过引入掩码重建、3D生成和2D生成三项预训练任务,实现信息的互补融合。模型利用掩码语言模型(MLM)机制重建被遮蔽的原子和坐标,采用对称子结构的置换不变损失确保结构的合理性,同时引入旋转平移不变性损失实现3D构象的稳健生成。具体流程包括:编码原子特征、距离和坐标,融合信息后通过多层GN块进行特征更新,最终输出每个原子的表征用于下游任务。

关键结果

  • 在11个下游分子性质预测任务中,模型在10项任务中达到了SOTA,平均提升8.3%,尤其在2D信息有限的任务中表现优异;在4个3D构象生成任务中,平均提升7.7%和3.6%的匹配得分,显著优于先前方法。
  • 在PCQM4Mv2数据集上预训练,模型有效融合2D图和3D构象信息,提升分子表征的表达能力,验证了多模态预训练的优势。
  • 通过引入对称子结构的置换不变性和旋转平移不变性损失,有效解决了结构对称性带来的挑战,增强模型的泛化能力。

研究意义

该研究突破了传统单一模态预训练的局限,首次提出统一模型同时处理2D和3D信息,极大丰富了分子表征的表达能力。其在药物设计、材料科学等领域具有广泛应用前景,能显著提升分子性质预测、构象生成等关键任务的准确性与效率,为深度学习在分子科学中的应用提供了新思路。模型的多任务预训练策略也为多模态学习提供了理论基础和实践范例,推动了分子表征学习的技术革新。

技术贡献

本研究提出了融合2D图和3D构象的统一预训练框架,创新性引入多模态信息融合机制和对称子结构的置换不变性损失,解决了结构对称性和旋转平移不变性问题。采用多任务学习策略,包括掩码重建、3D生成和2D生成,显著提升了模型的泛化能力。模型架构基于改进的图神经网络(GN)模块,能同时处理多模态输入,提供了端到端的分子表征学习方案。这些技术创新为未来多模态分子学习提供了坚实基础。

新颖性

本工作首次实现了单一模型同时编码和生成分子的2D图和3D构象,打破了以往多模态信息分离处理的局限。引入的对称子结构置换不变性和旋转平移不变性损失,确保了模型对结构对称性和空间变换的鲁棒性。这种融合策略在分子预训练领域尚属首创,为多模态信息的深度融合提供了新范式,具有重要的学术和应用价值。

局限性

  • 模型在处理极端对称或高度复杂的分子结构时仍存在一定挑战,可能由于结构多样性不足影响泛化能力。
  • 预训练依赖大量高质量的3D构象数据,数据获取成本较高,限制了模型的普适性。
  • 在极大规模分子库中的推理效率尚待优化,未来需提升模型的计算效率和可扩展性。

未来方向

未来将探索更高效的多模态信息融合机制,提升模型在大规模分子库中的推理速度。还计划引入自监督学习策略,减少对高质量3D数据的依赖,增强模型的泛化能力。此外,将结合药物设计中的反应预测和筛选任务,拓展模型在药物发现中的应用范围,推动分子表示学习的实际落地。

AI 总览摘要

近年来,分子表示学习成为药物设计和材料科学中的核心技术之一。传统方法多依赖单一模态信息,如2D图或3D构象,难以充分捕获分子的复杂结构特征。本文提出了一种创新的统一预训练框架,将分子的2D图和3D构象融合在一个模型中,显著提升了分子表征的表达能力。该模型基于图神经网络(GNN),通过引入掩码重建、3D生成和2D生成三项任务,实现多模态信息的互补融合。特别地,模型设计考虑了结构对称性和空间变换的不变性,采用对称子结构的置换不变性损失和旋转平移不变性损失,确保生成的结构合理且稳健。预训练在庞大的PCQM4Mv2数据集上进行,结果显示在11个分子性质预测任务中,模型在10项任务中达到了最新最优(SOTA),平均性能提升8.3%。在3D构象生成任务中,也取得了显著的性能提升,匹配得分分别提高7.7%和3.6%。这些成果验证了多模态预训练在分子科学中的潜力,为药物设计、材料开发等应用提供了强大工具。未来,模型将进一步优化计算效率,扩展到更大规模的分子库,并结合实际药物筛选流程,推动深度学习在分子领域的广泛应用。

深度分析

研究背景

分子表示学习经历了从传统的手工特征到深度学习的快速发展。早期方法多依赖分子指纹或手工设计的特征,难以捕获复杂结构信息。近年来,Graphormer、SchNet、DimeNet等模型引入图神经网络(GNN)处理分子图,显著提升了性能。同时,3D空间信息的引入进一步丰富了分子表征,诸如Schütt等提出的等变网络确保了空间变换的不变性。预训练技术在自然语言和计算机视觉中取得成功后,也被引入分子领域,推动了分子表征的深度学习应用。现有工作多偏重单一模态,缺乏对多模态信息的系统融合,限制了模型的表达能力。

核心问题

核心问题在于如何有效融合分子的2D图结构和3D空间构象信息,解决信息互补和空间变换不变性的问题。传统方法多采用多模型融合,存在信息分离、模型复杂、训练难度大的弊端。单一模型处理多模态信息的难点在于结构对称性、空间变换的鲁棒性,以及如何设计统一的训练目标。此外,缺乏端到端的预训练策略,限制了模型在下游任务中的性能表现。

核心创新

本研究提出了统一的多模态预训练框架,创新点包括:1)引入多任务学习策略,结合掩码重建、3D生成和2D生成任务,充分利用不同模态信息;2)设计对称子结构的置换不变性损失,解决结构对称性带来的挑战;3)引入旋转平移不变性损失,确保空间变换的鲁棒性;4)采用改进的图神经网络(GN)模块,有效融合原子特征、距离和空间信息。这些创新使模型在表达能力和泛化能力上优于现有方法,推动了多模态分子预训练的发展。

方法详解

  • �� 输入:分子2D图(节点和边)与对应的3D构象(原子坐标);
  • �� 编码:利用嵌入层将原子、键类型、距离和坐标映射到高维空间;
  • �� 融合:通过线性变换和Feed-Forward网络(FF)融合空间信息,形成原子和边的特征表示;
  • �� 多任务预训练:
  • ��– 掩码重建:随机遮蔽部分原子和坐标,利用MLM机制重建;
  • ��– 3D生成:在遮蔽空间中,基于2D图生成3D构象,考虑空间对称性和变换不变性;
  • ��– 2D生成:在遮蔽空间中,基于3D构象重建2D图结构。
  • �� 损失函数:结合掩码重建损失、对称子结构的置换不变性损失、空间变换不变性损失,确保模型稳健性。

实验设计

采用PCQM4Mv2数据集进行预训练,包含3.38M分子样本。下游任务包括7个纯2D性质预测任务(如分子能量、极性等)和4个结合2D/3D信息的毒性预测任务,以及2个3D构象生成任务。模型与传统单模态模型对比,采用平均绝对误差(MAE)、准确率等指标评估性能。进行消融实验验证不同损失和结构设计的贡献,调优超参数确保公平比较。

结果分析

模型在11项任务中获得10项SOTA,平均性能提升8.3%。在2D性质预测任务中,准确率提升显著,特别是在数据有限的任务中表现优异。3D构象生成任务中,匹配得分提升7.7%和3.6%,验证模型在空间结构生成中的优越性。消融实验显示,加入对称性和空间不变性损失显著改善模型鲁棒性和泛化能力。

应用场景

该模型可广泛应用于药物筛选、材料设计、反应预测等领域。只需输入分子2D图或3D构象,即可实现性质预测和构象生成,极大简化分子设计流程。模型的端到端训练和多模态融合能力,为工业界提供了高效、准确的工具,推动新药开发和新材料创新。

局限与展望

模型对极端对称或复杂结构的泛化能力仍有限,且训练依赖大量高质量3D数据,成本较高。推理速度在大规模分子库中有待优化,未来需提升计算效率和模型可扩展性。对某些特殊结构的适应性不足,未来需结合更多结构信息和优化算法。

通俗解读 非专业人士也能看懂

想象你在做一份复杂的拼图。每块拼图不仅有形状(对应分子的结构),还有颜色(对应空间信息)。以前,我们只用一种方式看拼图——只看形状或只看颜色,但这样很难拼出完整的图。现在,这个新方法就像用一种特别的工具,能同时看形状和颜色,把它们结合起来,帮你更快更准确地拼出完整的图。它还能理解对称的部分,比如左右对称的拼图块,即使你旋转或翻转它们,也能正确识别。这就像用一台智能拼图机,不仅能拼出漂亮的图,还能理解各种复杂的结构。这样一来,无论是设计新药还是制造新材料,都能用它来帮忙,变得更快、更准、更省力。

简单解释 像给14岁少年讲一样

想象你在玩一个超级复杂的拼图游戏。有时候,拼图块长得很像,左右对称,旋转一下就变得一样。以前的拼图工具只能看形状或者颜色,不能同时理解两者,所以拼得不太快,也容易出错。现在,这个新工具就像拥有魔法眼睛,能同时看到拼图的形状和颜色,还能理解对称和旋转带来的变化。它可以帮你找到正确的拼图位置,无论你怎么旋转或翻转,都能正确识别。这样一来,拼图变得简单多了,不仅能更快拼完,还能拼出更漂亮的图案。这就像给拼图加上了智能,让你变成拼图高手!在科学里,这个工具可以帮科学家设计新药、制造新材料,让未来变得更精彩。

原文摘要

Molecular representation learning has attracted much attention recently. A molecule can be viewed as a 2D graph with nodes/atoms connected by edges/bonds, and can also be represented by a 3D conformation with 3-dimensional coordinates of all atoms. We note that most previous work handles 2D and 3D information separately, while jointly leveraging these two sources may foster a more informative representation. In this work, we explore this appealing idea and propose a new representation learning method based on a unified 2D and 3D pre-training. Atom coordinates and interatomic distances are encoded and then fused with atomic representations through graph neural networks. The model is pre-trained on three tasks: reconstruction of masked atoms and coordinates, 3D conformation generation conditioned on 2D graph, and 2D graph generation conditioned on 3D conformation. We evaluate our method on 11 downstream molecular property prediction tasks: 7 with 2D information only and 4 with both 2D and 3D information. Our method achieves state-of-the-art results on 10 tasks, and the average improvement on 2D-only tasks is 8.3%. Our method also achieves significant improvement on two 3D conformation generation tasks.

cs.LG cs.AI