Research map / Graph learning / GNN architectures
Aligning Transformers with Weisfeiler-Leman
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a tokenization for standard (pure) transformers that aligns them with the k-WL hierarchy, yielding more expressive, feasible pure transformers (k,s)-tuples and LPE/SPE-based structural embeddings, evaluated on PCQM4Mv2, small molecular datasets and BREC.
Problem
Pure transformers (no architectural modification) for graphs lack a feasible hierarchy with provable expressivity beyond 1-WL; prior k-WL-aligned transformers have prohibitive runtime and head counts (e.g., bell(2k) heads, 203 for k=3).
Approach
Tokenize graphs as k-tuple tokens built by concatenating node-level structural embeddings (node- and adjacency-identifying PEs such as LPE or SPE) plus atomic-type embeddings, aligning a standard transformer with k-WL and its (k,s)-tuple variant, needing only 2k heads; pre-train at low order and transfer weights to higher order.
Key results
- KGT(2,1)+LPE reaches 0.0870 validation MAE on PCQM4Mv2, close to Graphormer (0.0864) and better than TokenGT (0.0910).
- On Alchemy (12K), pre-trained KGT(2,1)+LPE achieves 0.101 MAE, beating the best GNN baseline SPE (0.108).
- On BREC, KGT(3,1)+SPE distinguishes 228 pairs versus 217 for GINE+SPE and 79 for Graphormer, while 3-WL reaches 270.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use (2,1) + LPE instead of Graphormer | PCQM4Mv2 validation · PCQM4Mv2 validation | MAE: 0.0864 -> 0.0870 | Hurt |
| Use (2,1) + LPE instead of TokenGT | PCQM4Mv2 validation · PCQM4Mv2 validation | MAE: 0.0910 -> 0.0870 | Helped |
| Use (2,1) + LPE instead of TokenGT | CS · CS | accuracy: 0.903 $\pm$ 0.004 -> 0.924 $\pm$ 0.008 | Helped |
| Use (2,1) + LPE instead of TokenGT | Photo · Photo | accuracy: 0.949 $\pm$ 0.007 -> 0.933 $\pm$ 0.013 | Hurt |
| Pre-train (2,1) + LPE rather than train from scratch (ablation) | (2,1) + LPE · Alchemy (12K) | MAE: 0.124 $\pm$ 0.001 -> 0.101 $\pm$ 0.001 | Helped |
| Pre-train (2,1) + SPE rather than train from scratch (ablation) | (2,1) + SPE · Alchemy (12K) | MAE: 0.112 $\pm$ 0.000 -> 0.103 $\pm$ 0.002 | Helped |
| Transfer pre-trained (2,1) weights to (3,1) instead of fine-tuning (2,1) (ablation) | LPE models · Alchemy (12K) | MAE: 0.101 $\pm$ 0.001 -> 0.114 $\pm$ 0.001 | Hurt |
| Transfer pre-trained (2,1) weights to (3,1) instead of fine-tuning (2,1) (ablation) | SPE models · Alchemy (12K) | MAE: 0.103 $\pm$ 0.002 -> 0.108 $\pm$ 0.001 | Mixed |
| Use SPE rather than LPE for scratch training (ablation) | (2,1) · Alchemy (12K) | MAE: 0.124 $\pm$ 0.001 -> 0.112 $\pm$ 0.000 | Helped |
| Use LPE rather than SPE after pre-training (ablation) | (2,1) · Alchemy (12K) | MAE: 0.103 $\pm$ 0.002 -> 0.101 $\pm$ 0.001 | Helped |
| Increase order from (2,1) to (3,1) with LPE (ablation) | Order-transferred transformer · BREC All | solved graph pairs: 118 -> 159 | Helped |
| Use SPE instead of LPE (ablation) | (2,1) transformer · BREC All | solved graph pairs: 118 -> 226 | Helped |
Surprising findings
- Pre-training improves Alchemy performance but gives no benefit when order-transferring to KGT(3,1), which the authors attribute to overfitting from the added tokens.
- On ToxCast, the pre-trained KGT(2,1) underperforms even GIN without pre-training.
Limitations
- Only provides expressivity lower bounds; KGT(k) remains strictly below 3-WL on BREC.
- Pre-training on PCQM4Mv2 uses only 2 A100 GPUs for 2M steps, and gap to Graphormer on PCQM4Mv2 is not closed.
- Higher-order KGT(3,1) gives no gain over KGT(2,1) in the Alchemy fine-tuning, suggesting overfitting.
Experimental setup
Datasets: PCQM4Mv2 Alchemy (12K) BBBP BACE ClinTox Tox21 ToxCast CS Photo BREC
Benchmarks: PCQM4Mv2 CS Photo Alchemy (12K) BBBP BACE ClinTox Tox21 ToxCast BREC
Metrics: MAE accuracy AUROC BREC count of distinguished pairs
Baselines: Graphormer TokenGT GIN GINE SignNet BasisNet SPE 3-WL
Builds on: TokenGT k-order Transformer SPE Laplacian positional encodings (k,s)-WL
Closest papers by meaning
- Towards Principled Graph Transformers (NeurIPS 2024 · similarity 0.92)
- Weisfeiler-Leman at the margin: When more expressivity matters (ICML 2024 · similarity 0.90)
- On the Stability of Expressive Positional Encodings for Graphs (ICLR 2024 · similarity 0.88)
- Homomorphism Counts as Structural Encodings for Graph Learning (ICLR 2025 · similarity 0.87)
- From Theory to Practice: Rethinking Green and Martin Kernels for Unleashing Graph Transformers (ICML 2025 · similarity 0.87)
- Polynormer: Polynomial-Expressive Graph Transformer in Linear Time (ICLR 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.