[论文解读] GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction

Dominic Masters, Josef Dean+8 more|arXiv (Cornell University)|Nov 18, 2022

Machine Learning in Materials Science被引用 10

一句话总结

GPS++ 是一种混合 MPNN/Transformer 模型，专为 PCQM4Mv2 数据集上的分子性质预测而优化，通过整合 3D 原子坐标和去噪损失来提升性能。其在测试集上的平均绝对误差为 0.0719，位列 OGB-LSC 2022 挑战赛第一，使用 Graphcore IPU 加速和 112 个模型集成，推理耗时 1 小时 32 分钟。

ABSTRACT

This technical report presents GPS++, the first-place solution to the Open Graph Benchmark Large-Scale Challenge (OGB-LSC 2022) for the PCQM4Mv2 molecular property prediction task. Our approach implements several key principles from the prior literature. At its core our GPS++ method is a hybrid MPNN/Transformer model that incorporates 3D atom positions and an auxiliary denoising task. The effectiveness of GPS++ is demonstrated by achieving 0.0719 mean absolute error on the independent test-challenge PCQM4Mv2 split. Thanks to Graphcore IPU acceleration, GPS++ scales to deep architectures (16 layers), training at 3 minutes per epoch, and large ensemble (112 models), completing the final predictions in 1 hour 32 minutes, well under the 4 hour inference budget allocated. Our implementation is publicly available at: https://github.com/graphcore/ogb-lsc-pcqm4mv2.

研究动机与目标

• 开发一种在大规模 PCQM4Mv2 数据集上高效且准确的分子性质预测模型。
• 利用 3D 原子坐标和归纳偏置，提升图神经网络的泛化能力并减少过平滑现象。
• 在严格的推理时间约束（4 小时预算）下，通过硬件加速的训练与推理实现最先进性能。
• 证明混合 MPNN/Transformer 架构可在参数量显著更少的情况下，达到或超越纯 Transformer 模型的性能。
• 构建由 112 个模型组成的多样化集成，以最大化在测试挑战集上的预测鲁棒性与泛化能力。

提出的方法

• GPS++ 将深层消息传递神经网络（MPNN）与带偏置的自注意力机制相结合，以平衡局部归纳偏置与全局信息流动。
• 通过分组输入掩码策略将 3D 原子坐标整合进注意力机制，实现距离特征的融合。
• 在预训练阶段应用去噪损失，以缓解深层架构中的过平滑问题并提升表征学习能力。
• 采用混合架构设计：MPNN 层负责局部消息传递，注意力层实现长距离依赖，且注意力机制偏向邻近原子。
• 利用 Graphcore 的 IPU 硬件加速训练与推理，实现每轮 3 分钟的训练速度，并在单个 IPU 和 AMD EPYC CPU 上完成 112 个模型集成的推理，耗时 1 小时 32 分钟。
• 通过调整输入特征、dropout 率和掩码策略，构建由 112 个多样化模型组成的集成，以提升泛化能力与鲁棒性。

实验结果

研究问题

• RQ1 混合 MPNN/Transformer 架构是否能在参数量更少的情况下，优于纯 Transformer 模型，实现更优的分子性质预测性能？
• RQ2 通过分组输入掩码整合 3D 原子坐标，在提升模型性能与泛化能力方面效果如何？
• RQ3 去噪预训练目标在多大程度上能缓解深度图神经网络在分子图上的过平滑问题？
• RQ4 使用 IPU 的硬件加速是否能实现在严格时间约束下，对深度模型和大规模集成的可扩展训练与推理？
• RQ5 通过超参数变化实现的模型集成多样性，是否能显著提升在 PCQM4Mv2 测试挑战集上的最终预测性能？

主要发现

• GPS++ 在 PCQM4Mv2 数据集上实现了测试挑战集平均绝对误差（MAE）0.0719，位列 OGB-LSC 2022 挑战赛前三名。
• 最终的 112 个模型集成在单个 IPU 和 AMD EPYC CPU 上完成推理，耗时 1 小时 32 分钟，远低于 4 小时的预算限制。
• 该模型以仅 64% 的参数量（4430 万 vs. 6900 万），达到了与最先进 Transformer 模型（Luo et al., 2022）相当的性能。
• 通过分组掩码引入 3D 距离特征显著提升了注意力机制的有效性，消融实验证明，若省略 3D 特征，性能提升微乎其微。
• 通过整合多样化配置（包括输入特征、dropout 和掩码策略）的模型，MAE 从单个模型的 0.0755 降低至代理集的 0.0722，最终在测试集上达到 0.0719。
• 模型在 IPU 硬件上实现每轮 3 分钟的训练速度，支持快速迭代，并在 24 小时内完成深度（16 层）架构的训练。