[论文解读] On Graph Neural Network Ensembles for Large-Scale Molecular Property Prediction

Edward Elson Kosasih, Joaquín Cabezas+6 more|arXiv (Cornell University)|Jun 29, 2021

Machine Learning in Materials Science被引用 7

一句话总结

本文提出了一种由三种图神经网络（GNN）模型组成的集成方法——GIN-Virtual、贝叶斯GNN和DiffPool增强型GNN——通过多次随机初始化进行训练，以提升大规模分子性质预测性能。该集成方法在PCQM4M-LSC数据集上将平均绝对误差（MAE）降低至0.1290，相比基线模型相对提升了7.6%，并表明预测不确定性与预测误差之间存在强烈相关性（皮尔逊相关系数R = 0.5181），为主动学习提供了潜在应用价值。

ABSTRACT

In order to advance large-scale graph machine learning, the Open Graph Benchmark Large Scale Challenge (OGB-LSC) was proposed at the KDD Cup 2021. The PCQM4M-LSC dataset defines a molecular HOMO-LUMO property prediction task on about 3.8M graphs. In this short paper, we show our current work-in-progress solution which builds an ensemble of three graph neural networks models based on GIN, Bayesian Neural Networks and DiffPool. Our approach outperforms the provided baseline by 7.6%. Moreover, using uncertainty in our ensemble's prediction, we can identify molecules whose HOMO-LUMO gaps are harder to predict (with Pearson's correlation of 0.5181). We anticipate that this will facilitate active learning.

研究动机与目标

• 为了提升在包含380万个分子的OGB-LSC PCQM4M-LSC数据集上的大规模分子性质预测性能。
• 为解决图机器学习在真实世界大规模数据集中的可扩展性与性能挑战。
• 探究集成预测中的不确定性是否能指示高预测误差区域，以支持主动学习。
• 开发一种可扩展的集成框架，结合多种GNN架构与训练运行，以提升回归性能。

提出的方法

• 集成多个GNN弱学习器：GIN-Virtual、贝叶斯GNN（BNN）和DiffPool增强型GNN，每种模型均从不同随机初始化训练三次。
• 通过各模型多次运行的平均预测结果形成独立集成（例如GIN-Virtual-ensemble）。
• 将所有模型与所有运行的预测结果整合为最终的ensemble-all，采用9次独立实验的加权平均。
• 在GNN的读出层中应用贝叶斯反向传播（BBP），通过贝叶斯线性层引入不确定性估计。
• 通过DiffPool实现图池化，将图压缩为5个节点，以保留结构信息并提升表征能力。
• 将不确定性定义为9次集成实验中预测结果的标准差，用作模型置信度的代理指标。

实验结果

研究问题

• RQ1 通过集成多种GNN架构，是否能在大规模分子图上超越单个模型，提升回归性能？
• RQ2 从集成差异中估计的不确定性是否与HOMO-LUMO能隙预测的实际误差存在相关性？
• RQ3 所提出的集成方法是否能在PCQM4M-LSC数据集上显著优于OGB-LSC基线？
• RQ4 不确定性估计在多大程度上可指导主动学习，识别难以预测的分子？

主要发现

• ensemble-all模型在验证集上的平均绝对误差（MAE）为0.1290，相比基线MAE 0.1396实现了7.6%的相对提升。
• 不确定性估计（即9次集成实验中预测结果的标准差）与实际预测误差之间的皮尔逊相关系数为0.5181。
• GIN-Virtual-BNN-ensemble变体在所有独立集成类型中取得了最低的验证集MAE，为0.1329。
• ensemble-all模型优于所有单个模型及基线模型，包括在挑战赛中表现最佳的GIN-Virtual模型。
• 不确定性与误差之间的相关性表明，不确定性估计可识别出更难预测的分子，为未来主动学习应用提供支持。
• 该模型在300万个分子上进行训练，使用高内存云GPU实例，证明了其在大规模分子数据集上的可扩展性。