Labeling Trick: A Theory of Using Graph Neural Networks for Multi-Node Representation Learning (2010.16103v5)

Abstract: In this paper, we provide a theory of using graph neural networks (GNNs) for multi-node representation learning (where we are interested in learning a representation for a set of more than one node, such as link). We know that GNN is designed to learn single-node representations. When we want to learn a node set representation involving multiple nodes, a common practice in previous works is to directly aggregate the single-node representations obtained by a GNN into a joint node set representation. In this paper, we show a fundamental constraint of such an approach, namely the inability to capture the dependence between nodes in the node set, and argue that directly aggregating individual node representations does not lead to an effective joint representation for multiple nodes. Then, we notice that a few previous successful works for multi-node representation learning, including SEAL, Distance Encoding, and ID-GNN, all used node labeling. These methods first label nodes in the graph according to their relationships with the target node set before applying a GNN. Then, the node representations obtained in the labeled graph are aggregated into a node set representation. By investigating their inner mechanisms, we unify these node labeling techniques into a single and most general form -- labeling trick. We prove that with labeling trick a sufficiently expressive GNN learns the most expressive node set representations, thus in principle solves any joint learning tasks over node sets. Experiments on one important two-node representation learning task, link prediction, verified our theory. Our work explains the superior performance of previous node-labeling-based methods, and establishes a theoretical foundation of using GNNs for multi-node representation learning.

• The paper introduces the 'labeling trick' to overcome limitations in aggregating node representations for multi-node learning.
• It provides theoretical proofs showing that labeled node representations better capture inter-node dependencies in GNNs.
• Experimental evaluations on link prediction tasks validate the improved accuracy and broader applicability of labeling-based GNNs.

Labeling Trick: A Theory of Using Graph Neural Networks for Multi-Node Representation Learning

The paper provides a comprehensive theory on enhancing the capabilities of Graph Neural Networks (GNNs) for multi-node representation learning tasks, such as link prediction. Traditional GNNs are optimized primarily for single-node representations, and extending these methods to effectively predict outcomes for node sets, including links, has been challenging. The authors identify a fundamental limitation with directly aggregating independently learned node representations to form joint representations, which often fails to capture the dependencies among nodes within a set.

The paper introduces the concept of the "labeling trick," a technique that unifies node-labeling strategies used in successful models like SEAL. This method involves first labeling nodes in a graph based on their relationship to the target node set before training the GNN. The labeled node representations are then aggregated to construct the joint representation. The key argument is that this method better captures node dependencies and improves prediction accuracy.

Main Contributions

1. Limitations of Current GNN Aggregation Methods:
   • The paper critiques the conventional approach of aggregating node representations without considering the interdependencies between nodes in a set. Such methods, exemplified by Graph AutoEncoder (GAE) models, have shown limitations as they can fail to distinguish between structurally different node sets.
2. Labeling Trick and Its Theoretical Underpinnings:
   • The authors detail the "labeling trick," which involves adding node labels to enhance GNNs' expressive power. They assert that with sufficiently expressive GNNs, this approach can generate the most expressive node set representations, thus overcoming previous aggregation method limitations.
3. Comprehensive Evaluation:
   • Experimental validation is performed using link prediction tasks, demonstrating the superior performance of labeling-based GNNs over traditional aggregation methods.
4. Theoretical Proofs:
   • The paper provides rigorous proofs that labeling tricks ensure GNNs can learn structural representations of node sets. The proofs establish that node-most-expressive GNNs, when combined with an injective aggregation function, map isomorphic node sets to identical representations.
5. Broader Implications:
   • By addressing the gap in GNNs' ability to handle multi-node inputs robustly, the research suggests new methodologies for applying GNNs in diverse AI domains requiring effective multi-node reasoning, from social networks to biochemical graphs.

Implications and Future Directions

Practically, this research offers a path forward for employing GNNs in multi-node prediction tasks by leveraging node labeling techniques. The findings make GNNs viable for more complex tasks than previously feasible, such as link prediction in heterogeneous graphs, knowledge graph completion, and recommendation systems.

Theoretically, the insights into permutation equivariance and the necessity of capturing inter-node dependencies set a foundation for further innovations in GNN architectures. This work invites future research to explore alternative implementations of the labeling trick and identify other potential applications of this methodology.

In summary, this paper significantly enhances the theoretical understanding and practical application of GNNs in multi-node representation learning. It bridges a crucial gap, paving the way for broader adoption and efficacy of GNNs in complex graph-based tasks.