PaperlayerView exact revision on arXiv

Source-linked AI summary

A Generalization of ViT/MLP-Mixer to Graphs

Xiaoxin He, Bryan Hooi, Thomas Laurent, Adam Perold, Yann LeCun, Xavier Bresson

arXiv:2212.13350v2cs.CV

TL;DR

Message-passing GNNs have limitations in long-range modeling and graph-isomorphism expressivity, while higher-order WL-GNNs have high complexity. Graph ViT/MLP-Mixer generalizes ViT/MLP-Mixer to graphs, lifting performance over tested base MP-GNNs across datasets and reporting new state-of-the-art scores on larger molecular graphs.

  • Problem

    Message-passing GNNs have limited long-range modeling and graph-isomorphism expressivity, while higher-order WL-GNNs require high memory and speed complexity.

    Source: §1 · pp. 1–2

  • Method

    Graph ViT/MLP-Mixer generalizes ViT/MLP-Mixer to graphs with a GNN architecture designed for long-range dependency, linear speed and memory complexity, and high graph-isomorphism expressivity.

    Source: §2 · p. 3

  • Results

    Graph ViT/MLP-Mixer lifts performance over all tested base MP-GNNs across datasets and reports new state-of-the-art scores on larger molecular graphs.

    Source: §5.1 · p. 7, §5.2 · p. 8

  • Takeaways & Limitations

    The architecture can be applied to different base MP-GNNs in practice.

    Source: §5.1 · p. 7

    The work focuses on alleviating over-squashing, although poor long-range modeling in deep GNNs can also involve vanishing gradients and limited isomorphism expressivity.

    Source: §1 · p. 2

Abstract

from arXiv · show

Graph Neural Networks (GNNs) have shown great potential in the field of graph representation learning. Standard GNNs define a local message-passing mechanism which propagates information over the whole graph domain by stacking multiple layers. This paradigm suffers from two major limitations, over-squashing and poor long-range dependencies, that can be solved using global attention but significantly increases the computational cost to quadratic complexity. In this work, we propose an alternative approach to overcome these structural limitations by leveraging the ViT/MLP-Mixer architectures introduced in computer vision. We introduce a new class of GNNs, called Graph ViT/MLP-Mixer, that holds three key properties. First, they capture long-range dependency and mitigate the issue of over-squashing as demonstrated on Long Range Graph Benchmark and TreeNeighbourMatch datasets. Second, they offer better speed and memory efficiency with a complexity linear to the number of nodes and edges, surpassing the related Graph Transformer and expressive GNN models. Third, they show high expressivity in terms of graph isomorphism as they can distinguish at least 3-WL non-isomorphic graphs. We test our architecture on 4 simulated datasets and 7 real-world benchmarks, and show highly competitive results on all of them. The source code is available for reproducibility at: \url{https://github.com/XiaoxinHe/Graph-ViT-MLPMixer}.

1 Message-Passing GNNs and the Limitations

Standard message-passing GNNs aggregate local neighborhoods, but have limited graph-isomorphism expressivity and struggle to model long-range dependencies. Existing approaches address these limitations through higher-order WL methods, positional encodings, or graph transformers, each with tradeoffs in complexity or positional ambiguity.

Message-Passing GNNs (MP-GNNs)

  • Standard MP-GNNs aggregate 1-hop neighborhoods and stack L layers to transmit information across L hops, but cannot distinguish some non-isomorphic graphs.Their local message-passing mechanism is used across graph domains including recommendation [1] [2], chemistry [3] [4], and physics [5].Source: §1 · p. 1

Weisfeiler-Leman GNNs

  • Higher-order k-WL GNNs can distinguish broader classes of non-isomorphic graphs, but their O(N^k) memory and speed costs—or reduced O(N^2) and O(N^3) variants—remain costly versus O(E) MP-GNNs.Substructure-based designs, including subgraph isomorphism, routing, cellular WL, and k-hop egonets, seek to reduce WL-GNN costs while retaining expressivity.Source: §1 · pp. 1–2

Graph Positional Encoding (PE)

  • Graph positional encodings can increase MP-GNN representation power at linear complexity, but none provides globally unique, equivariant, and distance-sensitive node positions for arbitrary graphs.For example, Laplacian eigenvectors can flip signs, creating symmetries that networks must learn to handle.Source: §1 · p. 2

Issue of long-range dependencies

  • Over-squashing limits long-range modeling because L message-passing layers propagate information only L hops, compressing exponentially growing receptive fields into fixed-length vectors.For binary trees, the receptive field can grow as O(2^L); vanishing gradients and limited isomorphism expressivity can also impair deep GNNs.Source: §1 · p. 2

Graph Transformers

  • Graph Transformers address over-squashing and limited long-range dependencies, but generally raise complexity from O(E) to O(N^2), creating a computational bottleneck.Examples include Laplacian- or random-walk-based encodings, structural attention, and hybrid MPNN+Transformer architectures.Source: §1 · p. 2

2 Generalizing ViT/MLP-Mixer to Graphs

Graph ViT/MLP-Mixer adapts ViT/MLP-Mixer to graphs to capture long-range dependencies and achieve high graph-isomorphism expressivity while retaining linear speed and memory complexity. The section presents benchmark results and positions the approach as a bridge between computer vision, NLP, and graphs.

2 Generalizing ViT/MLP-Mixer to Graphs

  • MLP-Mixer replaces attention with MLPs applied alternately across image patches or tokens and features, offering an efficient alternative that can perform competitively with vision Transformers.Its input tokens are linearly encoded before the alternating feed-forward layers.Source: §2 · p. 2
  • Standard MP-GNNs are computationally efficient but have limited representation power and long-range dependency, whereas Graph Transformers address those limitations at quadratic computational cost.Source: §2 · p. 3
  • Graph ViT/MLP-Mixer is designed to capture long-range dependencies, maintain linear speed and memory complexity, and achieve high graph-isomorphism expressivity.The authors identify the challenges of generalizing ViT/MLP-Mixer from images to irregular, variable graphs.Source: §2 · p. 3
  • It reports 0.073 MAE on ZINC and 0.7997 ROCAUC on MolHIV, with competitive results across Benchmarking GNNs and OGB benchmarks.Source: §2 · p. 3
  • The model achieves SOTA performance on LRGB, mitigates over-squashing on TreeNeighbourMatch, and reaches 3-WL expressive power on SR25.The cited contribution also describes the LRGB result as low-complexity.Source: §2 · p. 3
  • The approach is presented as a unified architecture linking computer vision, NLP, and graphs to support cross-domain network design.Source: §2 · p. 3

3 Generalization Challenges

Adapting MLP-Mixer to graphs raises challenges in defining and extracting tokens, encoding variable subgraphs, preserving positional consistency, and reducing over-fitting through graph augmentation.

3 Generalization Challenges

  • Graphs lack a canonical grid and vary in node and edge counts, making meaningful, uniform graph-token definition and fast extraction challenging.Unlike image patches, graph patches cannot be uniformly divided across examples or uniquely defined without a canonical graph embedding.Source: §3 · p. 4
  • Graph tokens vary in topology and size, so their encoding must produce fixed-length vectors while remaining invariant to node re-indexing.Graph patches can differ in their numbers of nodes and edges and their connectivity.Source: §3 · p. 4
  • Because graph patches and their nodes are unordered, preserving local and global positional consistency is a challenge.Unlike image patches and pixels, graphs have no unique alignment or ordering for patches and nodes.Source: §3 · p. 4
  • ViT/MLP-Mixer architectures are strong over-fitters [62], while graph augmentation is challenging because node counts, edge counts, and connectivity vary.The cited augmentation and regularization approaches include cropping, random horizontal flipping, RandAugment [68], and mixup [69]; graph augmentation methods are not yet as effective [70].Source: §3 · p. 4

4 Proposed Architecture

Graph ViT/MLP-Mixer predicts graph-level targets by extracting graph patches, encoding them as tokens, mixing token and channel information, and pooling the result.

4.1 Overview

  • The architecture combines patch extraction and embedding, mixer layers, global average pooling, and a classifier head for graph-level prediction.Each mixer layer alternates residual token mixing across patches with channel mixing within each patch.Source: §4.1 · p. 4, Figure 1 · p. 5

4.5 Mixer Layer

  • The graph mixer alternates token mixing over patches and channel mixing within patches, while a graph-based multi-head attention variant can model dependencies from graph topology.The MLP-Mixer layer uses O(P) memory and computation, compared with O(P^2) memory and computation for ViT self-attention.Source: Figure 1 · p. 5, §4.5 · p. 7

4.2 Patch Extraction

  • METIS partitions graphs into clusters that are expanded to overlapping one-hop neighborhoods, preserving between-patch edges while meeting the O(E) extraction-complexity requirement.The extraction criteria also call for an algorithm applicable to arbitrary graphs and patches whose nodes are more closely connected internally than externally.Source: §4.2 · p. 5

4.3 Patch Encoder

  • A graph patch encoder uses a GNN to handle permutation invariance, heterogeneous neighborhoods, variable patch sizes, graph convolution, and graph-isomorphism expressivity.It updates node and edge representations with MP-GNN layers and produces a fixed-size patch embedding through mean pooling and an MLP; patch graphs are small, with mean sizes and diameters of 3.15 and 1.82 for molecular datasets and 17.20 and 3.07 for image datasets.Source: §4.3 · pp. 5–6

4.4 Positional Information

  • The model uses absolute positional encodings for patch nodes and relative positional encodings for graph patches to supply positional information absent from general graphs.It uses RWSE for molecular data and Laplacian eigenvectors for image superpixels, randomly flipping eigenvector signs during training; patch-level relative information comes from the coarsened adjacency matrix.Source: §4.4 · p. 6

4.6 Data augmentation

  • At each epoch, the method drops a small set of edges before METIS partitioning to vary the patches, then extracts patches from the original graph to retain all its nodes and edges.The perturbation is proposed to reduce overfitting associated with MLP-Mixer architectures.Source: §4.6 · p. 7

5 Experiments

Across simulated, small real-world, and large real-world graph benchmarks, Graph ViT/MLP-Mixer achieves competitive results while addressing limitations of message-passing GNNs.

5.1 Comparison with MP-GNNs

  • ViT/Graph MLP-Mixer improves all evaluated base MP-GNNs, with average gains of 0.056 Average Precision on Peptides-func and a 0.028 MAE decrease on Peptides-struct.The base models were augmented with the same type of positional encoding for a fair comparison.Source: §5.1 · p. 7, Table 2 · p. 8

5.2 Comparison with SOTAs

  • Graph ViT/MLP-Mixer sets new SOTA scores of 0.6970 on Peptides-func and 0.2449 on Peptides-struct, and achieves 0.073 on ZINC and 0.7997 on MolHIV.The comparisons include Graph Transformers and expressive GNNs.Source: Table 3 · p. 8, §5.2 · p. 8
  • On datasets with hundreds of nodes, SAN+LapPE and SUN require 9.4× and 43.8× the training time per epoch, and 12.4× and 18.8× the memory, respectively, versus Graph ViT/MLP-Mixer.The paper attributes the scalability advantage to avoiding full attention and per-node inner GNN computation.Source: §5.2 · p. 8

5.3 Graph ViT/MLP-Mixer can mitigate over-squashing

  • Standard MP-GNNs fail to generalize on TreeNeighbourMatch from r = 4, while Graph ViT/MLP-Mixer generalizes through r = 7 by transmitting long-range information directly without squashing it.The task labels a target node using the label of a leaf node with the same degree.Source: §5.3 · pp. 8–9, Figure 2 · p. 9

5.4 Graph ViT/MLP-Mixer can achieve empirical high expressivity

  • Graph ViT/MLP-Mixer achieves perfect accuracy on CSL, EXP, and SR25 where MP-GNNs fail, including SR25 graphs indistinguishable by a 3-WL test.These results are empirical; the paper notes that a direct analytical comparison with Weisfeiler-Lehman tests is challenging.Source: Table 4 · p. 9, §5.4 · p. 9

5.5 Ablation Studies

  • The ablations examine patch extraction and its partitioning, size, and overlap; node and patch positional encodings; augmentation trade-offs; and Graph ViT’s gMHA designs.The paper also provides complexity analysis and discusses limitations.Source: §5.5 · p. 9

6 Conclusion

The paper proposes a GNN architecture to address standard MP-GNNs’ limited expressivity and poor long-range dependency, reporting promising results on benchmark graph datasets.

6 Conclusion

  • The proposed GNN architecture targets standard MP-GNNs’ low expressivity and poor long-range dependency.The conclusion identifies these as the limitations the architecture aims to improve.Source: §6 · p. 9
  • The architecture achieves promising results on several benchmark graph datasets.The conclusion does not specify benchmark names or quantitative results.Source: §6 · p. 9
  • Future work will explore graph-token and vision Transformer-like inductive biases for fundamental node and link prediction tasks.The paper also proposes potentially avoiding specialized libraries such as PyG [76] or DGL [77] by replacing sparse graph-token linear algebra with dense operations.Source: §6 · p. 9

A Related Work

The paper distinguishes Graph ViT/MLP-Mixer from hierarchical graph models through its MLP-based backbone and its dynamic, overlapping graph patches. It presents this architecture as an alternative to Graph Transformers and a possible direction for GNN design.

Hierarchical graph models

  • Coarformer and Exphormer [79] combine MPNNs with Transformers, ANS-GT [80] uses node sampling and hierarchical attention, and NAGphormer [81] attends over node-wise multi-hop token sequences.Source: A Related Work · p. 15

Main differences

  • Unlike hierarchical models using hybrid MPNN+Transformer or Transformer backbones, this work uses an MLP backbone, applies MPNNs to small graph patches, and processes them sequentially.Coarformer and Exphormer combine message passing and Transformer modules in parallel, while SAT and this architecture are sequential.Source: A Related Work · p. 15

Graph coarsening

  • Instead of preprocessing static, non-overlapping patches, this work uses stochastic Metis on-the-fly to create dynamic, overlapping graph patches.Source: A Related Work · p. 15

Summary

  • The authors argue that replacing Graph Transformers with a Graph ViT/MLP-Mixer backbone could open new directions for GNN design.They identify graph clustering, graph embedding, mixer layers, positional encoding, and (pre-)training as possible building blocks to enhance.Source: A Related Work · p. 15

B Datasets Description

The study evaluates Graph MLP-Mixer on synthetic expressivity tests, molecular and image-derived benchmarks, and datasets designed to test long-range reasoning and over-squashing.

B Datasets Description

  • CSL, EXP, and SR25 test graph classification on isomorphism classes, pairs of 1&2-WL-failed graphs, and 3-WL-failed strongly regular graphs, respectively.CSL has 150 graphs across 10 classes; EXP has 600 graph pairs; SR25 has 15 graphs with 25 nodes each.Source: B Datasets Description · p. 15
  • ZINC is a 12K-molecule subset used to regress constrained solubility, with molecules of 9–37 nodes and a predefined 10K/1K/1K train/validation/test split.Nodes represent heavy atoms and edges represent bonds.Source: B Datasets Description · p. 15
  • MNIST and CIFAR10 are 10-class image benchmarks converted into 8-nearest-neighbor graphs of SLIC superpixels, with graph sizes of 40–75 and 85–150 nodes, respectively.Their standard train/validation/test splits follow the original image datasets; the authors describe them as sanity checks.Source: B Datasets Description · pp. 15–16
  • MolTOX21 and MolHIV are molecular property prediction datasets with atom-and-bond graphs, 9-dimensional node features, and predefined scaffold splits.Their train/validation/test splits are 6K/0.78K/0.78K for MolTOX21 and 32K/4K/4K for MolHIV.Source: B Datasets Description · p. 16
  • Peptides-func and Peptides-struct share graphs from 15,535 peptides and differ in task; their large size and diameter make them suitable for long-range interaction reasoning benchmarks.The graphs total 2.3 million nodes and average 150.94 nodes and 56.99 diameter per graph.Source: B Datasets Description · p. 16
  • TreeNeighbourMatch tests over-squashing using binary trees of depth r, where predicting a target-node label requires propagating information from all leaves.The problem radius r controls the intensity of over-squashing.Source: B Datasets Description · p. 16

C Experiment Details … F Study on Different Designs of Graph-Based MHA

The experiments show that patch design and attention choice materially affect performance, while positional encodings improve expressivity but do not consistently improve generalization. The study also reports the evaluation protocols and implementation settings used for its comparisons.

C Experiment Details

  • Experiments average four seeds at the best-validation epoch; Adam fluctuations on OGB molecular datasets motivate using ASAM, and simulation and real-world datasets follow their stated evaluation protocols.The simulation protocols include stratified five-fold cross-validation for CSL and EXP and whole-dataset training and validation for SR25; real-world benchmarks use standard train/validation/test splits and specified parameter budgets.Source: C Experiment Details · pp. 16–17
  • The default setup uses 128 hidden units and four layers, with two Mixer layers for LRGB; baseline models receive the same node positional encodings for fair comparison.Speed and memory comparisons use the same machine and batch size 128, halving the batch size when needed to avoid out-of-memory errors.Source: C Experiment Details · p. 18

D.1 Effect of Patch Extraction

  • Patch extraction is critical to performance; the authors attribute its benefit to METIS partitioning and 1-hop extension capturing important local graph information.The ablation treats each individual node as a patch when patch extraction is removed.Source: Table 9 · p. 18, §D.1 · p. 18

D.2 Effect of Graph Partition Algorithm

  • METIS partitioning consistently outperforms random partitioning, especially on larger graphs, although random partitions still yield reasonable results.The authors attribute the advantage to patches whose nodes and edges share similar semantic or informational content.Source: Table 10 · p. 19, §D.2 · p. 19

D.3 Effect of Number of Patches; D.5 Patch Size and WL Expressivity

  • Performance rises with patch count before flattening around 24 patches, and the study finds that greater patch counts increase WL expressivity, especially relevant to larger graphs.The model uses 32 patches by default; because the available 2-WL/3-WL datasets contain small graphs, the need for more patches on larger graphs is proposed rather than directly tested.Source: Figure 4 · p. 19, §D.3 · p. 19, Table 11 · p. 20

D.4 Effect of Patch Overlapping

  • Overlapping patches improve performance over non-overlapping patches, while larger k-hop extensions eventually flatten or reduce performance; the default is k = 1.The reported flattening or decline occurs at k = 2 for ZINC and k = 3 for Peptides-func.Source: Figure 5 · p. 19, §D.4 · p. 19

E.1 Effect of Positional Encoding; E.2 Positional Encoding and Patch Size

  • Positional encodings may increase GNN expressivity without improving generalization, while increasing patch count helps with or without them and node PE helps more than patch PE.Using both node and patch PE is generally more helpful at larger patch counts.Source: Table 12 · p. 20, §E.1 · p. 20, Table 13 · p. 20

F Study on Different Designs of Graph-Based MHA

G Effect of Data Augmentation

Data augmentation improves performance and adds only a small amount of training time by regenerating graph patches during training.

G Effect of Data Augmentation

H Long Range Graph Benchmark · I Mitigating Oversquashing in TreeNeighbourMatch · J Complexity Analysis

The sections report Graph MLP-Mixer’s LRGB benchmark results and its generalization to TreeNeighbourMatch depth 7, where standard MP-GNNs fail from depth 4. They describe degree-based token mixing as the mechanism for matching labels and give patch and mixer complexity bounds.

H Long Range Graph Benchmark

  • SUN scores 0.6730 on Peptides-func and 0.2498 on Peptides-struct, similar to Graph MLP-Mixer, but requires 44x memory and 19x training time.Source: Table 17 · p. 22, H Long Range Graph Benchmark · p. 22
  • Peptides-func and Peptides-struct contain 15,535 graphs and 2.3 million nodes, averaging 151 nodes per graph with mean diameter 57, making them suited to test long-range models.Source: H Long Range Graph Benchmark · p. 22
  • Graph MLP-Mixer sets SOTA at 0.6970 on Peptides-fun and 0.2449 on Peptides-struct, beating MP-GNNs by average 0.056 Average Precision on Peptides-func and 0.028 MAE decrease on Peptides-struct.Source: H Long Range Graph Benchmark · p. 22

I Mitigating Oversquashing in TreeNeighbourMatch

  • Standard MP-GNNs fail to generalize on TreeNeighbourMatch from r = 4, whereas Graph MLP-Mixer generalizes through r = 7.The standard models tested are GCN, GGCN, GAT, and GIN.Source: I Mitigating Oversquashing in… · p. 22
  • Unlike MP-GNNs’ graph-wide compression [44], Graph MLP-Mixer embeds each node’s degree and label per patch, compares degrees with the target via Token Mixer, then outputs the matching node’s label.The Patch Encoder extracts the degree and alphabetical label; the Fully Connected layer combines the label with the degree-match indicator.Source: I Mitigating Oversquashing in… · pp. 22–23

J Complexity Analysis

  • METIS patch extraction runs in O(E); patch embedding adds O(N' + E') runtime and memory over a base O(N + E) GNN, while mixer layers have O(P) complexity.The patch embedding module introduces a constant overhead over the base GNN.Source: J Complexity Analysis · p. 23

K Limitations

The paper identifies limitations in patch-count selection, the empirical basis of its expressivity results, and untested large-scale training and pre-training settings.

K Limitations

  • Patch counts must be chosen separately across datasets, and using a fixed count for variable-size graphs changes graph resolution and may affect performance.The model uses Metis clustering, making the number of patches an arbitrary choice.Source: K Limitations · p. 23
  • Graph MLP-Mixer’s WL-expressivity results are empirical; theoretical analysis at higher WL degrees remains valuable but non-trivial.Source: K Limitations · p. 23
  • Large-scale supervised training on MalNet [94], PascalVOC-SP, COCO-SP [66], and PCQM4Mv2 [65], and pre-training on small- and large-graph datasets, remain unstudied.The paper leaves these experiments for future work.Source: K Limitations · p. 23
Loading 2212.13350v2…