[论文解读] 3D Infomax improves GNNs for Molecular Property Prediction

H. Stärk, Dominique Beaini+5 more|arXiv (Cornell University)|Oct 8, 2021

Machine Learning in Materials Science参考文献 53被引用 85

一句话总结

本文提出了 3D Infomax，一种自监督预训练方法，教会 2D GNNs 对 3D 分子几何进行隐式推理，从而改善下游的量子性质预测，并实现跨数据集的泛化。

ABSTRACT

Molecular property prediction is one of the fastest-growing applications of deep learning with critical real-world impacts. Including 3D molecular structure as input to learned models improves their performance for many molecular tasks. However, this information is infeasible to compute at the scale required by several real-world applications. We propose pre-training a model to reason about the geometry of molecules given only their 2D molecular graphs. Using methods from self-supervised learning, we maximize the mutual information between 3D summary vectors and the representations of a Graph Neural Network (GNN) such that they contain latent 3D information. During fine-tuning on molecules with unknown geometry, the GNN still generates implicit 3D information and can use it to improve downstream tasks. We show that 3D pre-training provides significant improvements for a wide range of properties, such as a 22% average MAE reduction on eight quantum mechanical properties. Moreover, the learned representations can be effectively transferred between datasets in different molecular spaces.

研究动机与目标

• 在无法获得显式 3D 结构的大规模分子属性预测上提供动力。
• 开发一种预训练方案，使 2D GNNs 能从 2D 图中编码潜在的 3D 信息。
• 证明所学习的 3D 感知表示能提升下游的量子性质预测。
• 展示 3D Infomax 表示在不同分子空间和数据集间具有泛化能力。
• 探讨在预训练过程中使用多个构象的潜在收益，以进一步提升性能。）
• method':['在 2D 分子图上预训练一个 2D GNN f^a，通过最大化与来自构象的 3D 表示 z^b 的互信息来产生 z^a。','使用对比性的 NT-Xent 风格损失来对齐同一分子的 z^a 与 z^b，并排斥不匹配的对。','通过在计算对比损失 (L^multi3D) 时将一个分子的所有构象作为正样本来引入多个构象（c）。','3D 网络将 3D 坐标处理为点云，并通过基于距离的 MPNN 及正弦特征映射 γ(d_uv) 将成对距离编码为 SE(3)-不变的 z^b。','仅使用 2D 输入对 2D 网络在下游任务上进行微调，利用隐式学习到的 3D 信息来提升预测性能。'],
• research_questions':['2D GNN 是否能通过与 3D 构象的对比预训练来从 2D 分子图中编码潜在的 3D 信息？','与仅使用 2D 的预训练基线相比，3D Infomax 预训练模型是否能改进量子性质预测（如 QM9）？','所学习的表示是否在具有不同尺寸和组成的分子空间之间具有可迁移性？','在预训练中利用多个构象是否能在下游性质上带来额外提升？'],
• key_findings':['3D Infomax 在量子性质预测方面取得显著提升，QM9 属性的平均 MAE 降幅约为 22%。','预训练表示在不同分子空间间具备泛化性，在从小分子迁移到类似药物分子的场景以及相反场景时均显示出收益。','在预训练中使用多个构象可带来额外的性能提升，但超过三个构象后收益递减。','在所评估的数据集中，3D Infomax 的表现超越了诸如 GraphCL 和距离预测等传统 SSL 基线。','该方法未表现出负迁移，与某些先前的预训练方法不同。'],
• table_headers: []
• table_rows: []

提出的方法

• Pre-train a 2D GNN f^a on 2D molecular graphs to produce z^a by maximizing mutual information with a 3D representation z^b derived from conformers.
• Use a contrastive NT-Xent style loss to align z^a and z^b for same molecules and repel non-matching pairs.
• Incorporate multiple conformers (c) by treating all conformers of a molecule as positive examples when computing the contrastive loss (L^multi3D).
• The 3D network processes 3D coordinates as a point cloud and encodes pairwise distances into an SE(3)-invariant z^b via a distance-based MPNN with sinusoidal feature mappings γ(d_uv).
• Fine-tune the 2D network on downstream tasks using only 2D inputs, leveraging the implicitly learned 3D information for improved predictions.

实验结果

研究问题

• RQ1 Can a 2D GNN learn to encode latent 3D information from 2D molecular graphs through contrastive pre-training with 3D conformers?
• RQ2 Do 3D Infomax pre-trained models improve quantum property predictions (e.g., QM9) compared to 2D-only pre-training baselines?
• RQ3 Are the learned representations transferable across molecular spaces with different sizes and compositions?
• RQ4 Does leveraging multiple conformers during pre-training yield additional gains in downstream properties?

主要发现

• 3D Infomax achieves large improvements in quantum property predictions, including a ~22% average MAE reduction on QM9 properties.
• Pre-trained representations generalize across distinct molecular spaces, showing gains when transferring from small to drug-like molecules and vice versa.
• Using multiple conformers during pre-training provides additional performance gains, with diminishing returns beyond three conformers.
• 3D Infomax outperforms conventional SSL baselines like GraphCL and distance-prediction pre-training across evaluated datasets.
• The method does not exhibit negative transfer, unlike some prior pre-training approaches.