Pre-Training Graph Neural Networks: Strategies for Scientific Discovery

# Pre-Training Graph Neural Networks: Strategies for Scientific Discovery **Video Category:** Machine Learning / Graph Neural Networks ## 📋 0. Video Metadata **Video Title:** Stanford CS224W: Pre-Training Graph Neural Networks **YouTube Channel:** Stanford Engineering **Publication Date:** 3/18/21 (Extracted from presentation slides) **Video Duration:** ~20 minutes ## 📝 1. Core Summary (TL;DR) Graph Machine Learning holds massive potential for scientific domains like chemistry and biology, but faces severe bottlenecks due to label scarcity and out-of-distribution (OOD) test examples. Standard deep learning models overfit small datasets and rely on spurious correlations, causing them to fail when extrapolating to novel scientific discoveries. By systematically pre-training Graph Neural Networks (GNNs) on abundant relevant data using a combination of node-level self-supervised methods and graph-level supervised methods, we can successfully inject robust domain knowledge into the model. This dual-level pre-training strategy significantly improves downstream task performance, prevents the negative transfer commonly seen in naive pre-training approaches, and enables effective out-of-distribution generalization. ## 2. Core Concepts & Frameworks * **Graph Neural Networks (GNNs) for Classification:** -> **Meaning:** Neural architectures designed to process graph structures by iteratively aggregating neighboring information to obtain node embeddings, and subsequently pooling those node embeddings to obtain a global graph embedding. -> **Application:** Used to predict global properties, such as whether a specific molecular graph is toxic or whether a protein association sub-graph exhibits biological activity. * **Label Scarcity & Overfitting:** -> **Meaning:** A fundamental challenge in scientific machine learning where acquiring labels requires expensive wet-lab experiments, resulting in training datasets that are vastly smaller than the millions of parameters in a deep learning model. -> **Application:** Because the number of training data points is much less than the parameters, the model easily memorizes the small training set but fails to generalize. * **Out-of-Distribution (OOD) Prediction:** -> **Meaning:** A scenario where the test data fundamentally differs from the training data distribution. In scientific discovery, the goal is often to find novel entities (e.g., new drug molecules) that do not resemble known training examples. -> **Application:** Models trained purely on existing data often learn "spurious correlations" (easy shortcuts) rather than causal mechanisms, causing them to extrapolate poorly when faced with OOD data. * **Negative Transfer:** -> **Meaning:** A phenomenon where initializing a model with pre-trained parameters actually degrades performance on a downstream task compared to starting with a randomly initialized, non-pre-trained baseline. -> **Application:** Occurs when naive pre-training strategies (like graph-level multi-task training only) fail to generate high-quality, robust local node embeddings. * **Distributional Hypothesis in Graphs:** -> **Meaning:** Adapted from Natural Language Processing (e.g., word2vec), this hypothesis assumes that sub-graphs surrounded by similar context graphs are semantically similar. -> **Application:** Used as the foundation for the "Context Prediction" self-supervised pre-training algorithm to learn graph semantics without external labels. ## 3. Evidence & Examples (Hyper-Specific Details) * **Scientific Application Examples:** The video highlights two primary graph ML use cases: - **Chemistry (Molecular Graphs):** Nodes represent atoms, edges represent chemical bonds. The function $f(graph)$ predicts properties like toxicity. - **Biology (Protein-Protein Association):** Nodes represent proteins, edges represent associations. Given a sub-graph centered around a protein node, the function $f(graph)$ predicts biological activity. * **Spurious Correlations (Image Classification Analogy):** To illustrate OOD failure, the speaker uses an image classification example: distinguishing polar bears from brown bears. During training, most polar bears have a snow background, and brown bears have a grass background. A deep learning model will likely use the spurious correlation of the background (snow vs. grass) to make predictions rather than the animal itself. At test time, if presented with a polar bear on grass (OOD), the model fails completely. * **Naive Pre-Training Strategy & Negative Transfer Evidence:** - **Strategy:** Multi-task supervised pre-training on relevant graph-level labels (e.g., predicting Toxicity A, Toxicity B, Bioactivity A simultaneously). - **Experimental Setup:** Pre-trained on 1310 diverse binary bioassays annotated over ~450K molecules. Tested on 8 downstream molecular classification datasets (1K to 100K molecules) using ROC-AUC as the metric. - **Data Split:** A "Scaffold split" was used to ensure test molecules are structurally different (out-of-distribution) from training molecules. - **Result:** A scatter plot demonstrated that this naive strategy yielded limited improvements and frequently resulted in **negative transfer**. On several datasets, the pre-trained model performed worse than the randomly initialized baseline. * **Attribute Masking Algorithm (Node-Level Pre-Training):** A visual demonstration showed a molecule where a specific oxygen node ('O') was replaced with a 'Masked node' ('X'). The GNN processes the graph to generate node embeddings, and the embedding of node 'X' is used to predict its true identity from a list of possibilities [C, N, O, S...]. This forces the GNN to learn basic chemistry rules (valency, common structures) to solve the task. * **Context Prediction Algorithm (Node-Level Pre-Training):** The algorithm samples a center node and extracts its K-hop neighborhood (e.g., $r_1$) and the surrounding context graph (e.g., $r_2$). Two separate GNNs encode the neighborhood and context into vectors. The model is trained via contrastive learning to maximize the inner product between true (neighborhood, context) pairs and minimize the inner product between the neighborhood and randomly sampled "false" contexts from other graphs. * **Proposed Dual-Level Pre-Training Results:** By combining node-level self-supervised pre-training (Attribute Masking/Context Prediction) with graph-level supervised pre-training, the resulting model consistently avoided negative transfer and significantly improved performance across all 8 downstream molecular classification datasets compared to the baseline. * **Comparison of GNN Architectures (GIN vs. GAT):** A critical experiment tested how different GNN architectures benefit from the proposed pre-training strategy: - **GIN (Graph Isomorphism Network - Most expressive):** Saw massive gains of +7.2 ROC-AUC on Chemistry tasks and +9.4 on Biology tasks. - **GAT (Graph Attention Network - Less expressive):** Experienced negative transfer, losing -6.5 ROC-AUC on Chemistry and -0.4 on Biology after pre-training. - **GCN (Graph Convolutional Network):** Saw moderate gains (+3.4 Chem, +7.7 Bio). - **GraphSAGE:** Saw minor gains (+2.0 Chem, +2.8 Bio). ## 4. Actionable Takeaways (Implementation Rules) * **Rule 1: Evaluate using Scaffold Splits for scientific data** - When building models for discovery, do not use random data splits. Use scaffold splits to ensure the test set contains novel structures (out-of-distribution), accurately simulating real-world scientific discovery scenarios. * **Rule 2: Do not rely exclusively on graph-level pre-training** - Simply pre-training a GNN on a multi-task dataset of global graph labels is insufficient. It fails to guarantee high-quality local node embeddings prior to the pooling step, often resulting in negative transfer on downstream tasks. * **Rule 3: Implement Node-Level Attribute Masking to learn fundamental rules** - Mask random node attributes (e.g., atom types) in your unlabeled graph data and train the GNN to predict them based on their local neighborhood. This acts as a self-supervised quiz forcing the network to internalize domain-specific structural rules. * **Rule 4: Implement Node-Level Context Prediction to learn sub-graph semantics** - Isolate K-hop neighborhoods and their surrounding context graphs. Train two GNN encoders to maximize the similarity of valid neighborhood/context pairs and minimize similarity for randomized pairs to capture broader semantic meaning. * **Rule 5: Execute a 3-Step Pre-Training Pipeline** - For maximum performance: (1) Perform self-supervised node-level pre-training. (2) Use those initialized parameters to perform supervised graph-level pre-training. (3) Fine-tune the resulting model on your specific, small-data downstream task. * **Rule 6: Use highly expressive GNN architectures for pre-training** - Select expressive models like GIN (Graph Isomorphism Network) over models like GAT or GraphSAGE. Expressive models have the mathematical capacity to capture, retain, and transfer the complex domain knowledge injected during pre-training. ## 5. Pitfalls & Limitations (Anti-Patterns) * **Pitfall:** Training deep models on scarce scientific data. -> **Why it fails:** Deep learning models possess millions of parameters. When trained on datasets requiring expensive lab experiments (e.g., a few thousand examples), the models simply memorize the data (overfitting). -> **Warning sign:** High performance on validation sets drawn from the exact same distribution as the training data, but catastrophic failure when deployed in the wild. * **Pitfall:** Assuming models learn causal mechanisms. -> **Why it fails:** Deep networks default to learning the easiest predictive shortcut, which is often a spurious correlation (e.g., classifying a bear by the background color instead of its fur). -> **Warning sign:** The model's predictive accuracy drops to near-random guessing when tested on out-of-distribution examples where the background/shortcut no longer applies. * **Pitfall:** Naive multi-task supervised pre-training on graphs. -> **Why it fails:** While the model learns to predict global graph properties, there is no structural constraint forcing the model to generate robust, meaningful *node-level* embeddings prior to the pooling layer. -> **Warning sign:** Experiencing "negative transfer," where a model initialized with pre-trained weights performs significantly worse on a new task than a randomly initialized model. * **Pitfall:** Using low-expressivity models (like GAT) for pre-training. -> **Why it fails:** The architecture lacks the representational capacity to store the complex domain knowledge generated across diverse pre-training tasks. -> **Warning sign:** The pre-training phase yields negligible benefits or actual performance degradation (e.g., the -6.5 ROC-AUC drop observed with GAT) compared to non-pre-trained baselines. ## 6. Key Quote / Core Insight "To unlock the full potential of pre-training, you must utilize highly expressive neural architectures; a model lacking structural capacity cannot retain the complex domain knowledge you are attempting to inject." ## 7. Additional Resources & References * **Resource:** Sagawa et al. ICML 2020 - **Type:** Paper - **Relevance:** Explains how deep learning models often make predictions based on spurious correlations in datasets. * **Resource:** Hendrycks et al. ICML 2019 - **Type:** Paper - **Relevance:** Provides foundational evidence that pre-training improves out-of-distribution performance in vision and NLP. * **Resource:** Mikolov et al. NIPS 2013 - **Type:** Paper - **Relevance:** Introduces word2vec and the distributional hypothesis, which inspired the context prediction pre-training strategy for graphs. * **Resource:** Hu et al. ICLR 2020 - **Type:** Paper - **Relevance:** The primary research paper detailing the proposed dual-level GNN pre-training strategies discussed in the video.