PaperlayerView exact revision on arXiv

Source-linked AI summary

Pure Transformers are Powerful Graph Learners

Jinwoo Kim, Tien Dat Nguyen, Seonwoo Min, Sungjun Cho, Moontae Lee, Honglak Lee, Seunghoon Hong

arXiv:2207.02505v2cs.LGcs.AI

TL;DR

Graph learning often relies on graph-specific architectural bias, but the paper asks whether a standard Transformer can learn graphs from tokenized nodes and edges instead. It introduces TokenGT with node and type identifiers, proves strong expressiveness guarantees, and reports better performance than GNN baselines on PCQM4Mv2 while remaining competitive with modified Transformers.

  • Problem

    Graph Transformers typically require graph-specific modifications because global self-attention does not directly reflect graph structure.

  • Method

    TokenGT treats nodes and edges as independent tokens, augments them with node and type identifiers, and processes them with a standard Transformer.

  • Results

    TokenGT is at least as expressive as k-IGN and k-WL, more expressive than message-passing GNNs, and performs significantly better than GNNs while competing with graph-specific Transformers.

  • Takeaways & Limitations

    Simple token-wise embeddings can provide a standard Transformer with a powerful graph-learning mechanism without graph-specific architectural components.

  • Takeaways & Limitations

    Self-attention over n + m graph tokens has O((n + m)^2) cost, and the method remains slightly below state-of-the-art performance.

Abstract

from arXiv · show

We show that standard Transformers without graph-specific modifications can lead to promising results in graph learning both in theory and practice. Given a graph, we simply treat all nodes and edges as independent tokens, augment them with token embeddings, and feed them to a Transformer. With an appropriate choice of token embeddings, we prove that this approach is theoretically at least as expressive as an invariant graph network (2-IGN) composed of equivariant linear layers, which is already more expressive than all message-passing Graph Neural Networks (GNN). When trained on a large-scale graph dataset (PCQM4Mv2), our method coined Tokenized Graph Transformer (TokenGT) achieves significantly better results compared to GNN baselines and competitive results compared to Transformer variants with sophisticated graph-specific inductive bias. Our implementation is available at https://github.com/jw9730/tokengt.

1 Introduction

TokenGT applies a standard Transformer directly to graphs by representing nodes and edges as tokens with suitable embeddings. The paper argues this simple design is theoretically powerful and performs strongly against GNN and graph-Transformer baselines.

  • Graph Transformers commonly add local attention, message passing, or edge-aware attention biases because global self-attention alone does not reflect graph structure.
  • TokenGT instead treats every node and edge as an independent token, augments them with token-wise embeddings, and feeds them to a standard Transformer.
  • Appropriate node and type identifiers let self-attention approximate permutation equivariant graph operations, giving expressiveness at least matching 2-IGN and 2-WL.
  • On PCQM4Mv2, TokenGT performs significantly better than GNN baselines and competitively with Transformer variants using stronger graph-specific modifications.

2 Tokenized Graph Transformer (TokenGT)

TokenGT represents graph nodes and edges as independent tokens and supplies connectivity and token-type information through embeddings. A standard Transformer then produces graph representations with minimal graph-specific architectural bias.

  • TokenGT constructs a token sequence from all graph nodes and edges, using node identifiers and type identifiers before standard Transformer encoding.
  • Orthonormal node identifiers represent connectivity by assigning each node a distinct vector and attaching endpoint vectors to edge tokens.
  • Comparing endpoint and node identifiers through attention reveals incidence, allowing the Transformer to emphasize incident token pairs when needed.
  • Trainable type identifiers distinguish node tokens from edge tokens, enabling attention heads to select one token type and ignore the other.
  • For graph-level prediction, a trainable [graph] token supplies the graph representation used by a linear prediction head.
  • TokenGT uses less built-in graph inductive bias than GNNs, so it must learn to interpret token embeddings from data.

3 Theoretical Analysis

The theory connects self-attention with permutation-equivariant linear graph layers by using generalized node and type identifiers. It proves that suitably augmented Transformers can match k-IGN expressiveness and therefore exceed message-passing GNNs in expressiveness.

  • 3.1 Preliminary: Permutation Symmetry and Invariant Graph Networks: The analysis represents graphs and hypergraphs as order-k tensors whose features remain invariant under node renumbering.
  • 3.1 Preliminary: Permutation Symmetry and Invariant Graph Networks: k-IGNs use invariant and equivariant linear layers that respect permutation symmetry, with 2-IGN already more expressive than message-passing GNNs.
  • 3.2 Can Self-Attention Approximate Equivariant Basis?: Self-attention can approximate equivariant basis tensors because its weighted value aggregation parallels the masked sums defining equivariant linear layers.
  • 3.2 Can Self-Attention Approximate Equivariant Basis?: Orthonormal identifiers make attention scores encode desired basis patterns, such as identity attention, when query and key projections emphasize the auxiliary identifiers.
  • 3.3 Pure Transformers are Powerful Graph Learners: A Transformer layer with bell(2k) self-attention heads can approximate any order-k equivariant linear layer to arbitrary precision.
  • 3.3 Pure Transformers are Powerful Graph Learners: A Transformer with augmented node and type identifiers can approximate a k-IGN and is therefore at least as powerful as k-WL and more expressive than message-passing GNNs.

4 Related Work

Prior work has established Transformers’ broad expressive power and explored their connection to equivariance, while graph Transformers add architectural mechanisms to handle connectivity and node-ordering challenges.

  • Transformers’ generality has motivated efforts to minimize domain-specific inductive bias while relying on expressive power and scale.
  • Theory has shown that Transformers can model some equivariant functions, including set-equivariant layers and image convolutions with relative positional encodings.
  • Graph Transformers must address edge connectivity and the absence of canonical node ordering, which complicate direct adoption of standard positional encodings.
  • Existing graph-Transformer methods restrict attention locally, combine global attention with message passing, or inject edge information through attention bias.

5 Experiments

Experiments evaluate TokenGT’s equivariant-basis approximation, attention behavior, and large-scale graph-learning performance. Results show that identifier design is essential, while TokenGT performs strongly on PCQM4Mv2.

  • Experiments: TokenGT experiments combine a synthetic equivariant-basis test with PCQM4Mv2 regression and additional large-graph node-classification evaluations.
  • Approximating Second-Order Equivariant Basis: Accurate equivariant-basis approximation occurs only when orthonormal node identifiers and type identifiers are both supplied.
  • Approximating Second-Order Equivariant Basis: Laplacian eigenvector identifiers often outperform ORF identifiers and transfer learned pseudo-equivariant attention structure to unseen graphs near perfectly.
  • Attention Analysis: Attention distance is measured across heads and network depth using mean token distance in graph hops, with training-set analysis showing similar trends.
  • Approximating Second-Order Equivariant Basis: Non-orthogonal random embeddings and independently sampled first-order ORF embeddings produce inaccurate approximation, emphasizing orthogonality and incidence-aware token construction.
  • Large-Scale Graph Learning: 0.0962 MAE: TokenGT (ORF) outperforms all GNN baselines, while 0.0910 MAE: TokenGT (Lap) is competitive with graph-modified Transformers.
  • Large-Scale Graph Learning: 0.0935 MAE: TokenGT (Lap) + Performer achieves the best performance among O(n + m) models through kernelized pure self-attention.

6 Conclusion

The paper concludes that standard Transformers with suitable token embeddings can be effective graph learners in theory and practice. It also identifies quadratic cost, slightly lower state-of-the-art performance, and interpretability as open challenges.

  • Conclusion: TokenGT is at least as expressive as k-IGN and k-WL, making it more expressive than all message-passing GNNs.
  • Conclusion: Simple orthonormal node identifiers combined with trainable type identifiers suffice for the supported theoretical construction and synthetic verification.
  • Conclusion: On PCQM4Mv2, TokenGT performs significantly better than GNNs and competitively with Transformers using strong graph-specific components.
  • Limitations and Future Work: O((n + m)2) self-attention cost remains a challenge, although kernelization reduces the cost to O(n + m).
  • Limitations and Future Work: Future work includes efficient Transformer variants, scaling and architectural improvements, self-supervision, and interpretability research for graph self-attention.
  • Future Directions: Viewing graphs as (n + m) tokens motivates directions including autoregressive decoding, in-context learning, prompting, and multimodal learning.

A.1.1 Extended Preliminary (Cont. from Section 3.1)

The preliminary section formalizes multihead self-attention, Transformers, equivalence classes, basis tensors, and invariant graph networks. These definitions parameterize equivariant linear layers through permutation-based equality patterns.

  • Transformer Definitions: A multihead self-attention layer maps R^n×d to R^n×d, and a Transformer layer combines it with a tokenwise MLP.
  • Equivalence Classes and Basis Tensors: An order-l equivalence class groups multi-indices related by a node permutation, preserving their permutation-invariant equality pattern.
  • Equivalence Classes and Basis Tensors: Each equivalence class defines a binary basis tensor, and equivariant linear layers use these tensors with weights and biases indexed by equivalence classes.
  • Equivalence Classes and Basis Tensors: For order l, bell(l) equivalence classes and basis tensors exist independently of graph size n.
  • Examples: For first-order layers, the two second-order classes correspond to I and 11⊤−I, representing equal and unequal index pairs.
  • Examples: For graph layers with k = l = 2, the construction has bell(4) = 15 weights and bell(2) = 2 biases.

A.1.2 Proof of Lemma 1 (Section 3.3)

The proof decomposes membership in an equivalence class into simpler tests on the two multi-indices and their pairwise equality patterns. Node and type identifiers encode these tests in self-attention, whose coefficients can approximate each equivariant basis tensor arbitrarily accurately up to normalization.

  • Attention construction: Self-attention implements the inclusion test by expressing these Boolean conditions as query-key dot products followed by softmax.This construction uses simpler tests encoded through node and type identifiers.
  • Equivalence-class decomposition: Membership (i, j) ∈µ is reduced to independent tests i ∈γl, j ∈γk, and pairwise equality conditions between their entries.The equivalence classes γl and γk are represented as graph components, while cross-partition edges encode the required pairwise equalities.
  • Equivalence-class decomposition: The scoring function δ(i, j; µ, ϵ) reaches its maximum exactly when (i, j) ∈µ.Violating any membership or equality condition lowers the score by a positive amount.
  • Attention construction: Node identifiers and type identifiers augment each tensor entry before projection into the self-attention layer.The augmentation contains the original feature, node identifiers for each index, and the identifier of its equivalence class.

A.1.3 Proof of Theorem 1 (Section 3.3)

Theorem 1 extends the basis-tensor approximation from Lemma 1 to a complete equivariant linear layer. Multiple attention heads approximate the layer’s basis components, while value, output, and MLP transformations combine and denormalize them.

  • Theorem statement: A Transformer layer with bell(2k) self-attention heads can approximate any order-k equivariant linear layer Lk→k(X) arbitrarily accurately.Each head is assigned to approximate one normalized basis tensor, and the construction uses the augmented input Xinwin.
  • Attention heads: Each value projection preserves the original input features, which are then routed through output projections associated with the layer’s basis-tensor weights.The output projections place each head’s contribution into separate channel blocks.
  • Output construction: The elementwise MLP denormalizes head-wise attention outputs and adds the bias terms required by the equivariant linear layer.Its auxiliary functions recover the relevant equivalence class from the type identifier and compute denormalization factors and biases.
  • Output construction: A final projection selects the first d output channels, yielding the target equivariant layer approximation.The remaining channels serve as auxiliary computation channels and are discarded by the output projection.

A.1.4 Proof of Theorem 2 (Section 3.3)

Theorem 2 composes the Transformer-layer construction across multiple layers and completes it with sum-pooling and an MLP. This composition can approximate an k-IGN arbitrarily accurately.

  • Theorem statement: A T-layer Transformer followed by sum-pooling and an MLP can approximate any k-IGN Fk(X) arbitrarily accurately.The construction uses Transformer layers to implement equivariant layers before applying the invariant output stage.
  • Layer composition: The first T−1 Transformer layers are stacked to approximate successive equivariant layers while updating the first d channels.Elementwise MLPs absorb the activation functions between these layers.
  • Layer composition: The final Transformer layer copies each output entry into reserved channels indexed by its equivalence class.These reserved channels preserve the components needed by the subsequent invariant linear layer.
  • Invariant output: Sum-pooling aggregates the equivariant output over all indices before the final MLP approximates the invariant linear map and bias.The pooled representation is transformed into the graph-level output required by the k-IGN.
  • Invariant output: The resulting composition MLP◦sumpool◦T(T)◦...◦T(1)(X′) approximates Fk(X) to arbitrary precision.This follows from combining the stacked equivariant approximations with the final invariant approximation.

A.2 Additional Discussion on Linear Attention for Graph Transformers (Section 4)

Graph-specific attention biases can block efficient linear-attention methods because they require explicit pairwise attention matrices. TokenGT avoids this incompatibility and, in experiments, supports efficient attention while retaining strong graph-learning performance.

  • Graphormer injects graph structure through an attention bias matrix because node-only self-attention cannot recognize edge connectivity.
  • Kernelized attention reduces time and memory costs to O(n) by factorizing the exponential attention kernel into pairwise feature products.
  • Graphormer’s graph-dependent bias must be added before exponentiation, requiring explicit computation of the full attention matrix.
  • Graphormer and related variations therefore remain bound to O(n^2) and cannot use the described efficient attention methods.
  • Other efficient Transformers, including Set Transformer, LUNA, Linformer, Nyströmformer, Perceiver, and Perceiver-IO, are likewise reported as incompatible with Graphormer.
  • TokenGT variants use Performer attention with O(m + n) cost, enabling larger models and achieving the best performance on all but one reported dataset.
  • TokenGT’s current PCQM4Mv2 performance is slightly below Graphormer and its successors, motivating engineering improvements that may depart from the equivariance theory.
Loading 2207.02505v2…