Source-linked AI summary
Representation Learning on Graphs with Jumping Knowledge Networks
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, Stefanie Jegelka
TL;DR
Neighborhood aggregation can expose nodes to very different effective neighborhoods because influence ranges depend on graph structure, while fixed-depth models may degrade as they become deeper. The paper introduces Jumping Knowledge Networks, which learn to combine information from different neighborhood ranges for each node and task. JK networks improve representations on diverse graph structures and consistently improve several existing aggregation models.
Problem
Neighborhood aggregation has structure-dependent influence ranges, and deeper GCNs can perform worse than two-layer models, motivating adaptive representations for graphs with diverse substructures.
Method
Jumping Knowledge Networks selectively combine representations from different aggregation depths at the last layer to learn node- and task-specific neighborhood ranges.
Results
JK networks improve representations on graphs with diverse local structure and consistently improve models including GCN, GraphSAGE, and GAT.
Takeaways & Limitations
Adaptive neighborhood ranges are particularly useful when fixed numbers of neighborhood aggregations do not capture diverse graph substructures.
Takeaways & Limitations
After O(log |V|) aggregation iterations, representations inside expander regions can be influenced almost equally by every node there, limiting individual-node information.
Abstract
from arXiv · showhide
Recent deep learning approaches for representation learning on graphs follow a neighborhood aggregation procedure. We analyze some important properties of these models, and propose a strategy to overcome those. In particular, the range of "neighboring" nodes that a node's representation draws from strongly depends on the graph structure, analogous to the spread of a random walk. To adapt to local neighborhood properties and tasks, we explore an architecture -- jumping knowledge (JK) networks -- that flexibly leverages, for each node, different neighborhood ranges to enable better structure-aware representation. In a number of experiments on social, bioinformatics and citation networks, we demonstrate that our model achieves state-of-the-art performance. Furthermore, combining the JK framework with models like Graph Convolutional Networks, GraphSAGE and Graph Attention Networks consistently improves those models' performance.
1. Introduction
Neighborhood aggregation models draw information from ranges whose locality varies with graph structure, making fixed-depth representations poorly suited to heterogeneous subgraphs. The paper proposes Jumping Knowledge Networks to learn node- and task-specific influence radii by selectively combining representations from different layers.
- Neighborhood aggregation iteratively combines each node’s hidden features with those of adjacent nodes, with each iteration corresponding to a neural-network layer.
- Deeper GCNs can perform worse than two-layer GCNs despite accessing more information, even when residual connections are added.
- Model analysis: The paper analyzes influence distributions as the effective ranges from which node representations draw information and relates them to random-walk spread.Short-range influence can sometimes be more informative than higher-order features whose information is washed out by averaging.
- Changing locality: A node’s effective influence range depends strongly on graph structure, so the same number of layers can produce very different neighborhood localities.Random-walk expansion is rapid in well-connected cores but remains local in tree-like regions.
- JK networks: Jumping Knowledge Networks selectively combine aggregations from different localities at the last layer, allowing representations to adapt to each node and task.The authors report that the networks learn different representation orders for different graph substructures and improve several neighborhood-aggregation models.
2. Background and Neighborhood aggregation schemes
Neighborhood aggregation schemes update node representations by aggregating neighboring features and applying model-specific transformations. Existing skip connections and directional biases address related design choices but do not independently adapt final-layer neighborhood sizes for different output nodes.
- A graph is represented as G = (V, E) with node features Xv, while h(l)v denotes node v’s hidden feature at layer l.The augmented graph eG adds a self-loop to every node, and its neighborhood includes the node itself.
- A generic k-layer aggregation updates each node at every layer using an aggregation function, a trainable layer weight matrix, and a nonlinear activation.
- Graph Convolutional Networks: GCNs instantiate this framework using normalized neighborhood averaging, with a related normalization enabling inductive use on previously unseen nodes.
- Neighborhood Aggregation with Skip Connections: Skip-connection schemes aggregate neighbors and then combine that representation with the node’s previous representation using model-specific operations such as concatenation, interpolation, or a GRU.
- Neighborhood Aggregation with Skip Connections: Because skip connections are input- but not output-unit specific, they cannot independently adapt final-layer neighborhood sizes across representations.
- Neighborhood Aggregation with Directional Biases: Directional-bias models weight neighbors unequally, whereas JK networks change locality of expansion; the two approaches are orthogonal and can be combined.
3. Influence Distribution and Random Walks
The paper analyzes how node representations draw information from other nodes by defining influence distributions and relating them to random-walk behavior. This connection exposes how graph structure and aggregation choices shape locality, including rapid global mixing in some subgraphs and retained local information in others.
- Influence distributions: Influence distributions quantify how changes in one node’s input feature affect another node’s final representation.The influence score is based on absolute Jacobian entries and is normalized across source nodes.
- Connection to random walks: Random-walk distributions spread as steps increase, with convergence rates determined by subgraph structure and spectral properties.The same iteration count can therefore produce different influence localities on different subgraphs.
- Connection to random walks: A k-layer GCN’s influence distribution is equivalent in expectation to a k-step random-walk distribution on the self-looped graph, under the stated activation assumption.The result assumes all computation paths are activated with the same probability ρ.
- Aggregation variants: Directional neighborhood aggregation schemes resemble biased random walks, while residual connections approximately correspond to lazy random walks that retain more probability at the current node.The residual-network correspondence is observed empirically, and the lazy-walk comparison uses a higher stay probability.
- Implications: These findings motivate adaptive aggregation because fixed numbers of layers cannot preserve equally appropriate locality across diverse graph substructures.The paper identifies uniform local-information retention across nodes as a limitation of some aggregation schemes.
- Implications: On expander subgraphs, influence can collapse toward an almost-uniform global distribution after O(log |V|) aggregation steps, whereas bounded-tree-width regions retain locality longer.Fixed-range models inherit these contrasting expansion speeds, which may produce unsuitable representations across nodes.
4. Jumping Knowledge Networks
JK-Nets address structure-dependent and fixed locality in neighborhood aggregation by selecting and combining intermediate representations separately for each node. Their aggregation mechanisms can adapt neighborhood ranges to subgraph structure, including tree-like regions and hub-affiliated nodes.
- Architecture: JK-Nets use jump connections to expose intermediate neighborhood representations at the final layer for node-specific selection or combination.This lets the effective neighborhood size vary by node rather than being fixed solely by network depth.
- Layer aggregation: The framework supports concatenation, max-pooling, and LSTM-attention as layer-aggregation mechanisms.Concatenation shares weights across nodes, max-pooling selects feature-wise maxima without additional parameters, and LSTM-attention assigns node-specific importance scores to layers.
- Layer aggregation: LSTM-attention can perform well on large complex graphs but may overfit on small graphs because of its relatively higher complexity.This limitation is stated specifically for the LSTM-attention implementation.
- Adaptive locality: Layer-wise max-pooling implicitly learns node-adaptive influence locality by determining feature importance across learned representations from different ranges.The design evaluates layer features before determining their relative importance, rather than fixing identical combination weights for all nodes.
- Adaptive locality: Under equal activation probabilities for same-length paths, a k-layer JK-Net with max-pooling has an expected influence equivalent to a mixture of 0- through k-step random-walk distributions.The mixture coefficients depend on the layer features.
- Empirical behavior: On citation-network visualizations, JK-Net keeps influence within small communities in tree-like regions and emphasizes the node itself or a reasonable neighborhood near hubs.The comparison indicates that fixed random-walk influence can spread too far in tree-like structures or dilute the importance of hub-affiliated nodes’ own features.
5. Other Related Work
Spectral graph convolution uses graph-Laplacian eigenvectors as Fourier atoms but requires the Laplacian in advance. Consequently, it cannot generalize to unseen graphs; the supplied dataset-statistics passage only identifies a table.
- Spectral methods: Spectral graph convolutional networks define graph convolution using graph-Laplacian eigenvectors as Fourier atoms.The approach is contrasted with spatial neighborhood-aggregation methods.
- Spectral methods: Because the graph Laplacian must be known in advance, spectral methods cannot generalize to unseen graphs.This is identified as a major drawback relative to spatial approaches.
6. Experiments
The experiments evaluate JK-Nets across citation, Reddit, and PPI benchmarks using GCN, GraphSAGE, and GAT base models. JK-Nets generally outperform the baselines, with adaptive aggregation especially effective on structurally diverse PPI graphs.
- Datasets and settings: JK-Nets are evaluated on Citeseer, Cora, Reddit, and PPI using transductive settings for the first three datasets and an inductive setting for PPI.The tasks are paper classification, Reddit community prediction, and protein-function classification.
- Citeseer & Cora: JK-Nets outperform both GCN and GAT baselines in prediction accuracy on Citeseer and Cora, although no variant wins consistently across datasets.Performance varies slightly across datasets.
- Citeseer & Cora: On Cora, JK-Nets achieve their best performance with 6 layers, whereas GCN and GAT perform best with only 2 or 3 layers.The results suggest that combining global and local information can improve over relying primarily on local information.
- Reddit: On Reddit, JK-Net with GraphSAGE MaxPool node aggregation and concatenation layer aggregation achieves the best Micro-F1 among GraphSAGE and JK-Net variants.The original GraphSAGE baseline has Micro-F1 0.95, while JK-Net reduces the error by 30%.
- PPI: On PPI, JK-Nets with LSTM-attention outperform GraphSAGE, GAT, and concatenation-based JK-Nets, with JK-LSTM exceeding GraphSAGE by 0.128 micro-F1 after 30 epochs.The authors identify structure-aware adaptive models as especially beneficial for complex graphs with diverse structures.
7. Conclusion
The paper proposes JK-Nets, an aggregation scheme that adapts neighborhood ranges to individual nodes. It argues that this flexibility can improve representations on graphs whose subgraphs have diverse local structures.
- 7. Conclusion: JK-Nets adapt neighborhood ranges to individual nodes instead of using a fixed number of neighborhood aggregations.The approach is motivated by differences in neighborhood information ranges for graph node embeddings.
- 7. Conclusion: The proposed flexibility is particularly useful for graphs with diverse local subgraph structures that fixed aggregation depths may not capture well.The authors identify exploring other layer aggregators and combinations of layer-wise and node-wise aggregators as future work.
A. Proof for Theorem 1
The proof connects a node’s expected influence under neighborhood aggregation to random-walk probabilities. It does so by decomposing derivatives into computation paths and normalizing the resulting influence scores.
- A. Proof for Theorem 1: For paths of length k, the proof relates the derivative contributions to k-step paths from node x to node y.The path formulation uses the fact that random-walk probability at y is obtained by summing probabilities over all length-k paths from x to y.
- A. Proof for Theorem 1: The proof represents derivatives through paths in a directed acyclic computation graph, tracking path weights and ReLU activation indicators.The activation variables are assumed to be Bernoulli with a common success probability.
- A. Proof for Theorem 1: After normalization, the expected influence score at node z equals the k-step random-walk probability from x to z multiplied by a factor shared across z.Random-walk probabilities from x sum to 1, enabling the normalization step.
B. Proof for Proposition 1
The proposition characterizes layer-wise max-pooling in a k-layer JK-Net as an expected mixture of random-walk distributions across neighborhood ranges. The mixture coefficients depend on the layer features.
- B. Proof for Proposition 1: The proof expresses the layer-aggregated feature at node x in terms of influence contributions associated with different layer depths.The fraction of entries selected by max-pooling appears in the influence expression.
- B. Proof for Proposition 1: The resulting influence term is equivalent to the l-step random-walk probability at node y starting from x, up to the stated normalization factors.This connects max-pooling’s layer selection to locality at multiple neighborhood ranges.
C. Visualization Results
The visualizations compare GCN influence distributions with random-walk distributions and show why fixed neighborhood ranges can help or hurt predictions across subgraph structures.
- Figure 6 uses node colors to represent probability masses from influence or random-walk distributions.Shallower colors indicate smaller probability masses; probabilities above 0.2 share one color for visualization.
- Table 6 compares influence distributions from GCN and residual GCN models with corresponding random-walk distributions.The visualizations use GCN models with 2, 4, and 6 layers and label influenced nodes and walk starting nodes with squares.
- The random-walk comparisons use a lazy factor of 0.4, adding a 0.4 probability of remaining at the current node each step.GCN and random-walk colors can differ for high-degree nodes because the GCN weighting assigns them less weight.
- Figures 7 and 8 show real subgraphs where 2-layer GCNs succeed or fail relative to 3- and 4-layer GCNs.These examples complement the paper’s theoretical analysis of effective neighborhood ranges.
- Fixed 2-hop ranges can miss useful 3- or 4-hop information, while fixed 3- or 4-hop ranges can include irrelevant or cross-community neighbors.Rapid random-walk expansion can make deeper models incorporate nearly every node, whereas bridge-like structures can connect to another community.