Aligning Transformers with Weisfeiler-Leman

TL;DR

This paper enhances transformer architectures aligned with the Weisfeiler-Leman hierarchy, improving their expressivity and practicality for graph tasks, and demonstrates competitive performance on large-scale and molecular datasets.

Contribution

It advances the alignment of transformers with the $k$-WL hierarchy, providing stronger theoretical expressivity results and practical feasibility, along with a framework for studying positional encodings.

Findings

Stronger expressivity results for transformers aligned with $k$-WL.

Competitive performance on PCQM4Mv2 dataset.

Effective fine-tuning on small molecular datasets.

Abstract

Graph neural network architectures aligned with the $k$-dimensional Weisfeiler--Leman ($k$-WL) hierarchy offer theoretically well-understood expressive power. However, these architectures often fail to deliver state-of-the-art predictive performance on real-world graphs, limiting their practical utility. While recent works aligning graph transformer architectures with the $k$-WL hierarchy have shown promising empirical results, employing transformers for higher orders of $k$ remains challenging due to a prohibitive runtime and memory complexity of self-attention as well as impractical architectural assumptions, such as an infeasible number of attention heads. Here, we advance the alignment of transformers with the $k$-WL hierarchy, showing stronger expressivity results for each $k$, making them more feasible in practice. In addition, we develop a theoretical framework that allows the study of established positional encodings such as Laplacian PEs and SPE. We evaluate our transformers on the large-scale PCQM4Mv2 dataset, showing competitive predictive performance with the state-of-the-art and demonstrating strong downstream performance when fine-tuning them on small-scale molecular datasets. Our code is available at https://github.com/luis-mueller/wl-transformers.