GML26

Graph Machine Learning Challenge Repository

Protein Function Prediction from Graph Structure

A config‑driven pipeline for classifying protein graphs as enzyme or non‑enzyme using the PROTEINS benchmark from TUDataset.
The pipeline integrates techniques from graph representation learning, self‑supervised pre‑training, optimisation, calibration and evaluation.

Val. Accuracy: 79.6 %   |   Val. F1: 74.5 %   |   Test ROC‑AUC: 79.2 %   |   Test PR‑AUC: 69.4 %

---

Table of Contents

1. Overview
2. Key Features
3. Repository Structure
4. Quick Start
5. Usage & Run Modes
6. Pipeline Walkthrough
7. Configuration Reference
8. Reproducibility
9. Results
10. Dependencies
11. Acknowledgments & References
12. License

---

Overview

The code implements a full graph‑level classification workflow:

1. Data loading – PROTEINS_full from TUDataset (automatic fallback to synthetic data if offline).
2. Feature engineering – topological features (degree, clustering coefficient) and positional encodings (Laplacian eigenvectors, random‑walk structural encodings).
3. Self‑supervised pre‑training (SSL) – structure‑only training on the train split:
   • Adversarial VGAE‑CL (default) – variational autoencoder with adversarial prior regularisation and a contrastive objective between augmented views. Also generates synthetic graphs to augment the training set.
   • VGAE+GraphMAE – adjacency reconstruction + masked node feature reconstruction.
   • GraphCL‑style contrastive – SimCLR on two graph views.
4. Hybrid GIN + Graph‑Transformer/GATv2 classifier – Jumping Knowledge, attention pooling, residual connections, GraphNorm/LayerNorm/PairNorm.
5. Training – class‑weighted focal loss, warmup + cosine annealing with restarts, curriculum‑scheduled augmentation and dropout, optional SAM (Sharpness‑Aware Minimization) and Lookahead, AMP.
6. Post‑training refinement – Stochastic Weight Averaging (SWA), validation‑ROC‑tuned decision threshold, temperature scaling calibration.
7. Inference – test‑time augmentation (TTA).
8. Extended evaluation modes – k‑fold CV, multi‑seed ensemble, Optuna hyper‑parameter search, and ablation sweep grid.

All techniques are selected via a single Config dataclass; the pipeline is fully reproducible (deterministic CUDA, fixed seeds).

---

Key Features

• Self‑supervised pre‑training on raw structure – no label leakage.
• Synthetic graph generation from the adversarial VGAE latent space.
• Curriculum augmentation – linearly ramped feature masking, edge dropping, and DropNode.
• Degree‑scaled adaptive edge dropping – stronger perturbation on denser graphs.
• Hybrid GNN architecture – injective GIN + expressive attention (TransformerConv or GATv2) in parallel or sequential mode.
• Optimisation – Lion optimizer, SAM, Lookahead, gradient clipping, AMP.
• Calibration & threshold tuning – temperature scaling, cost‑weighted Youden’s J.
• Multiple run modes – single run, 5‑fold CV, ensemble, Optuna HPO, ablation grid sweep.
• Rich logging & visualisation – ROC curves, t‑SNE embeddings, CSV output for sweeps.

---

Repository Structure

protein_gin_classifier.py   # Main script – single entry point
outputs/                    # Saved model, diagnostic plots
data/TUDataset/             # Local copy of PROTEINS (downloaded automatically)
README.md

All logic is contained in protein_gin_classifier.py for simplicity – no separate package hierarchy.

---

Quick Start

1. Clone the repository

git clone https://github.com/Penzo00/GML26.git
cd GML26

2. Install dependencies

pip install torch torch-geometric scikit-learn matplotlib networkx

Optional for HPO:

pip install optuna

3. Run a full training + evaluation

python protein_gin_classifier.py

The script will:

• Download PROTEINS_full (or create a synthetic fallback),
• Pre‑train an adversarial VGAE‑CL SSL model,
• Generate synthetic graphs,
• Train the GINTransformer classifier,
• Perform SWA, threshold tuning, calibration, TTA,
• Save model weights and diagnostic plots to outputs/.

---

Usage & Run Modes

The script accepts a --mode argument:

Mode	Command	Description
single (default)	python protein_gin_classifier.py	Standard train/val/test run
kfold	python protein_gin_classifier.py --mode kfold --k 5	Stratified 5‑fold cross‑validation
ensemble	python protein_gin_classifier.py --mode ensemble --runs 5	Multi‑seed averaged predictions
hpo	python protein_gin_classifier.py --mode hpo --trials 20	Optuna hyper‑parameter search
sweep	python protein_gin_classifier.py --mode sweep --sweep-epochs 60	Ablation grid over optimiser/SSL method, writes outputs/ablation_results.csv

All modes reuse the same core Config and train_and_evaluate() function, guaranteeing consistent evaluation.

---

Pipeline Walkthrough

Data Loading & Feature Engineering

• PROTEINS_full: 1113 graphs, 32 node features, binary classification (enzyme / non‑enzyme).
• Topological features: normalised degree + clustering coefficient (2 dims).
• Positional Encodings (8+8 dims): Laplacian eigenvectors and/or random‑walk self‑return probabilities.
• All graph features are reset before each SSL pre‑training to avoid leakage.

SSL Pre‑training (train graphs only)

• Adversarial VGAE‑CL (default ssl_method='adv_vgae_cl'):
  • Encoder: GATv2Conv + residual connection.
  • Trains with reconstruction loss, KL divergence, feature decoder loss, contrastive NT‑Xent loss, and adversarial loss to push latents towards a standard Gaussian.
  • After training, generates synthetic_augmentation_ratio * len(train_set) synthetic graphs to augment the training split.
• VGAE + GraphMAE ('vgae_graphmae'): GCN encoder, adjacency + masked feature reconstruction.
• Contrastive ('contrastive'): GIN encoder + projection head, NT‑Xent on two augmented views.
• The resulting frozen per‑node embeddings (16‑dim) are concatenated to node features for all splits.

Classifier Architecture (GINTransformer)

• Input projection + normalisation.
• num_layers blocks of:
  • ManualGINLayer (injective update with trainable ε, internal GraphNorm/LayerNorm/PairNorm).
  • Attention convolution (TransformerConv or GATv2) – applied either in parallel (summed with GIN output) or sequentially after GIN.
• Jumping Knowledge: concatenation of outputs from every layer, projected to hidden dim.
• Attention pooling (AttentionalAggregation) + sum pooling → concatenated graph embedding.
• Final classification head.

Training Recipe

• Focal loss (γ=2) with inverse frequency class weights and label smoothing (0.005).
• Optimiser: Lion (default) or AdamW, wrapped in SAM (ρ=0.03) for flat‑minima generalisation.
• Learning rate schedule: linear warmup (15 epochs) → cosine annealing with warm restarts (T₀=80, T_mult=2, min_lr=1e‑5).
• Curriculum augmentation: feature mask probability, edge drop probability, node drop probability linearly ramped from 0 to final values over 100 epochs.
• Curriculum dropout: linearly ramped from 0 to 0.5 over 100 epochs.
• Gradient clipping (max norm 5.0).
• Early stopping with patience=30 after the ramp completes, monitored on validation F1.

Post‑training

• SWA: 30 additional epochs with reduced LR, averaged weights – used only if it improves validation F1.
• Threshold tuning: cost‑weighted Youden’s J on validation ROC (cost_fn_over_fp=1.0 treats FN/FP equally).
• Temperature scaling: fits a single scalar on validation logits with LBFGS to improve probability calibration.
• Test‑time augmentation: averages softmax probabilities over 5 lightly augmented views + 1 clean pass.

Evaluation Metrics

• Accuracy, precision, recall, F1 (binary), MCC.
• PR‑AUC, ROC‑AUC.
• Brier score before/after calibration.

---

Configuration Reference

All hyper‑parameters are declared in the Config dataclass (line ~85). The most important ones:

Parameter	Default	Description
seed	42	Random seed for full reproducibility
batch_size	512	Graphs per batch
ssl_method	'adv_vgae_cl'	SSL strategy (vgae_graphmae, contrastive, none)
synthetic_augmentation_ratio	0.28	Fraction of train graphs to add as synthetic
hidden_dim	128	Hidden dimension of GNN
num_layers	2	Number of GIN+attention blocks
conv_type	'transformer'	Attention convolution (transformer or gatv2)
hybrid_mode	'parallel'	How GIN and attention are combined (parallel or sequential)
norm_type	'graph'	Normalisation (graph, layer, pairnorm)
optimizer_type	'lion'	Optimiser family (lion or adamw)
use_sam	True	Enable Sharpness‑Aware Minimization
lion_lr	1e‑4	Lion learning rate
sam_rho	0.03	SAM perturbation radius
weight_decay	0.1	Weight decay (decoupled for Lion)
warmup_epochs	20	Linear warmup duration
n_epochs	200	Maximum training epochs
patience	30	Early stopping patience (after ramp)
swa_epochs	30	Number of SWA refinement steps
feat_mask_prob_final	0.2	Max. feature mask probability (augmentation)
edge_drop_prob_final	0.2	Max. edge drop probability
node_drop_prob_final	0.1	Max. node drop probability
cost_fn_over_fp	1.0	Relative cost of false negatives for threshold tuning
use_tta	True	Enable test‑time augmentation

Modify these directly in the script or override programmatically when calling train_and_evaluate().

---

Reproducibility

• Deterministic CUDA operations (torch.backends.cudnn.deterministic = True, benchmark = False) – exact same results on the same hardware.
• All random seeds set for Python, NumPy, PyTorch.
• Data splits are stratified and seeded (random_state=cfg.seed).
• Full end‑to‑end pipeline can be re‑run with python protein_gin_classifier.py.

---

Results

The logged output shows:

SSL epoch 20 loss: 12.1504
Synthetic graphs generated: 224
RandomForest baseline (on SSL embeddings): acc=0.7054, f1=0.6374
Final node feature dim: 32 (raw) + 18 (topo+PE) + 16 (SSL) = 66
Class weights: [0.8403, 1.2346]

Best validation at epoch 130:
  val_acc=0.7960, val_f1=0.7453
SWA refinement: val_f1 0.7485 → SWA weights selected.

Test (0.5 threshold):
  accuracy=0.7232, precision=0.6667, recall=0.6222, f1=0.6437, MCC=0.4185
PR‑AUC=0.6938, ROC‑AUC=0.7924
Test TTA (5 views + clean): accuracy=0.7054
Temperature scaling: T=0.742, Brier 0.2158 → 0.2093

Note: The test set is relatively small (112 graphs) and has a class imbalance (~58% non‑enzyme). The metrics are consistent with the original literature on PROTEINS; the pipeline prioritises robust generalisation over raw test‑set maximisation by using SAM, SWA, and careful calibration.

---

Visualizations

The pipeline saves two diagnostic plots after each successful single run.

t‑SNE of Learned Graph Embeddings
How the model organises the 1113 protein graphs in its embedding space.

Test ROC Curve
Performance of the classifier with the validation‑tuned threshold marked in red.

---

Dependencies

• Python 3.8+
• PyTorch ≥ 1.12
• PyTorch Geometric (and its dependencies)
• scikit‑learn
• matplotlib
• networkx
• (Optional) Optuna for HPO mode

Install with:

pip install torch torch-geometric scikit-learn matplotlib networkx optuna

---

Acknowledgments & References

The pipeline synthesises concepts from:

• Graph Representation Learning – William L. Hamilton (Ch. 5–7)
• GIN / Graph Isomorphism Networks – Xu et al., ICLR 2019
• VGAE – Kipf & Welling, NeurIPS 2016
• GraphMAE – Hou et al., KDD 2022
• GraphCL – You et al., NeurIPS 2020
• Lion Optimiser – Chen et al., 2023
• Sharpness‑Aware Minimization – Foret et al., ICLR 2021
• SWA – Izmailov et al., UAI 2018
• Temperature Scaling – Guo et al., ICML 2017
• PairNorm – Zhao & Akoglu, NeurIPS 2020
• Lookahead – Zhang et al., NeurIPS 2019

---

AI Usage

This README was generated with the assistance of DeepSeek, a large language model, based on the codebase and training logs provided. The main pipeline code was developed with the help of Claude (Anthropic), another large language model, which contributed to the implementation of the architecture, training procedures, and various techniques.