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.py → simple_model.py → utils.py.
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
.pklfiles.
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.
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”:
- Training pairs (my sampled subset)
Created byutils.create_training_data_biased(...):
- Pick “eligible” nodes with degree ≥
vertex_degree_cutoff. - Positives: edges that appear by
year_start+deltabut didn’t exist atyear_start(and passmin_edges). - Negatives: random eligible pairs that are not edges at
year_startand do not become positive byyear_start+delta. - Target is 50/50:
edges_used//2positives + 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_aucyou see in my logs).
- Official evaluation pairs (provided by the dataset)
Loaded from the.pklfile (unconnected_vertex_pairs+ ground truth labels). This is whatevaluate_model.pyevaluates at the end.
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:
-
Per‑year encoder (GIN)
For each year, I run a small Graph Isomorphism Network that produces an embedding for every node. Each embeddingh_t(u)already sees the k‑hop neighborhood of nodeuin that year through message passing. -
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 embeddingshu = h_t(u)andhv = 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). -
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. -
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. -
Binary edge head
On top I put a small MLP head that outputs a single logit per candidate pair: “willu–vbecome 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 forT+1, treat them as part of the graph, then predictT+2, and so on.
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([hu‖hv])
logit_vu = head([hv‖hu])
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.
For training I use a fairly standard setup:
- Optimizer: AdamW with weight decay and gradient clipping (see
train_modelinsimple_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_usedin 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.
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.
- 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 clip1.0 - Training subset size:
edges_used=50kor100k(as written in the run title/log)
Below are three comparison plots (4 runs on each figure):
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 |
Setup:
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txtRun one dataset:
python evaluate_model.py datasets/SemanticGraph_delta_1_cutoff_5_minedge_1.pkl --no-plotsExample 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

