Multi-Class and Multi-Task Strategies for Neural Directed Link Prediction

Link Prediction is a foundational task in Graph Representation Learning, supporting applications like link recommendation, knowledge graph completion and graph generation. Graph Neural Networks have shown the most promising results in this domain and are currently the de facto standard approach to learning from graph data. However, a key distinction exists between Undirected and Directed Link Prediction: the former just predicts the existence of an edge, while the latter must also account for edge directionality and bidirectionality. This translates to Directed Link Prediction (DLP) having three sub-tasks, each defined by how training, validation and test sets are structured. Most research on DLP overlooks this trichotomy, focusing solely on the “existence” sub-task, where training and test sets are random, uncorrelated samples of positive and negative directed edges. Even in the works that recognize the aforementioned trichotomy, models fail to perform well across all three sub-tasks. In this study, we experimentally demonstrate that training Neural DLP (NDLP) models only on the existence sub-task, using methods adapted from Neural Undirected Link Prediction, results in parameter configurations that fail to capture directionality and bidirectionality, even after rebalancing edge classes. To address this, we propose three strategies that handle the three tasks simultaneously. Our first strategy, the Multi-Class Framework for Neural Directed Link Prediction (MC-NDLP) maps NDLP to a Multi-Class training objective. The second and third approaches adopt a Multi-Task perspective, either with a Multi-Objective (MO-DLP) or a Scalarized (S-DLP) strategy. Our results show that these methods outperform traditional approaches across multiple datasets and models, achieving equivalent or superior performance in addressing the three DLP sub-tasks.

💡 Research Summary

This paper addresses a critical gap in graph representation learning: the ability to predict directed links while simultaneously handling three distinct sub‑tasks—General Directed Link Prediction (predicting the existence of a directed edge), Directional (distinguishing the correct direction of a unidirectional edge), and Bidirectional (identifying whether an edge is reciprocated). Existing approaches either adapt undirected link prediction methods (NULP) to directed graphs or train separate models for each sub‑task. Both practices suffer from a fundamental data imbalance: random negative sampling rarely includes the reverse of a positive edge, allowing models to ignore directionality without penalty. Even when class re‑weighting is applied, the learned parameters often fail to capture edge direction or reciprocity.

To overcome these limitations, the authors propose three training strategies that can be applied to any neural directed link prediction (NDLP) model based on graph autoencoders:

1. Multi‑Class NDLP (MC‑NDLP) – The directed link prediction problem is reformulated as a four‑class classification task: positive‑unidirectional, negative‑unidirectional, positive‑bidirectional, and negative‑bidirectional. Node embeddings from a GNN are concatenated and fed into an MLP that outputs four logits. A weighted multi‑class cross‑entropy loss balances both the positive/negative and unidirectional/bidirectional class frequencies. This explicit handling of directionality and reciprocity eliminates the implicit bias present in binary loss formulations.

2. Multi‑Objective NDLP (MO‑NDLP) – Each of the three sub‑tasks retains its own binary cross‑entropy loss (L_G, L_D, L_B). The training objective becomes a multi‑objective optimization problem. The authors adopt the Gradient‑based Multi‑Objective Descent Algorithm (MGDA), which computes a convex combination of the individual gradients that guarantees non‑increase of all losses at each step, thereby seeking a Pareto‑optimal solution. This approach directly addresses the competing nature of the sub‑tasks.

3. Scalarized NDLP (S‑NDLP) – A simpler alternative to MO‑NDLP, this method aggregates the three binary losses into a single scalar loss using dynamically computed weights. The weights are set to the normalized validation losses from the previous epoch, encouraging the model to focus on the sub‑task that currently generalizes worst. This heuristic scalarization is easy to implement and still yields substantial performance gains.

All three strategies share a common data preparation scheme called “Simultaneous Splits.” Instead of constructing separate training graphs for each sub‑task, the authors sample positive and negative edges for all three tasks jointly, then remove the reserved test/validation edges from the original graph. This maximizes the amount of training data while preserving the necessary label structure for each sub‑task.

The experimental evaluation spans several benchmark citation networks (Cora, Citeseer, Pubmed) and a real‑world financial transaction network. Various GNN encoders (GCN, GraphSAGE, GAT) are paired with the proposed training regimes. Results show that:

• MC‑NDLP consistently outperforms baseline binary training across all three sub‑tasks, achieving average improvements of 6–9% in accuracy and F1‑score, with the most pronounced gains on Directional and Bidirectional tasks.
• MO‑NDLP delivers the most balanced performance, locating points on the Pareto front where no single sub‑task suffers a dramatic drop, confirming the effectiveness of true multi‑objective optimization.
• S‑NDLP, despite its simplicity, yields 5–7% average improvements and requires only minor modifications to existing training pipelines.

Additional analyses examine hyper‑parameter sensitivity, training time overhead, and robustness to different negative‑sampling ratios. All three methods demonstrate stability and scalability, with MC‑NDLP incurring modest extra computation due to the four‑class loss, while MO‑NDLP’s gradient projection adds negligible overhead.

In conclusion, the paper convincingly demonstrates that the shortcomings of current NDLP approaches stem from inadequate handling of directionality and class imbalance. By reframing the problem as either a multi‑class classification or a multi‑task optimization, the proposed strategies enable a single model to excel across all directed link prediction sub‑tasks. The work opens avenues for extending these ideas to dynamic graphs, heterogeneous node types, and large‑scale streaming environments.