OGB-LSC: A Large-Scale Challenge for
Machine Learning on Graphs

Practical relevance and dataset overview. The volume of scientific publication has been increasing exponentially, doubling every 12 years (Dong et al., 2017). Currently, subject areas of arXiv papers are manually determined by the paper’s authors and arXiv moderators. An accurate automatic predictor of papers’ subject categories not only reduces the significant burden of manual labeling, but can also be used to classify the vast number of non-arXiv papers, thereby allowing better search and organization of academic papers.

MAG240M is a heterogeneous academic graph extracted from the Microsoft Academic Graph (MAG) (Wang et al., 2020). Given arXiv papers situated in the heterogeneous graph, whose schema diagram is illustrated in Figure 2, we aim to automatically annotate their topics, i.e., predicting the primary subject area of each arXiv paper.

Graph. We extract 121M academic papers in English from MAG (version: 2020-11-23) to construct a heterogeneous academic graph. The resultant paper set is written by 122M author entities, who are affiliated with 26K institutes. Among these papers, there are 1.3 billion citation links captured by MAG. Each paper is associated with its natural language title and most papers’ abstracts are also available. We concatenate the title and abstract by period and pass it to a RoBERTa sentence encoder (Liu et al., 2019; Reimers and Gurevych, 2019), generating a 768-dimensional vector for each paper node. Among the 121M paper nodes, approximately 1.4M nodes are arXiv papers annotated with 153 arXiv subject areas, e.g., cs.LG (Machine Learning). On the paper nodes, we attach the publication years as meta information.

Prediction task and evaluation metric. The task is to predict the primary subject areas of the given arXiv papers, which is cast as an ordinary multi-class classification problem. The metric is the classification accuracy.

To understand the relation between the prediction task and the heterogeneous graph structure, we analyze the graph homophily (McPherson et al., 2001)—tendency of two adjacent nodes to share the same labels—to better understand the interplay between heterogeneous graph connectivity and the prediction task. Homophily is normally analyzed over a homogeneous graph, but we extend the analysis to the heterogenous graph by considering meta-paths (Sun et al., 2011)—a path consisting of a sequence of relations defined between different node types. Given a meta-path, we can say two nodes are adjacent if they are connected by the meta-path. Table 3 shows the homophily for different kinds of meta-paths with different levels of connection strength. Compared to the direct citation connection (i.e., P-P), certain meta-paths (i.e., P-A-P) give rise to much higher degrees of homophiliness, while other meta-paths (i.e., P-A-I-A-P) provide much less homophily. As homophily is the central graph property exploited by many graph ML models, we believe that discovering essential heterogeneous connectivity is important to achieve good performance on this dataset.

Dataset split. We split the data according to time. Specifically, we train models on arXiv papers published until 2018, validate the performance on the 2019 papers, and finally test the performance on the 2020 papers. The split reflects the practical scenario of helping the authors and moderators annotate the subject areas of the newly-published arXiv papers.

Baseline. We benchmark a broad range of graph ML models in both homogeneous (where only paper to paper relations are considered) and full heterogeneous settings. For both settings, we convert the directed graph into an undirected graph for simplicity. First, for the homogeneous setting, we benchmark the simple baseline models: graph-agnostic MLP, Label Propagation, and the recently-proposed simplified graph methods: SGC (Wu et al., 2019), SIGN (Rossi et al., 2020) and MLP+C&S (Huang et al., 2020), which are inherently scalable by decoupling predictions from propagation. Furthermore, we benchmark state-of-the-art expressive GNNs trained with neighborhood sampling (NS) (Hamilton et al., 2017), where we recursively sample 25 neighbors in the first layer and 15 neighbors in the second layer during training time. At inference time, we sample at most 160 neighbors for each layer. Here, we benchmark two types of strong models: the GraphSAGE (Hamilton et al., 2017) model (performing mean aggregation and utilizing skip-connections), and the more advanced Graph Attention Network (GAT) model (Velickovic et al., 2018). For the full heterogeneous setting, we follow Schlichtkrull et al. (2018) and learn distinct weights for each individual relation type (denoted by R-GraphSAGE and R-GAT, where “R” stands for “Relational”). We obtain the input features of authors and institutions by averaging the features of papers belonging to the same author and institution, respectively. The models are trained with NS. We note that the expressive GNNs trained with NS require more efforts to scale up, but are more expressive than the simple baselines.

Hyper-parameters. Hyper-parameters are selected based on their best validation performance. For all the models without NS, we tuned the hidden dimensionality $\in\{128,256,512,1024\}$, MLP depth $\in\{1,2,3,4\}$, dropout ratio $\in\{0,0.25,0.5\}$, propagation layers (for SGC, SIGN, and C&S) $\in\{2,3\}$. For all the GNN models with NS, we use a hidden dimensionality of 1024. We make use of batch normalization (Ioffe and Szegedy, 2015) and ReLU activation in all models.

Discussion. Validation and test performances of all models considered are shown in Table 2. First, the graph-agnostic MLP and Label Propagation algorithm perform poorly, indicating that both graph structure and feature information are indeed important for the given task. Across the graph ML models operating on the homogeneous paper graph, GNNs with NS perform the best, with slight gains compared to their simplified versions. In particular, the advanced expressive graph attention aggregation is favourable compared to the uniform mean aggregation in GraphSAGE. Furthermore, considering all available heterogeneous relational structure in the heterogeneous graph setting yields significant improvements, with performance gains up to 3 percentage points. Again, the advanced attention aggregation provides favorable performance. Overall, our experiments highlight the benefits of developing and evaluating advanced expressive models on the larger scale.

KDD Cup 2021 summary. In Table 2, we show the results of the top 3 winners of the KDD Cup: BD-PGL Team (Shi et al., 2021), Academic Team (Addanki et al., 2021), and Synerise AI Team (Daniluk et al., 2021). All the solutions outperform our baselines significantly, yielding 5–6% gain in test accuracy. For a more direct comparison, with a single model (no model ensembling), the BD-PGL Team reports a validation accuracy of 73.71% (Shi et al., 2021), improving our best R-GAT baseline by 3.7%.

Notably, all the winning solutions used the target labels as input to their models, which allows the models to propagate labels together with the features. Regarding the GNN architectures, the BD-PGL adopted the expressive Transformer-based UniMP architecture (Shi et al., 2020), while the Academic adopted the standard MPNN (Gilmer et al., 2017) but trained it with self-supervised contrastive learning on unlabeled paper nodes (Thakoor et al., 2021). These results suggest that expressive GNNs are indeed promising for this dataset. Finally, both the BD-PGL and Academic teams exploited the temporal aspect of the academic graph by using the publication years either as input positional encoding (Shi et al., 2021) or as a way to sample mini-batch subgraphs for GNNs (Addanki et al., 2021). As real-world large-scale graphs are almost always dynamic, exploiting the temporal information is a promising direction of future research.

Practical relevance and dataset overview. Large encyclopedic Knowledge Graphs (KGs), such as Wikidata (Vrandečić and Krötzsch, 2014) and Freebase (Bollacker et al., 2008), represent factual knowledge about the world through triplets connecting different entities, e.g., Hinton $\xrightarrow{citizen~{}of}$ Canada. They provide rich structured information about many entities, aiding a variety of knowledge-intensive downstream applications such as information retrieval, question answering (Singhal, 2012), and recommender systems (Guo et al., 2020). However, these large KGs are known to be far from complete (Min et al., 2013), missing many relational information between entities.

WikiKG90M is a Knowledge Graph (KG) extracted from the entire Wikidata knowledge base. The task is to automatically impute missing triplets that are not yet present in the current KG. Accurate imputation models can be readily deployed on the Wikidata to improve its coverage.

Graph. Each triplet (head, relation, tail) in WikiKG90M represents an Wikidata claim, where head and tail are the Wikidata items, and relation is the Wikidata predicate. We extracted triplets from the public Wikidata dump downloaded at three time-stamps: September 28, October 26, and November 23 of 2020, for training, validation, and testing, respectively. We retain all the entities and relations in the September dump, resulting in 87,143,637 entities, 1,315 relations, and 504,220,369 triplets in total.

In addition to extracting triplets, we provide text features for entities and relations. Specifically, each entity/relation in Wikidata is associated with a title and a short description, e.g., one entity is associated with the title ‘Geoffrey Hinton‘ and the description ‘computer scientist and psychologist‘. Similar to MAG240M, we provide RoBERTa embeddings (Reimers and Gurevych, 2019; Liu et al., 2019) as node and edge features.³³3We concatenate the title and description with comma, e.g., ‘Geoffrey Hinton, computer scientist and psychologist‘, and pass the sentence to a RoBERTa sentence encoder (Note that the RoBERTa model was trained before September 2020, so there is no obvious information leak). The title or/and description are sometimes missing, in which case we simply use the blank sentence to replace it.

Prediction task and evaluation metric. The task is the KG completion, i.e., given a set of training triplets, predict a set of test triplets. For evaluation, we follow the protocol similar to how KG completion is evaluated (Bordes et al., 2013). Specifically, for each validation/test triplet, (head, relation, tail), we corrupt tail with randomly-sampled 1000 negative entities, e.g., tail_neg, such that (head, relation, tail_neg) does not appear in the train/validation/test KG. The model is asked to rank the 1001 candidates (consisting of 1 positive and 1000 negatives) for each triplet and predict the top 10 entities that are most likely to be positive. The goal is to rank the ground-truth positive entity as high in the rank as possible, which is measured by Mean Reciprocal Rank (MRR). ⁴⁴4Note that this is more strict than the standard MRR since there is no partial score for positive entities being ranked outside of top 10.

Dataset split. We split the triplets according to time, simulating a realistic KG completion scenario of imputing missing triplets not present at a certain timestamp. Specifically, we construct three KGs using the aforementioned September, October, and November KGs, where we only retain entities and relation types that appear in the earliest September KG. We use the triplets in the September KG for training, and use the additional triplets in the October and November KGs for validation and test, respectively.

We analyze the effect of the time split. We find that head entities of validation triplets tend to be less popular entities; on average, they only have 6.5 out-degrees in the training KG, which is less than a quarter of the out-degree averaged over training triplets (i.e., 28.0). This suggests that learning signals for predicting validation (and test) triplets are sparse. Nonetheless, even for the sparsely-connected triplets, we find the textual information provides important clues, as illustrated in Table 4. Hence, we expect that advanced graph models that effectively incorporate textual information will be key to achieve good performance on the challenging time split.

Baseline. We consider two representative KG embedding models: TransE (Bordes et al., 2013) and ComplEx (Trouillon et al., 2016). These models define their own decoders to score knowledge triplets using the corresponding entity and relation embeddings. For instance, TransE uses $-\|\bm{h}+\bm{r}-\bm{t}\|_{2}$ as the decoder, where $\bm{h}$, $\bm{r}$, and $\bm{t}$ are embeddings of head, relation, and tail, respectively. For the encoder function (mapping each entity and relation to its embedding), we consider the following three options. Shallow: We use the distinct embedding for each entity and relation, as normally done in KG embedding models. RoBERTa: We use two MLP encoders (one for entity and another for relation) that transform the RoBERTa features into entity and relation embeddings. Concat: To enhance the expressive power of the previous encoder, we concatenate the shallow learnable embeddings into the RoBERTa features, and use the MLPs to transform the concatenated vectors to get the final embeddings. This way, the MLP encoders can adaptively utilize the RoBERTa features and the shallow embeddings to fit the large amount of triplet data. To implement our baselines, we utilize DGL-KE (Zheng et al., 2020).

Hyper-parameters. For the loss function, we use the negative sampling loss from Sun et al. (2019), where we pick margin $\gamma$ from {1,4,8,10,100}. In order to balance the performance and the memory cost, we use the embedding dimensionality of 200 for all the models.

Discussion. Table 5 shows the validation and test performance of the six different models, i.e., combination of two decoders (TransE and ComplEx) and three encoders (Shallow, RoBERTa, and Concat). Notably, in terms of the encoders, we see that the most expressive Concat outperforms both Shallow and RoBERTa, indicating that both the textual information (captured by the RoBERTa embeddings) and structural information (captured by node-wise learnable embeddings) are useful in predicting validation and test triplets. In terms of the decoders, TransE and ComplEx show similar performance with the Concat encoder, while they show somewhat mixed results with the Shallow and RoBERTa encoders.

Overall, our experiments suggest that the expressive encoder that combines both textual information and structural information gives the most promising performance. In the KG completion literature, the design of the encoder has been much less studied compared to the decoder designs. Therefore, we believe there is a huge opportunity in scaling up more advanced encoders, especially GNNs (Schlichtkrull et al., 2018), to further improve the performance on this dataset.

KDD Cup 2021 summary. Table 5 shows the results of the top 3 winners of the KDD Cup: BD-PGL Team (Su et al., 2021), OhMyGod Team (Peng et al., 2021), and the GraphMIRAcles Team (Cai et al., 2021). All the winning solutions outperform our strong baselines significantly, achieving near-perfect test MRR score of 0.97. For a more direct comparison, with a single model (no model ensembling), the BD-PGL Team reports a validation MRR of 0.92 (Su et al., 2021), improving our best ComplEx-Concat baseline by 0.07 points in validation MRR. Similar to our baselines, all the winners utilize the KG embedding approach as the backbone, and adopt the encoder that takes both shallow embedding and textual embeddings into account. Specifically, BD-PGL proposed the NOTE model (Su et al., 2021) which makes the RotatE model more expressive, while OhMyGod adopted the ensemble of several existing KG embedding models. On the other hand, GraphMIRAcles explored different design choices for the encoder and found that adding residual connection for shallow embeddings significantly improved the model performance.

In addition to the model advances, all the winners exploited some statistical property of candidate tail entities. Most notably, Yang et al. (2021) found that simply by sorting the candidate tails by the frequency they appear in the training KG, it was possible to achieve validation MRR of 0.75, rivaling our TransE-Shallow baseline. This highlights that our negative tail candidates are mostly rare entities that can be easily distinguished from the true tail entity. On the other hand, the practical KG completion presents a much harder challenge: the candidate tails are not provided, and a model needs to predict the true tail entity out of all the possible 87M entities. As the performance on WikiKG90M has already saturated, we have updated WikiKG90M to WikiKG90Mv2 to reflect the realistic setting in large-scale KG completion. See Section 3 for further details.

Practical relevance and dataset overview. Density Functional Theory (DFT) is a powerful and widely-used quantum physics calculation that can accurately predict various molecular properties such as the shape of molecules, reactivity, responses by electromagnetic fields (Burke, 2012). However, DFT is time-consuming and takes up to several hours per small molecule. Using fast and accurate ML models to approximate DFT enables diverse downstream applications, such as property prediction for organic photovaltaic devices (Cao and Xue, 2014) and structure-based virtual screening for drug discovery (Ferreira et al., 2015).

PCQM4M is a quantum chemistry dataset originally curated under the PubChemQC project (Nakata, 2015; Nakata and Shimazaki, 2017). Based on the PubChemQC, we define a meaningful ML task of predicting DFT-calculated HOMO-LUMO energy gap of molecules given their 2D molecular graphs. The HOMO-LUMO gap is one of the most practically-relevant quantum chemical properties of molecules since it is related to reactivity, photoexcitation, and charge transport (Griffith and Orgel, 1957). Moreover, predicting the quantum chemical property only from 2D molecular graphs without their 3D equilibrium structures is also practically favorable. This is because obtaining 3D equilibrium structures requires DFT-based geometry optimization, which is expensive on its own.

To ensure the resulting models are practically useful, we limit the average inference budget per molecule (including both pre-processing and model inference) to be less than 0.1 second using a single GPU and CPU (multi-threading on a multi-core CPU is allowed). This means that expensive (quantum) calculations cannot be used to perform inference. As our test set contains 377,423 molecules, we require the all the prediction to be made within 12 hours. Note that this time constraint is quite generous for ordinary GNNs—each of our baseline GNN only took about 3 minutes to perform inference over the entire test data.

Graph. We provide molecules as the SMILES strings (Weininger, 1988), from which 2D molecule graphs (nodes are atoms and edges are chemical bonds) as well as molecular fingerprints (hand-engineered molecular feature developed by the chemistry community) can be obtained. By default, we follow OGB (Hu et al., 2020a) to convert the SMILES string into a molecular graph representation, where each node is associated with a 9-dimensional feature (e.g., atomic number, chirality) and each edge comes with a 3-dimensional feature (e.g., bond type, bond stereochemistry), although the optimal set of input graph features remains to be explored.

Prediction task and evaluation metric. The task is graph regression: predicting the HOMO-LUMO energy gap in electronvolt (eV) given 2D molecular graphs. Mean Absolute Error (MAE) is used as evaluation metric.

Dataset split. We split molecules by their PubChem ID (CID) with ratio 80/10/10. Our original intention was to provide the scaffold split (Hu et al., 2020a; Wu et al., 2018), but the provided data turns out to be split by the CID due to some pre-processing bug. The CID number itself does not indicate particular meaning about the molecule, but splitting by CID may provide a moderate distribution shift (most likely not as severe as the scaffold split). We empirically compared the CID and scaffold splits and found the model performances were consistent between the two splits.⁵⁵5Detailed discussion can be found at https://github.com/snap-stanford/ogb/discussions/162

Baseline. We benchmark two types of models: a simple MLP over the Morgan fingerprint (Morgan, 1965) and more advanced GNN models. For GNNs, we use the four strong models developed for graph-level prediction: Graph Convolutional Network (GCN) (Kipf and Welling, 2017) and Graph Isomorphism Network (GIN) (Xu et al., 2019), as well as their variants, GCN-virtual and GIN-virtual, which augment graphs with a virtual node that is bidirectionally connected to all nodes in the original graph (Gilmer et al., 2017). Adding the virtual node is shown to be effective across a wide range of graph-level prediction datasets in OGB (Hu et al., 2020a). Edge features are incorporated following Hu et al. (2020b). At inference time, the model output is clamped between 0 and 50 to avoid model’s anomalously large/small prediction.

Hyper-parameters. For the MLP over Morgan fingerprint, we set the fingerprint dimensionality to be 2048, and tune the fingerprint radius $\in\{2,3\}$, as well as MLP’s hyper-parameters: hidden dimensionality $\in\{1200,1600\}$, number of hidden layers $\in\{2,4,6\}$, and dropout ratio $\in\{0,0.2\}$. For GNNs, we tune hidden dimensionality, i.e., width $\in\{300,600\}$, number of GNN layers, i.e., depth $\in\{3,5\}$. Simple summation is used for graph-level pooling. For all MLPs (including GIN’s), we use batch normalization (Ioffe and Szegedy, 2015) and ReLU activation.

Discussion. The validation and test results are shown in Table 6. We see both the GNN models significantly outperform the simple fingerprint baseline. Expressive GNNs (GIN and GIN-virtual) outperform less expressive ones (GCN and GCN-virtual); especially, the most advanced and expressive GIN-virtual model significantly outperforms the other GNNs. Nonetheless, the current performance is still much worse than the chemical accuracy of 0.043eV—an indicator of practical usefulness established by the chemistry community. In the same Table 6, we show our ablation, where we use only 10% of data to train the GIN-virtual model. We see the performance significantly deteriorate, indicating the importance of training the model on large data. Finally, in Table 7, we show the relation between model sizes and validation performance. We see that the largest models always achieve the best performance.

Overall, we find that advanced, expressive, and large GNN model gives the most promising performance on the PCQM4M dataset, although the performance still needs to be improved for practical use. We believe further advances in advanced modeling, expressive architectures, and larger model sizes could yield breakthrough in the large-scale molecular property prediction task.

KDD Cup 2021 summary. In Table 6, we show the results of the top 3 winners of the KDD Cup: Machine Learning Team (Ying et al., 2021b), SuperHelix Team (Zhang et al., 2021), and Quantum Team (Addanki et al., 2021). The winners have significantly reduced the MAE compared our baselines, yielding around 0.03 points improvement in test MAE. For a more direct comparison, with a single model, the Machine Learning reports the validation MAE of 0.097 for their Graphormer model (Ying et al., 2021a), which is 0.04 points lower than our best GIN-virtual baseline.

In terms of methodology, we find that the winning solutions share three important components in common. (1) Their winning GNN models are indeed large and deep; the number of learnable parameters (single model) ranges from 50M up to 450M, while the number of GNN layers ranges from 11 up to 50, being significantly larger than our baseline models. (2) All the GNNs perform global message passing at each layer, either through the virtual nodes (Gilmer et al., 2017) or fully-connected Transformer-style self-attention (Ying et al., 2021a). (3) All the winners utilize 3D structure of molecules to supervise their GNNs. As 3D structure was not provided at our KDD Cup, the winners generate the 3D structure themselves using RDkit (Landrum et al., 2006) or PySCF (Sun et al., 2020), both of which provide cheap but less accurate 3D structure of molecules.

As modeling 3D molecular graphs is a promising direction in graph ML (Schütt et al., 2017; Klicpera et al., 2020; Sanchez-Gonzalez et al., 2020; Hu et al., 2021), we have updated PCQM4M to PCQM4Mv2 to include DFT-calculated 3D structures for training molecules. Details are provided in Section 3.