[Paper Review] First Place Solution of KDD Cup 2021 & OGB Large-Scale Challenge Graph Prediction Track.

Chengxuan Ying, Mingqi Yang+8 more|arXiv (Cornell University)|Jun 15, 2021

Advanced Graph Neural Networks11 references4 citations

TL;DR

This paper presents the winning solution for the KDD Cup 2021 OGB-LSC PCQM4M-LSC graph prediction track, achieving 0.1200 MAE on the test set. The approach combines Graphormer and ExpC⋆ models, enhanced with domain-specific molecular features including RDKit-calculated 3D Euclidean distances and RBF-embedded bond lengths, trained via 8-fold cross-validation and ensembled via simple averaging to improve generalization and performance.

ABSTRACT

In this technical report, we present our solution of KDD Cup 2021 OGB Large-Scale Challenge - PCQM4M-LSC Track. We adopt Graphormer and ExpC as our basic models. We train each model by 8-fold cross-validation, and additionally train two Graphormer models on the union of training and validation sets with different random seeds. For final submission, we use a naive ensemble for these 18 models by taking average of their outputs. Using our method, our team MachineLearning achieved 0.1200 MAE on test set, which won the first place in KDD Cup graph prediction track.

Motivation & Objective

• To develop a high-accuracy graph neural network model for predicting HOMO-LUMO energy gaps in molecules using 2D molecular graphs.
• To improve generalization and reduce overfitting in large-scale graph regression by incorporating 3D molecular structural features derived from RDKit.
• To achieve state-of-the-art performance on the PCQM4M-LSC dataset, a large-scale quantum chemistry graph benchmark.
• To demonstrate the effectiveness of model ensembling and feature engineering in graph representation learning for molecular property prediction.

Proposed method

• Employed Graphormer with 12 layers, 768 hidden dimensions, and 32 attention heads, using Euclidean distances as spatial encoding features.
• Applied a Gaussian RBF kernel with 256 basis functions to transform raw Euclidean distances into more generalizable representations.
• Enhanced edge features by adding Laplace noise to bond distances, calibrated on QM9 data to reduce overfitting.
• Used a modified version of ExpC⋆ with expanded hidden dimensions (d′ = 1200) and nonlinear activation before aggregation to improve representational power.
• Trained models using 8-fold cross-validation on the union of training and validation sets, with additional training on full training+validation data with different random seeds.
• Ensembled 18 models (16 Graphormer and 2 ExpC⋆ variants) via simple averaging of predictions to improve robustness and final performance.

Experimental results

Research questions

• RQ1 How can graph neural networks be enhanced with 3D molecular structural features to improve prediction accuracy on quantum chemistry property tasks?
• RQ2 What impact do RBF-embedded spatial features have on model generalization in graph regression tasks?
• RQ3 Can model ensembling with diverse training runs and architectures outperform individual models on large-scale molecular graph prediction?
• RQ4 How do different hyperparameter settings, such as embedding dropout, affect performance on small molecular graphs with high sensitivity to node representations?
• RQ5 To what extent can RDKit-generated 3D coordinates improve model performance despite their lower accuracy compared to DFT-optimized structures?

Key findings

• The final model ensemble achieved a test MAE of 0.1200, securing first place in the KDD Cup 2021 OGB-LSC graph prediction track.
• Graphormer models fine-tuned with RBF-embedded Euclidean distances and Laplace-augmented bond features showed improved generalization and lower validation MAE.
• The use of 3D molecular structure features from RDKit, despite lower accuracy than DFT, provided meaningful inductive bias that improved model performance.
• Setting embedding dropout to 0.0 instead of the default 0.1 significantly improved performance, likely due to the small size of molecular graphs (median ~15 atoms).
• The naive ensemble of 18 models, including multiple Graphormer and ExpC⋆ variants trained with different seeds and data splits, yielded the best generalization and final score.
• Minor differences in model outputs (MAE ≤ 1e-4) were observed across different PyTorch Geometric versions, highlighting the importance of reproducibility in large-scale ML experiments.