12 Node classification and link prediction with GNNs

This chapter presents a practical roadmap for applying Graph Neural Networks to two cornerstone graph-learning problems: node classification and link prediction. It adopts a unified, end-to-end approach that transforms semi-structured data into graphs, encodes structure and attributes with GNNs, and decodes task-specific outputs, all implemented with PyTorch Geometric. Model quality is assessed with precision, recall, and F1-score, complemented by confusion-matrix analyses to surface trade-offs. Across both tasks, three encoders are compared—GCN, GAT, and GraphSAGE—highlighting differences in efficiency, parameterization, and generalization, and showing how design choices in data processing, masking/splitting, and negative sampling shape outcomes.

For node classification, the chapter builds an anti-money laundering detector on the Elliptic Bitcoin transaction graph. Preprocessing compacts node IDs, assembles edge_index and node feature tensors, and encodes labels while masking unknowns; PyG Data objects and boolean masks orchestrate train/validation/test splits. An encoder-decoder design uses interchangeable homogeneous GNNs (via a shared base class) and a log-softmax head trained with cross-entropy to predict licit vs. illicit nodes under heavy class imbalance. Empirically, GraphSAGE provides the best balance of precision and F1, particularly improving the identification of illicit nodes; GAT performs moderately but incurs higher computational cost; and GCN is the most efficient yet struggles with precision early and is less reliable for illicit detection. Weighted averaging and confusion-matrix insights clarify these trade-offs.

For link prediction in movie recommendation, the chapter models MovieLens as a heterogeneous user–movie graph. Movie genres are one-hot encoded, IDs are remapped, and PyG HeteroData captures typed nodes and relations with undirected edges. Dataset construction leverages RandomLinkSplit to create disjoint supervision edges and generate negatives, while LinkNeighborLoader yields scalable subgraph mini-batches. The encoder first builds embeddings—learned for users, and for movies a linear transform of genres combined with embeddings—then refines them with a to_hetero-wrapped GNN (heterogeneous GCN, GAT, or SAGE). A dot-product decoder with binary cross-entropy optimizes link likelihoods. Despite larger models and longer training than in classification, patterns persist: GraphSAGE achieves the highest precision and F1, GCN attains the best recall and a strong overall balance, and GAT lags on both precision and recall. The shared pipeline demonstrates how careful data handling and modular GNN design generalize across disparate domains.

End-to-end framework to develop a GNN-based system that can be adopted for multiple tasks and application domains. The starting point is represented by the input data, which is usually a collection of semi-structured sources describing different objects and the interactions between these objects. The goal is to process such data to produce a graph that becomes the input for an encoder/decoder architecture. In our scenarios, the encoder employs a GNN model, and the decoder is based on a function that allows us to optimize the model for a specific downstream task. The output of this framework is a trained model that can be adopted for the inference phase.

This diagram illustrates the end-to-end framework tailored for node classification in the context of AML. Transaction data and node labels are transformed into a homogeneous graph structure. The encoder, a homogeneous GNN, learns node embeddings by capturing its local graph structure. The decoder classifies nodes as licit or illicit by leveraging a log softmax function and cross-entropy loss, producing a trained model for detecting suspicious activities.

Overview of data processing applied to the Elliptic dataset. Source data are preprocessed to create node features, edge indices, and encoded labels. These tensors are incorporated into a PyTorch Geometric Data object, where a node masking approach partitions the nodes into training (80%), validation (10%), and testing (10%) sets.

Overview of the main steps of the encoder-decoder architecture. Homogeneous graph data, consisting of transactions (nodes) and Bitcoin flows (edges), is processed using three distinct GNN encoders: Graph Convolutional Network (GCN), Graph Attention Network (GAT), and GraphSAGE (SAGE). Each encoder generates unique node representations by aggregating neighbor features. Then, the Log Softmax function transforms these representations into probability values, which are used to classify transactions as licit or illicit.

Performance of GCN, GAT, and SAGE across precision, recall, and f1-score on the validation dataset during the training phase.

Confusion matrices to understand the performance of GCN, GAT, and SAGE models on each class (licit or illicit).

This diagram illustrates our end-to-end framework designed for link prediction in the context of recommendation systems. Interactions/preferences data are transformed into a heterogenous graph structure that includes two types of nodes: users and movies. The encoder, a heterogenous GNN, learns node embeddings by capturing its local graph structure. The decoder predicts the existence of links between users and movies by leveraging a dot product operation combined with a binary cross-entropy loss to produce a trained model for suggesting relevant movies to users.

This diagram illustrates the data processing pipeline for the MovieLens dataset, including incremental ID mapping, genre encoding, and edge creation. The heterogeneous graph data comprises two node types (users and movies) and one edge type (user-rates-movie). Edges are split into training (80%), validation (10%), and testing (10%) sets, with negative examples generated for validation and testing. A mini-batch loader then prepares subgraphs for GNN input, ensuring scalability for large graphs that exceed memory capacity.

This diagram illustrates the pipeline for processing heterogeneous graph data in a link prediction system for movie recommendation. The data consists of users and movies as node types and user-movie ratings as edges. Embeddings are generated, with user embeddings learned by the model and movie embeddings initialized from genre features. The data is processed using three heterogeneous GNN encoders (H-GraphConv, H-GAT, and H-SAGE), which aggregate neighbor features to generate node representations. A dot product then quantifies the compatibility between users and movies, with higher scores indicating a greater likelihood of interaction. The scores are finally converted into probabilities representing the likelihood of a link between users and movies.

Trend of precision, recall, and f1-score of the GNN models on the validation dataset during the training

Confusion matrices of GCN, GAT, and SAGE models on the testing dataset

• Graph Neural Networks (GNNs) allow us to address fundamental challenges in graph-based machine learning tasks, such as node classification and link prediction. These tasks are crucial for various real-world applications, including anti-money laundering (AML) and movie recommendation systems.
• The chapter also compares different GNN architectures, such as Graph Convolutional Networks (GCN), Graph Attention Networks (GAT), and GraphSAGE (SAGE), to evaluate their performance in these domains.
• Despite the differences in tasks and application domains, we can define a unique end-to-end framework that comprises multiple steps, from processing the input data as a graph to producing a trained model for the inference phase.
• This framework is based on the encoder-decoder model introduced in chapter 11. The encoder is a GNN that generates a unique representation of the nodes, and the decoder includes a set of operations that leverages this representation for the specific downstream task.
• This model is used to detect suspicious transactions in financial networks. Transactions are classified as licit, illicit, or unknown. The Elliptic Dataset, representing Bitcoin transactions, is utilized to perform this task. It includes 200,000 transactions (nodes) and 234,000 payment flows (edges). Each transaction contains 166 anonymized features and is labeled as licit, illicit, or unknown, making it ideal for training fraud detection models.
• To prepare the data for GNNs, transactions are assigned incremental IDs for a compact representation. Original transaction IDs are removed to create node feature tensors, and transaction labels are encoded numerically. Edge indices define relationships between transactions, structuring the data for learning.
• The model is trained using three homogeneous GNN architectures based on the PyTorch Geometric (PyG) implementations of GCN, GAT, and SAGE. The Log Softmax function transforms node embeddings into probabilities for classification, while cross-entropy loss optimizes model performance.
• SAGE outperforms GCN and GAT by achieving higher accuracy in identifying illicit transactions. The confusion matrices show that it minimizes misclassification errors, while the precision, recall, and F1-score confirm that it provides the best balance between detecting fraud and minimizing false positives.
• In the context of link prediction, the encoder-decoder model seeks to predict potential interactions between users and movies based on previous ratings. By organizing user-movie relationships as a heterogeneous graph, the goal is to develop a recommendation system that suggests movies users are likely to engage with. The MovieLens dataset is used to conduct this task.
• The MovieLens dataset comprises 100,000 user ratings and 3,600 tag applications for 9,000 movies. It includes 600 users, each linked to multiple movie ratings. Movies are characterized by genre labels, which serve as input features for the model, enabling it to capture content-based similarities between films.
• To prepare the data for graph-based learning, movie genres are encoded as numerical feature vectors, enabling representation in a format suitable for GNN processing. Incremental IDs are assigned to users and movies to facilitate efficient indexing. The dataset is then organized into a heterogeneous graph with two distinct node types (users and movies) and one edge type (user-rates-movie), capturing user-movie interactions.
• The model employs heterogeneous GNN architectures of GCN, GAT, and SAGE to learn user-movie relationships. To enhance feature expressiveness, an embedding layer is implemented for both users and movies. Since users lack intrinsic features, the model learns their embeddings directly. In contrast, movie embeddings are initialized using genre-based feature vectors to provide a meaningful starting representation.
• To predict user-movie interactions, the dot product operation measures compatibility between embeddings. These scores are converted into probabilities, reflecting the likelihood that a user will engage with a particular movie.
• Among the three distinct GNN architectures tested, SAGE strikes the best balance between precision and recall, making it the most effective choice for movie recommendations. GCN excels in recall, indicating it captures more potential user-movie interactions but sacrifices precision, resulting in some irrelevant suggestions. Conversely, GAT struggles with accuracy, yielding more incorrect predictions about user preferences. An analysis of the confusion matrix confirms that SAGE minimizes false recommendations, ensuring users receive more relevant and personalized movie suggestions.
• Based on these analyses, GNNs have proven to be highly effective for both node classification, such as fraud detection, and link prediction, as seen in recommendation systems. Among the tested architectures, SAGE consistently outperforms GCN and GAT, showing superior generalization capabilities across various tasks.
• The ability to capture complex relationships makes these models valuable for real-world applications. Furthermore, end-to-end GNN frameworks can be tailored to different domains and tasks, showcasing the versatility of graph-based machine learning models across various applications.

1. Hamilton, Will, Zhitao Ying, and Jure Leskovec. "Inductive representation learning on large graphs." Advances in neural information processing systems 30 (2017).
2. Kipf, Thomas N., and Max Welling. "Semi-supervised classification with graph convolutional networks." arXiv preprint arXiv:1609.02907 (2016).
3. Veličković, Petar, et al. "Graph attention networks." arXiv preprint arXiv:1710.10903 (2017).

FAQ