A Graph Neural Network (GNN) based movie recommendation system that uses link prediction to suggest movies to users. This project implements two state-of-the-art GNN architectures (GraphSAGE and GAT) on the MovieLens dataset to predict user-movie connections in a heterogeneous graph.
This project addresses the link prediction problem in recommendation systems: given a heterogeneous graph of users and movies connected by rating relationships, can we predict new user-movie connections? The system learns from existing user ratings to recommend movies that users are likely to enjoy.
- Two GNN Architectures: Compare GraphSAGE and Graph Attention Networks (GAT) for link prediction
- Heterogeneous Graph Support: Handles different node types (users and movies) with distinct feature spaces
- Scalable Training: Mini-batch training with neighborhood sampling for large graphs
- Multiple Dataset Sizes: Implementations for both small (ml-latest-small) and large (ml-latest) MovieLens datasets
- High Performance: Achieves 90% precision/recall (GraphSAGE) and 97.6% AUC (GAT)
- Full Recommendation Pipeline: From data loading to generating top-k movie recommendations
Link-prediction/
├── Link-prediction.ipynb # GraphSAGE implementation (ml-latest-small)
├── Link-prediction-GAT.ipynb # GAT implementation (ml-latest)
├── requirements.txt # Python dependencies
├── README.md # Project documentation
├── .gitignore # Git ignore configuration
└── ml-latest-small.zip # MovieLens small dataset (956 KB)
- Python 3.8+
- CUDA-capable GPU (recommended for large datasets)
- Clone the repository:
git clone <repository-url>
cd Link-prediction- Install required packages:
pip install -r requirements.txt- Install PyTorch Geometric dependencies:
pip install torch-scatter torch-sparse pyg-lib -f https://data.pyg.org/whl/torch-$(python -c "import torch; print(torch.__version__)").htmltorch>=2.4.0- Deep learning frameworktorch_geometric>=2.5.3- Graph neural network librarynumpy>=2.2.1- Numerical computingpandas>=2.1.4- Data manipulationscikit-learn>=1.6.1- Evaluation metricstqdm>=4.66.4- Progress bars
Open and run Link-prediction.ipynb:
# The notebook will:
# 1. Download and process the ml-latest-small dataset
# 2. Train a 2-layer GraphSAGE model
# 3. Evaluate with Precision@10 and Recall@10
# 4. Generate movie recommendations for all usersTraining Configuration:
- Dataset: ml-latest-small (609 users, 9,742 movies)
- Epochs: 9
- Batch size: 10,240
- Hidden dimensions: 256
Open and run Link-prediction-GAT.ipynb:
# The notebook will:
# 1. Download and process the ml-latest dataset
# 2. Train a 2-layer Graph Attention Network
# 3. Evaluate with AUC-ROC metric
# 4. Generate top-10 recommendations with movie metadataTraining Configuration:
- Dataset: ml-latest (330,975 users, 86,537 movies)
- Epochs: 30
- Batch size: 262,144
- Hidden dimensions: 64
Architecture:
- Node Embeddings: Users (256-dim), Movies (genre features → 256-dim)
- GNN Layers: 2-layer SAGEConv with ReLU activation
- Edge Classifier: Dot-product between user and movie embeddings
- Loss Function: Bayesian Personalized Ranking (BPR)
Key Features:
- Mean-pooling aggregation for neighborhood information
- Heterogeneous graph support with separate layers per edge type
- Efficient sampling: 20 and 10 neighbors per layer
Architecture:
- Node Embeddings: Users (64-dim), Movies (genre features → 64-dim)
- GNN Layers: 2-layer GATConv with attention mechanism
- Edge Scorer: MLP on concatenated embeddings (128 → 1)
- Loss Function: Bayesian Personalized Ranking (BPR)
Key Features:
- Attention-based aggregation for weighted neighbor importance
- MLP-based edge scoring for flexible prediction
- Scalable to large graphs (330K+ users)
- Source: GroupLens
- Statistics: 609 users, 9,742 movies
- Ratings Filtered: Only 4.0, 4.5, and 5.0 stars (48,580 ratings)
- Features: 20 movie genres as binary indicators
- Source: GroupLens
- Statistics: 330,975 users, 86,537 movies
- Ratings: 33,832,162 ratings (all ratings included)
- Features: 20 movie genres as binary indicators
Genres: Action, Adventure, Animation, Children, Comedy, Crime, Documentary, Drama, Fantasy, Film-Noir, Horror, IMAX, Musical, Mystery, Romance, Sci-Fi, Thriller, War, Western
| Model | Dataset | Metric | Score |
|---|---|---|---|
| GraphSAGE | ml-latest-small | Precision@10 | 0.9000 |
| GraphSAGE | ml-latest-small | Recall@10 | 0.9000 |
| GAT | ml-latest | AUC-ROC | 0.9760 |
- Load MovieLens ratings and movie metadata
- Extract genre features as 20-dimensional binary vectors
- Map IDs to consecutive integers for graph representation
- Create heterogeneous graph with User and Movie nodes
- Add edges from user-movie ratings
- Apply bidirectional edge transformation for message passing
- 70% training, 10% validation, 10% test
- Disjoint training: 30% of edges removed from graph during training
- Negative sampling: 2:1 ratio for contrastive learning
- LinkNeighborLoader samples local neighborhoods
- Neighborhood sampling: 20 neighbors (layer 1), 10 neighbors (layer 2)
- BPR loss maximizes ranking of positive edges over negative edges
- Perform inference on all user-movie pairs
- Rank movies by predicted edge scores
- Return top-k recommendations with movie metadata
Input Graph (Users + Movies)
↓
Node Feature Extraction
- Users: Learned embeddings
- Movies: Genre features + embeddings
↓
GNN Message Passing (2 layers)
- GraphSAGE: Mean aggregation
- GAT: Attention-weighted aggregation
↓
Edge-Level Prediction
- Dot-product (GraphSAGE)
- MLP scorer (GAT)
↓
BPR Loss Optimization
↓
Top-K Recommendations
Predicting missing or future edges in a graph. In this context, predicting which movies a user will rate highly.
A graph with multiple node types (users and movies) and edge types (ratings), where different types may have different feature spaces.
A pairwise ranking loss that maximizes the score difference between positive (observed) and negative (unobserved) edges:
Loss = -log(σ(score_positive - score_negative))
Instead of using the full graph, sample a fixed number of neighbors per node to create mini-batches for scalable training.
| Aspect | GraphSAGE | GAT |
|---|---|---|
| Aggregation | Mean pooling | Attention-weighted |
| Edge Scoring | Dot-product | MLP (concatenation) |
| Dataset | Small (609 users) | Large (330K users) |
| Hidden Dim | 256 | 64 |
| Batch Size | 10,240 | 262,144 |
| Metric | Precision/Recall | AUC-ROC |
- Implement temporal dynamics for evolving user preferences
- Add content-based features (movie plots, actors, directors)
- Incorporate implicit feedback (views, clicks)
- Experiment with deeper architectures (3+ layers)
- Add explainability for recommendations
- Deploy as REST API service
Contributions are welcome! Please feel free to submit a Pull Request.
- MovieLens Dataset by GroupLens Research
- PyTorch Geometric for GNN implementations
- GraphSAGE paper: Hamilton et al., 2017
- GAT paper: Veličković et al., 2018
This project is available for educational and research purposes.
If you use this code in your research, please cite:
@software{link_prediction_gnn,
title={Link Prediction for Movie Recommendations using Graph Neural Networks},
author={Your Name},
year={2025},
url={https://github.com/yourusername/Link-prediction}
}Contact: For questions or feedback, please open an issue in the repository.