Source-linked AI summary
Labeling Trick: A Theory of Using Graph Neural Networks for Multi-Node Representation Learning
Muhan Zhang, Pan Li, Yinglong Xia, Kai Wang, Long Jin
TL;DR
The paper asks whether GNNs can learn effective representations for sets of nodes when their individual representations are aggregated. It formalizes labeling trick, proves its expressive power for node sets, and evaluates the theory on link prediction, while identifying a dense-graph limitation in experiments.
Problem
Directly aggregating independently learned node representations cannot capture dependence between nodes in a target set, limiting multi-node representation learning.
Method
The paper unifies target-node-distinguishing and permutation-equivariant labeling methods as labeling trick and applies a sufficiently expressive GNN to the labeled graph.
Results
Labeling trick enables aggregation of labeled-graph node representations into structural node-set representations, with experiments verifying effectiveness on link prediction.
Takeaways & Limitations
The theory establishes a foundation for using GNNs in multi-node representation learning beyond directly aggregating independently learned node representations.
Takeaways & Limitations
On the dense ogbl-ddi dataset, labeling-trick methods fall behind GAE methods using trainable node embeddings.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper studies GNNs for multi-node representation learning, focusing on link prediction, and identifies a limitation in aggregating independently learned node representations. It proposes labeling trick as a general solution with theoretical and experimental support.
- Multi-node representation learning with GNNs is less studied than single-node and whole-graph tasks, although link prediction has broad practical applications.
- GAE first computes individual node representations with a GNN and then aggregates the target nodes’ representations to predict their link.
- Symmetric nodes can receive identical representations, causing aggregated GAE representations to assign the same prediction to structurally different links.
- GAE cannot capture neighborhood-overlap features such as common neighbors because each node representation is computed independently of the other target node.
- Node-discriminating input features may distinguish targets but sacrifice inductive generalization and violate GNN permutation invariance or equivariance.
- Labeling trick requires target-node distinction and permutation equivariance, enabling a sufficiently expressive GNN to learn the most expressive structural representations of node sets.
2 Preliminaries
The preliminaries formalize permutations, set isomorphism, permutation invariance and equivariance, and most expressive structural representations. These concepts define when node-set representations should agree or differ.
- A target node set S is transformed under a node permutation π by mapping every member i to π(i).
- Set isomorphism requires a permutation that maps both one graph’s target set and attributed graph structure to another’s.
- Set isomorphism is stricter than graph isomorphism because the permutation must also map the specified node set.
- Permutation invariance and equivariance ensure GNN representations do not depend improperly on node indexing.
- A most expressive structural representation assigns the same representation exactly to isomorphic node sets and different representations to non-isomorphic sets.
- Structural representations are studied because joint prediction tasks over node sets require no information beyond the structural representation.
3 The limitation of directly aggregating node representations
This section shows that directly aggregating GNN node representations cannot produce structural link representations, even when the node encoder is maximally expressive. The failure arises because node representations are learned independently rather than jointly conditioned on the target set.
- 3.1 GAE for link prediction: GAE predicts a link by applying a GNN independently to each endpoint and combining their representations with an aggregation function.
- 3.1 GAE for link prediction: GAE’s aggregation choices include inner product, mean, sum, bilinear product, concatenation, and Hadamard product.
- 3.2 Node representations: A 1-WL-GNN preserves equality for isomorphic nodes but may also assign equal representations to some non-isomorphic nodes.
- 3.2 Node representations: A node-most-expressive GNN is defined to assign different representations to every pair of non-isomorphic nodes.
- 3.3 GAE cannot learn structural link representations: Even with a node-most-expressive GNN, aggregating endpoint representations can give non-isomorphic links the same representation.
- 3.3 GAE cannot learn structural link representations: GAE cannot learn structural link representations regardless of the expressiveness of its node representations.
- 3.3 GAE cannot learn structural link representations: The same limitation extends beyond links because multi-node representation learning cannot be decomposed into independent single-node problems.
4 Labeling trick for multi-node representation learning
Labeling trick augments a graph with target-set-dependent labels before GNN processing, allowing node-set representations to preserve target-node relationships while respecting graph symmetries. Under sufficient expressiveness, this yields the most expressive structural representations, though relabeling can add computational cost.
- General formulation: Labeling trick unifies prior methods by adding a target-dependent labeling tensor to the graph before applying a GNN.The labeling must distinguish target nodes and remain permutation equivariant.
- General formulation: A valid labeling distinguishes target nodes and preserves labels consistently under graph permutations.These are the target-nodes-distinguishing and permutation-equivariance properties.
- Zero-one labeling: Zero-one labeling assigns label 1 to target nodes and label 0 elsewhere, providing additional node features to the GNN.This labeling is permutation equivariant and satisfies the required validity conditions.
- Zero-one labeling: With target nodes labeled, message passing becomes aware of the other nodes in the set and can encode relationships such as common-neighbor counts.For a link, common neighbors receive messages from both labeled endpoints and return them in later iterations.
- Theoretical guarantee: Theorem 1 states that a node-most-expressive GNN with injective set aggregation produces identical representations exactly for isomorphic labeled node sets.Thus, aggregating representations from the labeled graph yields a structural representation, unlike direct aggregation from the original graph.
- Theoretical guarantee: Labeling trick enables GNNs to learn most expressive structural representations for joint prediction over node sets, including multi-node tasks beyond links.The result addresses the limitation of GNNs that process single-node representations independently.
- SEAL and complexity: SEAL extracts an h-hop enclosing subgraph and applies a labeling trick there, while labeling every target node set may increase computation.For large graphs requiring mini-batches, the paper reports similar computation costs for GAE-type methods and labeling trick.
5 Local isomorphism: a more practical view of isomorphism
Local isomorphism restricts structural comparison to h-hop enclosing subgraphs around target node sets rather than requiring whole-graph isomorphism. This gives a more practical notion for representations while preserving the paper’s corresponding guarantees.
- Motivation: Exact isomorphism is rare, and 1-WL-GNNs may require up to O(n) layers to reach their discrimination limit.The paper motivates a local alternative because whole-graph structural equivalence is uncommon in practice.
- Definition: Local isomorphism compares whether the h-hop enclosing subgraphs around two target node sets are isomorphic.It replaces whole-graph comparison with a neighborhood-level criterion.
- Motivation: Local h-isomorphism can better group nodes or links with similar neighborhoods and may reduce overfitting to exact graph symmetries.The paper argues that exact-isomorphism methods may fail to identify structurally similar but non-isomorphic neighborhoods.
- Guarantee: A node-most-expressive GNN with injective aggregation can learn node-set representations that coincide exactly for locally h-isomorphic sets.This is the local analogue of the global structural-representation guarantee.
- SEAL connection: SEAL implements the local view by extracting each target set’s h-hop enclosing subgraph before applying labeling trick and a GNN.The paper identifies this procedure as exactly the operation needed to switch from global to local isomorphism.
6 Related work
Related work contrasts direct aggregation of independently learned node representations with labeling-based approaches that encode target-node dependence. The paper places SEAL, distance encoding, and ID-GNN within this broader labeling-based line of work.
- GNN representation power: Prior GNN studies largely analyze whole-graph representation power, while node- and link-set representation power receives less attention.Neighbor-aggregation GNNs are bounded by the 1-WL test, motivating related efforts to increase expressive power.
- Direct aggregation: Directly aggregating pairwise node representations can fail to distinguish many non-isomorphic links because independently computed node embeddings omit node dependence.The paper contrasts this limitation with structural link representations obtained through labeling trick.
- Position-aware methods: P-GNN produces position-aware embeddings from randomly selected anchor nodes but cannot learn structural node or link representations and does not scale to large datasets.Its embeddings reflect distances to anchors rather than neighborhood structure.
- Labeling-based methods: Specific labeling tricks preceded the paper’s formal framework: SEAL introduced one for link prediction, DE generalized it to sets larger than two, and later work applied it to other tasks.Applications include inductive knowledge-graph completion, matrix completion, and hyperedge prediction.
7 Experiments
Experiments evaluate labeling trick on four large-scale OGB link-prediction datasets using Hits@K and MRR. Labeling-based GNNs generally outperform plain GAE methods and heuristics, but underperform on the densest dataset.
- Setup: The evaluation uses ogbl-ppa, ogbl-collab, ogbl-ddi, and ogbl-citation2 with realistic splits and Hits@K and MRR metrics.The datasets contain up to 2.9M nodes and 30.6M edges.
- Main results: Plain GAE methods do not consistently outperform common-neighbor and Adamic-Adar heuristics on ogbl-ppa and ogbl-collab.This supports the paper’s argument that direct node-representation aggregation may miss neighborhood-overlap signals.
- Main results: Labeling-trick GNNs achieve state-of-the-art performance on three of four datasets and outperform GAE and positional-embedding methods in several comparisons.The paper reports these results on ogbl-ppa, ogbl-collab, and ogbl-citation2.
- Main results: 48.80 Hits@100: SEAL on ogbl-ppa, reported as 87%-195% higher than GAE methods without labeling trick.This is the paper’s representative large-margin comparison.
- Limitation: Labeling-trick methods perform poorly on dense ogbl-ddi, where they trail GAE methods using trainable node embeddings.The paper attributes this possibly to the difficulty of learning structural patterns inductively in the dense graph.
8 Conclusions
The paper formalizes labeling trick as a general solution for multi-node representation learning and characterizes its expressive power. It also establishes limits of 1-WL-GNNs and gives an asymptotic separation showing labeling trick can distinguish links that ordinary message passing cannot.
- 8 Conclusions: Direct aggregation of individual node representations is limited for multi-node representation learning, motivating a structural analysis.The conclusion states this is the key limitation addressed by the paper.
- 8 Conclusions: Labeling trick assigns target nodes distinct labels in a permutation-equivariant way and enables GNNs to learn most expressive structural representations of node sets.The construction is based on a labeling tensor satisfying target-set discrimination and permutation equivariance conditions.
- Theorem 1: The proof concludes by showing that the labeling conditions force target-set correspondence under graph permutation, while equivariance preserves corresponding labels.These two directions establish the structural equivalence characterized by the theorem.
- Theorem 2: For any constant ϵ > 0, there exist ω(n2ϵ) non-isomorphic link pairs that an h-layer 1-WL-GNN represents identically, while labeling trick distinguishes them.This result holds under the stated degree condition in the theorem.
- Theorem 2: The proof constructs ω(nϵ) locally h-isomorphic nodes and then forms at least ω(n2ϵ) non-isomorphic link pairs from them.The construction partitions locally isomorphic nodes and selects shared neighbors to create non-isomorphic links.
C Proof of Lemma 1
Lemma 1 bounds the number of message-passing layers needed for a 1-WL-GNN to realize all distinctions that 1-WL can make. The bound is linear in the graph size, and paths attain this order.
- Lemma 1: O(n) message-passing layers suffice for a 1-WL-GNN to discriminate all nodes that 1-WL can discriminate.The argument follows because 1-WL increases its total number of colors before convergence and has at most n final colors.
- Lemma 1: A path of n nodes requires ⌈n/2⌉ iterations for 1-WL to converge, establishing an O(n) worst-case example.The path construction supplies the matching linear-order convergence example.
- DRNL and DE: DRNL and DE encode distances from each node to the two link endpoints, and their encodings have theoretically equivalent expressive power.DRNL compresses the distance pair into an integer label, whereas DE retains a size-2 distance vector.
- DRNL and DE: DRNL masks the opposite endpoint when computing each distance, whereas DE may cap distances at dmax.The masking choice is intended to preserve the distance to each endpoint without influence from the other endpoint.
- DRNL and DE: DE+ catches up with DRNL on datasets where DE underperforms after adopting DRNL’s masking trick and removing DE’s maximum-distance cap.The paper attributes the practical difference between theoretically equivalent encodings to implementation differences.
E More discussion on the differences between DE’s theory and ours
The paper broadens the analysis beyond DE’s distance-based theory, showing that labeling trick applies to arbitrary graphs and need not rely on distance labels. It also identifies an unresolved transfer question for arbitrary labeling tricks on regular graphs.
- DE’s theory: DE’s theory shows that a 1-WL-GNN combined with distance encoding can discriminate uniformly sampled node-set pairs from r-regular graphs with high probability.Without DE, 1-WL-GNNs fail to discriminate node sets in this graph class.
- DE’s limitations: DE’s analysis is limited to random graphs, does not establish structural representation learning, and leaves non-distance labeling tricks unexamined.These limitations motivate the broader theory developed in this paper.
- This paper’s theory: The paper defines labeling trick through permutation equivariance and target-node-set discrimination, proves structural representation learning with sufficiently expressive GNNs, and applies the result to arbitrary graphs.Thus, its scope is not restricted to distance-based labeling or regular graphs.
- Open question: It remains uncertain whether DE’s regular-graph discrimination power transfers to every valid labeling trick, especially those not based on distance.The paper leaves a rigorous proof or disproof for future work.
F More details about the datasets
The evaluation uses four large-scale OGB link-prediction datasets covering protein associations, author collaborations, drug interactions, and citation links. Their realistic splits, open availability, standard metrics, and official leaderboard support benchmarking practical performance.
- The four datasets represent protein-protein association, author collaboration, drug-drug interaction, and citation-link prediction tasks.
- The datasets use realistic splits based on laboratory resource cost, time, or drug targets in the body.
- The benchmark is large-scale, reaching 2.9M nodes and 30.6M edges, and uses standard evaluation metrics.
- OGB provides an official leaderboard for fair comparison of link-prediction performance.
G More details about the baselines
The baseline comparison spans non-GNN methods, plain GAE models, and GNNs enhanced with labeling tricks. The implementations differ in representation construction, graph processing, training data usage, and computational requirements.
- Plain GAE methods form link representations from pairwise node representations without node labeling.
- All compared methods feed link representations to an MLP, while labeling-trick methods extract 1-hop enclosing subgraphs.
- On ogbl-ddi, GAE methods use jointly trained free-parameter node embeddings because the graph has no node features.
- GCN+LRGA required reduced hidden dimension and matrix rank on ogbl-citation2 because its default configuration exceeded 32GB GPU memory.The reduced settings may limit the baseline’s achievable performance.
- Labeling-trick methods used only 1% to 10% of available positive training edges and generally required longer runtimes than GAE methods.
H.1 How powerful is the zero-one labeling trick?
The ablations examine zero-one labeling, distance-label implementations, subgraph-level readout, and message-passing depth. They show that practical performance depends not only on labeling validity but also on GNN expressiveness and implementation choices.
- How powerful is the zero-one labeling trick?: Zero-one labeling improves over unlabeled GCN but performs below DRNL, especially on ogbl-citation2.
- How powerful is the zero-one labeling trick?: Replacing GCN with GIN substantially improves zero-one labeling and nearly closes the gap with GCN+DRNL.This supports the role of GNN expressive power in realizing the labeling trick’s theoretical potential.
- DE vs. DE+ vs. DRNL: DE+ substantially outperforms DE, indicating that masking the other target node is important for producing diverse relative-position labels.
- DE vs. DE+ vs. DRNL: DE+ without a maximum distance performs very similarly to DRNL, while masking benefits are not observed on smaller ogbl-collab.
- Is a subgraph-level readout useful?: Subgraph-level sum-pooling greatly improves ogbl-citation2 performance but slightly reduces ogbl-collab performance.The authors attribute its practical benefit to extracting richer subgraph features when message-passing depth is limited.
- Is it helpful to make number of layers larger than number of hops?: Using three message-passing layers with one-hop subgraphs outperforms using one layer across all three tested datasets.The performance gap is small on ogbl-collab but significant on ogbl-ppa and ogbl-citation2.
- Is it helpful to make number of layers larger than number of hops?: The benefit of using more layers than hops remains uncertain for larger subgraphs because increasing hop count exponentially increases subgraph size.Only ogbl-collab permitted broader hop-depth experiments, so it may not support a general conclusion.