Temporal Graph Networks (TGN) for Dynamic Interaction Modeling

To know more about the Temporal graph networks please read my blog on it

Graph networks Frozen in Time Are Useless. Temporal Graph Networks Aren’t.

This repository contains a high-performance implementation of a Temporal Graph Network (TGN) applied to the Wikipedia JODIE dataset. The model is designed to predict future interactions (link prediction) in a dynamic graph where edges are added over time.

Overview

Temporal Graph Networks are a class of Graph Neural Networks (GNNs) specifically designed for Dynamic Graphs. Unlike static graphs, dynamic graphs evolve over time, with nodes and edges appearing or changing state. TGNs maintain a "memory" for each node, allowing the model to capture long-term dependencies and the evolution of node interactions.

Key Features:

• TGN Memory: Persistent state for every node, updated after each interaction.
• Time2Vec: Non-linear time encoding to capture periodic and trend-based temporal patterns.
• Graph Attention: Multi-head attention mechanism to aggregate information from a node's temporal neighborhood.
• Chronological Training: Strict temporal splitting and batching to prevent data leakage.

---

Dataset: Wikipedia JODIE

The model is trained on the Wikipedia dataset from the JODIE benchmark.

• Nodes: 9,227 (Users and Pages)
• Events: 157,474 (Edits/Interactions)
• Features: 172-dimensional edge features representing the edit metadata.
• Task: Predict the next interaction given the history of the graph.

---

Architecture

The architecture consists of four main components:

1. Memory Module: A GRU-based memory that stores a compressed representation of each node's history.
2. Time Encoder (Time2Vec): Maps raw timestamps into a vector space using sine and cosine functions to learn periodic temporal behaviors.
3. Embedding Module (GNN): A TransformerConv layer that performs graph convolution over the temporal neighborhood, combining node memories and edge features.
4. Link Predictor: An MLP decoder that scores the likelihood of an interaction between two node embeddings.

---

Results

The model achieves high accuracy on the link prediction task. Below is the visualization of the training progress, showing the AUC (Area Under Curve) and Training Loss over 50 epochs.

Performance Summary:

• Best Validation AUC: Obtained during training and used for checkpointing.
• Test Performance: The model demonstrates robust generalization on unseen future interactions.

---

Project Structure

• TGN_best_Model.ipynb: The main Jupyter notebook containing the full pipeline (data loading, model definition, training, and evaluation).
• Artifacts/:
  • best_tgn.pth: Saved model weights for the best performing epoch.
  • tgn_training.png: Training curve visualization.
• data/: (Generated upon run) Directory for the Wikipedia dataset.

---

Usage

Prerequisites

Ensure you have the following installed:

• PyTorch
• Torch Geometric
• Scikit-learn
• Matplotlib

Running the Notebook

1. Open TGN_best_Model.ipynb in Jupyter or Kaggle.
2. Run the cells in order.
3. The model will automatically download the dataset, train for 50 epochs, and save the best weights to best_tgn.pth.

---

Acknowledgments

• Implementation based on the Temporal Graph Networks for Deep Learning on Dynamic Graphs paper.
• Dataset provided by the Stanford Network Analysis Project (SNAP).