Skip to content

Repository files navigation

Bonus task report — End-to-end temporal link prediction

This folder contains my end‑to‑end temporal link prediction model for the Science4Cast / FutureOfAIviaAI graph. I focus on learning directly from the graph structure and time, without building large sets of hand‑crafted graph statistics (centralities, node/edge scores, etc.).

I reused the minimal runner (evaluate_model.py, simple_model.py, utils.py) from the original repository because it is lightweight and convenient, but the actual GNN + temporal model and training logic I wrote from scratch for this project.

The code you run is evaluate_model.pysimple_model.pyutils.py.


1) Competition context

The original FutureOfAIviaAI / Science4Cast setup is a link prediction competition on a temporal citation-style graph:

  • The dataset is a dynamic graph: edges have timestamps.
  • For a chosen configuration (I use delta=1, cutoff=5, minedge=1) the task is to predict which currently unconnected pairs will become edges one year later.
  • The official metric is AUC-ROC on a fixed set of candidate pairs provided in the .pkl files.

In the official leaderboard, the strongest methods (M1, M2, …) rely on fairly heavy handcrafted feature engineering: dozens of centrality measures, similarity scores, and so on, which are then fed into classical ML or shallow neural nets. There are also deep models that still sit on top of precomputed features.

I decided to go in the opposite direction:

  • I am lazy and I believe in deep learning, so I wanted a model that learns as much as possible directly from the graph structure and time, without constructing tens of graph statistics by hand.
  • In particular, I did not use the M6 “15 hand‑crafted features” or the larger feature sets from M1/M2; instead I let a GNN + Selective SSM learn node representations and edge scores end‑to‑end.

2) Data pipeline

I focus on a single dataset due to time and compute constraints:

  • datasets/SemanticGraph_delta_1_cutoff_5_minedge_1.pkl
  • years_delta = 1, vertex_degree_cutoff = 5, min_edges = 1

Inside my code there are two different “pair sets”:

  1. Training pairs (my sampled subset)
    Created by utils.create_training_data_biased(...):
  • Pick “eligible” nodes with degree ≥ vertex_degree_cutoff.
  • Positives: edges that appear by year_start+delta but didn’t exist at year_start (and pass min_edges).
  • Negatives: random eligible pairs that are not edges at year_start and do not become positive by year_start+delta.
  • Target is 50/50: edges_used//2 positives + the rest negatives (falls back if not enough positives exist).

In simple_model.link_prediction_semnet(...) I:

  • Shuffle these pairs.
  • Split them into 90% train / 10% validation (this is the val_auc you see in my logs).
  1. Official evaluation pairs (provided by the dataset)
    Loaded from the .pkl file (unconnected_vertex_pairs + ground truth labels). This is what evaluate_model.py evaluates at the end.

3) Model: architecture and why I chose it

High‑level idea

My goal was an end‑to‑end model that, given the temporal graph, learns:

  • Node embeddings that encode a k‑hop neighborhood at each time slice.
  • A way to combine those embeddings into edge scores for candidate pairs.
  • A way to aggregate information from several previous years into a single prediction.

The overall architecture is sketched here:

I implemented this in simple_model.py as:

  1. Per‑year encoder (GIN)
    For each year, I run a small Graph Isomorphism Network that produces an embedding for every node. Each embedding h_t(u) already sees the k‑hop neighborhood of node u in that year through message passing.

  2. Edge representation via concatenation of two node embeddings
    For an undirected candidate pair (u,v) at a fixed time slice I take the two node embeddings hu = h_t(u) and hv = h_t(v) and concatenate them: [hu‖hv].
    This concatenation is not symmetric by itself: [hu‖hv] and [hv‖hu] are different even though the underlying edge has no direction. Initially I tried to fix symmetry directly in these features (see the next section); in the final version I simplified the architecture and enforce symmetry at the head level by averaging logits for (u,v) and (v,u).

  3. Hop aggregation (attention over hops)
    I keep all hop outputs from the GIN (0..k). For each hop I build an edge representation, and then I use attention over hops to get one vector per (u,v, t) per year. This replaces the hand‑designed “hop features” from some other solutions with a learned hop‑selector.

  4. Temporal aggregation (Selective SSM)
    Across years, instead of a transformer I use a Selective State Space Model to aggregate the sequence of per‑year edge representations. It is lighter than a transformer and handles longer histories better than a plain RNN.

  5. Binary edge head
    On top I put a small MLP head that outputs a single logit per candidate pair: “will uv become an edge next year or not”.
    I chose a pure binary head because it keeps the task simple (one‑year‑ahead yes/no prediction). If I ever need multi‑year forecasts, I can run the same model autoregressively: predict edges for T+1, treat them as part of the graph, then predict T+2, and so on.


4) How I handled symmetry

At the start, I used a naive asymmetric head:

  • I concatenated [hu‖hv] and fed it into the MLP head.
  • This worked reasonably well, but the representation itself was still orientation‑dependent.

Then I tried to symmetrize the edge representation by hand:

  • Instead of plain concatenation I added combinations like hu+hv, hu⊙hv, |hu−hv|.
  • This helped with symmetry, but in practice it turned into another small set of manually designed features at the head level and started to overfit very quickly on my validation split.

In the final version I dropped these handcrafted symmetric combinations and made the head symmetric at the logit / loss level:

logit_uv = head([huhv])
logit_vu = head([hvhu])
logit = 0.5 * (logit_uv + logit_vu)
loss = BCEWithLogits(logit, y)

This way I keep the head simple (no engineered features), but force it to be invariant to swapping (u,v). In practice this stopped the early overfitting I saw with plain concatenation and gave slightly higher and more stable validation AUC.


5) Training setup and regularization

For training I use a fairly standard setup:

  • Optimizer: AdamW with weight decay and gradient clipping (see train_model in simple_model.py).
  • Dropout in the GIN MLPs, hop attention and Selective SSM.

Compared to the original M6 code I:

  • Slightly lowered the learning rate (to make generalization a bit safer).
  • Increased dropout in the model blocks to reduce overfitting, given that I do not use strong hand‑crafted features.
  • Used edges_used in the range 50k–100k (instead of very large values) because on my hardware the model trains quite slowly; I wanted to finish several ablations in a reasonable time.

All experiments in this report are on one dataset: SemanticGraph_delta_1_cutoff_5_minedge_1.pkl.


6) Results and ablations (delta=1, cutoff=5, minedge=1)

Below are validation AUC-ROC curves from my internal 90/10 split on the sampled training pairs.
They are meant to compare different variants of my model to each other on the same dataset.

Experimental setup (for the plots below)

  • Dataset: datasets/SemanticGraph_delta_1_cutoff_5_minedge_1.pkl
  • Seed: 42
  • Time context: time_slices=6 (years)
  • Model defaults: d_model=128, gin_layers=3, mlp_hidden=512
  • Optimizer: AdamW + weight_decay=1e-2, grad clip 1.0
  • Training subset size: edges_used=50k or 100k (as written in the run title/log)

Plots

Below are three comparison plots (4 runs on each figure):

Link loss (train):

Validation loss:

Validation AUC-ROC:


7) Comparison with official methods (delta=1, cutoff=5, minedge=1)

Below is the leaderboard slice for delta=1, cutoff=5, minedge=1 from the original FutureOfAIviaAI README (AUC‑ROC on official evaluation pairs), plus my model.

Model AUC
M1 — Yichao’s NF+ML(+GNN) 0.9252
M2 — HashBrown (node feats + MLP) 0.9175
M4B — Common Neighbours 0.9016
M3 — LSTM over temporal feats 0.8980
My model — GIN + Selective SSM, symmetric loss, no handcrafted edge features 0.8947
M4A — Preferential Attachment 0.8862
M7A — Node2Vec + MLP 0.8558
M7B — ProNE + MLP 0.8538
M6 — Baseline (15 feats + MLP) 0.8526
M8 — Transformer (temporal embeddings) 0.8253

8) How to run

Setup:

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Run one dataset:

python evaluate_model.py datasets/SemanticGraph_delta_1_cutoff_5_minedge_1.pkl --no-plots

Example with my typical settings:

python evaluate_model.py datasets/SemanticGraph_delta_1_cutoff_5_minedge_1.pkl \
  --edges-used 100000 --batch-size 400 --lr 1e-4 --max-iters 10000

About

SSM-based GNN for citation graph link prediction

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages