PaperlayerView exact revision on arXiv

Source-linked AI summary

NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification

Qitian Wu, Wentao Zhao, Zenan Li, David Wipf, Junchi Yan

arXiv:2306.08385v1cs.LGcs.AI

TL;DR

Observed-graph GNN message passing is limited by heterophily, long-range dependencies, incomplete or absent graphs, while adaptive topology learning is difficult to scale. NodeFormer addresses this with differentiable all-pair message passing over learned latent graphs using a kernelized Gumbel-Softmax operator. Experiments report improvements across diverse tasks, scaling to graphs with up to 2M nodes and reducing resource use versus structure-learning competitors.

  • Problem

    GNN message passing can suffer from heterophily, over-squashing, long-range dependencies, graph incompleteness, and absent input graphs, while adaptive topology learning faces quadratic scalability costs.

  • Method

    NodeFormer uses a kernelized Gumbel-Softmax operator to learn layer-specific latent graphs and propagate messages between arbitrary node pairs with linear algorithmic complexity.

  • Results

    NodeFormer shows significant improvement across diverse node, image, and text classification tasks, scales to graphs with up to 2M nodes, and reduces competitors’ resource consumption.

  • Takeaways & Limitations

    The method supports scalable graph structure learning and all-pair message passing even when input graphs are incomplete or unavailable.

  • Takeaways & Limitations

    The temperature τ creates a trade-off: smaller values require very large kernel dimensions for precision, whereas overly large values make edge weights converge toward 1.

Abstract

from arXiv · show

Graph neural networks have been extensively studied for learning with inter-connected data. Despite this, recent evidence has revealed GNNs' deficiencies related to over-squashing, heterophily, handling long-range dependencies, edge incompleteness and particularly, the absence of graphs altogether. While a plausible solution is to learn new adaptive topology for message passing, issues concerning quadratic complexity hinder simultaneous guarantees for scalability and precision in large networks. In this paper, we introduce a novel all-pair message passing scheme for efficiently propagating node signals between arbitrary nodes, as an important building block for a pioneering Transformer-style network for node classification on large graphs, dubbed as \textsc{NodeFormer}. Specifically, the efficient computation is enabled by a kernerlized Gumbel-Softmax operator that reduces the algorithmic complexity to linearity w.r.t. node numbers for learning latent graph structures from large, potentially fully-connected graphs in a differentiable manner. We also provide accompanying theory as justification for our design. Extensive experiments demonstrate the promising efficacy of the method in various tasks including node classification on graphs (with up to 2M nodes) and graph-enhanced applications (e.g., image classification) where input graphs are missing.

1 Introduction

NodeFormer addresses limitations of fixed observed-graph message passing by learning scalable, layer-specific latent graphs for all-pair propagation. Its kernelized Gumbel-Softmax operator enables linear-complexity structure learning, while experiments report broad efficacy and scaling to graphs with up to 2M nodes.

  • GNNs can struggle with heterophily, over-squashing, long-range dependencies, graph incompleteness, and missing input graphs.
  • Learning adaptive topology is hindered by O(N^2) complexity, which is prohibitive for networks with 10K–1M nodes and can create over-fitting risks.Neighbor sampling, anchor-based approximations, and hashing reduce overhead but may sacrifice precision or remain unable to handle million-node graphs.
  • The kernelized Gumbel-Softmax operator approximates discrete latent structures and reduces per-layer topology-learning complexity from quadratic to linear in node numbers.Its approximation error is bounded by feature dimensions, according to the paper’s theory.
  • NodeFormer performs layer-wise message passing over latent graphs potentially connecting all nodes, with end-to-end differentiable topology sampling conditioned on node features and labels.The model is presented as a Transformer that scales all-pair message passing to large node-classification graphs.
  • Experiments show significant improvement over strong GNNs and state-of-the-art structure-learning methods across node, image, and text classification, scaling to 2M-node graphs.The introduction reports reductions of up to 93.1% in time/space consumption, while the detailed comparison reports separate reductions for training time, inference time, and memory cost.

2 Related Works

Related work spans expressive and scalable GNNs, similarity- or adjacency-based graph structure learning, and node-level prediction settings where arbitrary relationships and scalability matter. NodeFormer is positioned as layer-wise, input-graph-independent structure learning that scales to 2M nodes.

  • Graph Neural Networks: Existing GNNs reweight, sparsify, or sample input structures but cannot learn edges outside the input geometry.This can limit receptive fields to local neighbors and neglect global information.
  • Graph Structure Learning: Graph structure learning methods infer edges through similarity functions or optimize adjacency matrices using methods such as bi-level optimization, Bayesian inference, and projected gradient descent.
  • Graph Structure Learning: NodeFormer enables layer-wise structure learning without requiring input graphs and successfully scales to graphs with 2M nodes.The comparison concerns node-level tasks on graphs connecting all instance nodes.
  • Node-Level v.s. Graph-Level Prediction: The paper focuses on node-level prediction over a single large graph, where scalability is paramount when modeling arbitrary relationships across nodes.

3 NODEFORMER: A Transformer Graph Network at Scale

NODEFORMER learns layer-specific latent graphs and performs all-pair message passing while avoiding quadratic computation. Its kernelized Gumbel-Softmax operator provides linear complexity, differentiable discrete structure learning, and theoretical approximation guarantees.

  • General model: NODEFORMER estimates layer-specific latent graphs from node representations and propagates features over potentially arbitrary node pairs.The model can use an optional input adjacency matrix alongside node features.
  • Key challenges: Learning latent all-pair structures faces scalability and differentiability challenges because quadratic complexity and discrete graph variables hinder optimization.Latent graphs may connect thousands to millions of nodes, making precision and scalability difficult to guarantee simultaneously.
  • Efficient learning discrete structures: Kernelized message passing rewrites softmax attention with random-feature kernel approximations so shared summations yield O(N) computation without constructing an N × N similarity matrix.The random-feature map converts dot-then-exponentiate similarities into inner products, enabling reusable terms across nodes.
  • Differentiable stochastic structure learning: The kernelized operator combines random features with Gumbel-Softmax relaxation to learn discrete latent structures differentiably while retaining sampled-graph message passing with linear complexity.Multiple samples can be averaged in practice; the temperature controls closeness to hard discrete samples.
  • Well-posedness: Theoretical results show approximation error is independent of node number, while convergence to the original categorical distribution requires sufficiently large feature dimension and temperature approaching zero.The theory exposes a trade-off: smaller temperature can require a much larger random-feature dimension for adequate approximation precision.
  • Input structures and learning objective: When input graphs are available, relational bias reinforces observed edges and edge-level regularization uses them to guide adaptive topology learning.The regularizer addresses the increased degrees of freedom created by topology learning relative to the available training labels.

4 Evaluation

NODEFORMER is evaluated across graph-based, graph-free, and large-scale node-classification settings, with results emphasizing accuracy, scalability, and resource efficiency. Ablations and sensitivity analyses examine its stochastic operator, graph guidance components, temperature, and learned structures.

  • Transductive node classification: NODEFORMER achieves the best mean Accuracy/ROC-AUC across four graph datasets and performs especially strongly on the two heterophilic graphs.The compared datasets are Cora, Citeseer, Deezer, and Actor.
  • Larger graph datasets: NODEFORMER scales to Amazon2M with 2 million nodes, using only 4GB GPU memory for training with batch size 100000.The large-scale experiments also include OGB-Proteins with over 0.1 million nodes.
  • Graph-enhanced applications: On graph-free Mini-ImageNet, removing edge regularization and relational bias still yields competitive or superior results, suggesting learned latent structures can outperform k-NN guidance.The graph-free experiments use artificially constructed k-NN graphs for message passing.
  • Time and space consumption: 93.1%, 97.9%, and 75.6% are NODEFORMER’s maximum reductions in training time, inference time, and memory cost versus LDS, respectively.Against IDGL, the corresponding maximum reductions are 61.8%, 80.8%, and 80.6%.
  • Ablations: Removing the Gumbel components causes a performance drop, while edge-level regularization and relational bias each contribute positive effects.The stochastic ablation compares deterministic and temperature-controlled Softmax variants.
  • Sensitivity analysis: Test accuracy generally peaks near temperature 0.4 and depends on feature-map dimension differently at low versus high temperatures.At low temperature, increasing random-feature dimension improves approximation; at high temperature, accuracy can drop even with larger dimension.
  • Visualization and implications: Visualizations show larger weights among same-class nodes, sparse cross-class edges, and dominant nodes that may act as connectivity pivots.The learned latent structures can differ from observed graphs, especially for heterophilic graphs.

5 Why NODEFORMER Improves Downstream Prediction?

NODEFORMER frames latent graph learning probabilistically: estimated graphs define a variational distribution, while supervised propagation evaluates their predictive usefulness. The objective is an ELBO whose optimum aligns the learned topology with the task-conditioned posterior and maximizes the data likelihood.

  • Bayesian perspective: NODEFORMER models prediction through a latent graph, factorizing the predictive distribution into graph estimation and message-passing terms.The simplified analysis uses one latent graph, although the model practically learns layer-specific graphs.
  • Bayesian perspective: The ideal latent topology is characterized by the posterior p(Ã|Y, X, A), which is unknown and intractable to compute directly.The posterior follows from Bayes’ theorem and requires an intractable integration.
  • Variational perspective: The structure estimator is treated as a variational distribution q(Ã|X, A), and the training objective combines supervised prediction with KL regularization.Gumbel-Softmax supplies a surrogate for sampling-based estimates, while edge-level MLE provides the KL regularization term.
  • Theoretical implication: At the optimum, the variational distribution has zero KL divergence from the task-conditioned posterior and log p(Y|X, A) is maximized.This is the content of Proposition 1 under the assumption that q can exploit arbitrary distributions over Ã.
  • Theoretical implication: The proposition links training on labeled data to latent topologies that ideally propagate task-relevant adjacent information and can improve downstream performance.The claim concerns the topology produced by a well-trained NODEFORMER network.

6 Conclusion

The paper concludes that NODEFORMER enables scalable node-level Transformer message passing over layer-specific latent graphs, with linear complexity and effectiveness across diverse settings. Its results support applications involving heterophily, long-range dependencies, large graphs, incomplete graphs, and missing input graphs.

  • Conclusion: NODEFORMER propagates layer-wise node signals between arbitrary pairs beyond the input topology.Its kernelized Gumbel-Softmax operator learns layer-specific latent graphs with linear algorithmic complexity without compromising precision.
  • Conclusion: Experiments on diverse graph datasets and situations verify the method’s effectiveness, scalability, and stability.The conclusion points to settings including heterophily, long-range dependencies, large-scale graphs, graph incompleteness, and absent input graphs.

Checklist

The checklist records affirmative responses on claims and scope, theory assumptions and proofs, reproducibility materials, training details, error bars, compute reporting, and data-ethics documentation. The supplied implementation passages also describe sampled latent-edge propagation and its linear-complexity kernel approximation.

  • Checklist: The authors report that the paper’s main claims accurately reflect its contributions and scope, and that limitations and potential negative societal impacts are discussed.The checklist marks these items as Yes.
  • Checklist: The theoretical results are accompanied by stated assumptions and complete proofs in Appendix B.The checklist marks both theory-reporting items as Yes.
  • Checklist: The authors report public code, dataset information, training details, error bars, and compute-resource information for reproducibility.The checklist points to the supplemental sections and experiment section for these materials.
  • Checklist: The checklist reports citation of existing assets and addresses licensing, new assets, consent, personally identifiable information, offensive content, participant instructions, risks, and compensation.Several human-subject and asset items are marked N/A where applicable.
  • Implementation: Each node samples K latent neighbors per layer, and the resulting discontinuous sampling is relaxed with Gumbel-Softmax for differentiable training.The kernel approximation reduces the one-layer update from O(N^2) to O(N).

A.2 Model Implementation from the Matrix View

NODEFORMER updates all node representations simultaneously through kernelized Gumbel-Softmax message passing. The matrix implementation avoids explicit all-pair similarities, reducing each layer’s complexity to O(N).

  • The matrix view is used in practice to simultaneously update all nodes within one layer.
  • O(N) complexity is achieved by avoiding explicit computation of the all-pair similarity matrix.This replaces the cumbersome quadratic computation required by explicit all-pair similarities.
  • NODEFORMER’s feed-forward process updates node embeddings layer by layer through kernelized Gumbel-Softmax message passing.The implementation presents this process from a matrix perspective.

B.1 Proof for Theorem 1

The proof establishes that the kernelized Gumbel-Softmax preserves the relevant softmax and discrete-selection properties needed by NODEFORMER. It combines random-feature approximation results with Gumbel perturbation arguments and a zero-temperature limit.

  • Random-feature approximation: Positive Random Features provide an unbiased approximation to the softmax kernel with quantified variance.This approximation supplies the kernel-estimation result used later in the theorem proof.
  • Gumbel perturbation: Gumbel perturbation yields the softmax selection probability for pairwise comparisons between perturbed values.The result is expressed as 1 / (1 + exp(−(x_i − x_j))).
  • Zero-temperature limit: As τ → 0, a Concrete or Gumbel-Softmax sample becomes a one-hot vector selecting the index with the largest location parameter.This establishes the discrete-selection behavior of the original operator.
  • Theorem conclusion: The theorem proof shows that kernelized Gumbel-Softmax retains the one-hot limit property of the original Gumbel-Softmax operator.The argument compares kernelized scores formed from random feature maps and applies the approximation bounds.
  • Theorem conclusion: The approximation probability depends on feature-map dimension, temperature, and approximation precision, approaching one as the feature dimension becomes sufficiently large.

D.1 Dataset Information

The experiments cover small homophilous and non-homophilous node-classification graphs, million-node scalability benchmarks, and graph-free image and text applications. Dataset descriptions specify the node semantics, features, labels, and graph availability.

  • Transductive node classification: Cora and Citeseer are homophilous citation networks with document bag-of-words features, citation edges, and document class labels.
  • Transductive node classification: Actor and Deezer provide non-homophilous settings, representing actor co-occurrence and a social network with reported gender labels.Actor nodes use Wikipedia keyword features and five categories; Deezer is the largest of the four smaller datasets at nearly 20K nodes.
  • Large-scale datasets: OGB-Proteins and Amazon2M are million-node, million-edge datasets used to evaluate scalable training.OGB-Proteins is an undirected typed protein graph whose edges carry 8-dimensional features.
  • Graph-enhanced applications: 20News-Groups and Mini-ImageNet lack graph structure and are treated as node collections for graph-enhanced applications.The datasets contain approximately 20,000 documents and selected 84×84 images, respectively, with text or CNN-derived node features.

E More Experiment Results

Additional experiments examine NODEFORMER’s structural components, hyperparameters, and learned attention patterns. The reported studies vary temperature and random-feature dimension and compare latent structures with original input graphs.

  • Ablation studies: The ablation study compares NODEFORMER with and without relational bias and edge-level regularization on four transductive datasets.In the table notation, “rb” denotes relational bias and “reg” denotes edge-level regularization.
  • Hyperparameter impact: Figure 6 studies how temperature τ and random-feature-map dimension m affect performance on Cora.
  • Latent structure visualization: Figure 7 visualizes two NODEFORMER layers’ attention maps alongside the original input graphs for Cora, Citeseer, Deezer, and Actor.Colors represent the weights of the input and latent graph structures.

F Current Limitations, Outlooks and Potential Impacts

NODEFORMER is currently evaluated for node classification, while its broader encoder potential and applications beyond accuracy remain future directions.

  • Current Limitations: Experiments focus on node classification, although NODEFORMER could serve as a flexible graph encoder for graph classification and link prediction.
  • Current Limitations: Robustness and explainability are identified as future targets beyond testing accuracy.
  • Potential Impact: NODEFORMER may serve as a scalable structure-learning encoder for uncovering dependencies, identifying structures, and purifying noisy data.
  • Potential Impact: The kernelized Gumbel-Softmax operator is proposed as a plug-in module for efficient and expressive graph Transformers when graph size limits performance.
Loading 2306.08385v1…