Graph Transformers on PATTERN

From paper to experiment

The Transformer was built for sequences, not graphs. Adapting its attention to graph-structured data takes a handful of specific design choices, and it is not obvious which of them actually carry the result. This project starts from a paper that proposes four such choices: restrict attention to each node’s graph neighbours, use the graph’s own structure to give every node a position, swap one normalization layer for another, and (where available) carry edge features. It puts the first three to the test.^[01]01Dwivedi and Bresson (2021), “A Generalization of Transformer Networks to Graphs,” DLG-AAAI 2021.

The three testable choices are reimplemented from scratch in modern PyTorch Geometric and trained on PATTERN, a graph benchmark, changing one choice at a time. PATTERN has no edge features, so the fourth choice is left aside. The point is not a full reproduction but a clear read on how much each choice matters.

The page builds up the method (how attention becomes message passing, and how a graph gives each node a position), then compares the variants and looks inside one trained model. Interactive figures let you explore the attention mask, the positions, and the results directly.

Attention becomes message passing

A standard Transformer lets every token look at every other token and decide how much to borrow from each.^[02]02Vaswani et al. (2017), “Attention Is All You Need,” NeurIPS 2017. A graph already carries that information about which elements are related: its edges. Graph neural networks build directly on this: each node repeatedly gathers messages from the neighbours it is connected to and updates itself from them, so information flows along the graph’s own structure.^[03]03Sanchez-Lengeling et al. (2021), “A Gentle Introduction to Graph Neural Networks,” Distill. ^[04]04Daigavane, Ravindran, and Aggarwal (2021), “Understanding Convolutions on Graphs,” Distill.

A Graph Transformer sits between the two ideas. It keeps attention, weighing how much each neighbour matters rather than treating them all equally, but instead of attending to the whole graph, each node attends only to its neighbours. For a node $i$, the model first turns every node’s features into three vectors: a query, a key, and a value.

\[ q_i^k = Q^k h_i, \qquad k_j^k = K^k h_j, \qquad v_j^k = V^k h_j. \]

The query of node $i$ is compared with the key of each neighbour $j$ to score how relevant that neighbour is, and the scores are turned into weights that add up to one across the neighbourhood $\mathcal N(i)$:

\[ s_{ij}^k = \frac{q_i^k\cdot k_j^k}{\sqrt{d_k}}, \qquad \alpha_{ij}^k = \frac{\exp(s_{ij}^k)}{\sum_{r\in\mathcal N(i)}\exp(s_{ir}^k)}. \]

Finally, node $i$ collects its neighbours’ values in those proportions, across all $H$ attention heads, to form its update:

\[ \widehat h_i = O\mathop{\Vert}_{k=1}^{H} \sum_{j\in\mathcal N(i)}\alpha_{ij}^k v_j^k. \]

The one thing that changes between the two attention variants tested later is that neighbourhood $\mathcal N(i)$. Sparse attention uses a node’s real graph neighbours; full attention lets every node attend to every other node in the same graph, ignoring the edges entirely. The figure below shows both; switch the mask and hover a node or an edge to see which messages reach the query node.

Restricting attention to the edges does two useful things. It saves computation by never comparing unrelated node pairs, and, more importantly, it builds in the assumption that a node’s immediate neighbours are the ones that matter. Full attention throws that assumption away while keeping everything else the same, which is exactly what makes the later comparison between them informative.

Giving each node a position

Attention has no built-in sense of where a node sits. Renumbering the nodes just renumbers the outputs, which is the right behaviour for a graph, but it also means two nodes with identical features look identical to the model even when their place in the graph is exactly what the task depends on. The nodes need coordinates that come from the graph’s structure rather than from an arbitrary numbering.

Those coordinates come from the graph Laplacian, a matrix built from which nodes connect to which:

\[ L = I-D^{-1/2}AD^{-1/2}=U\Lambda U^\top, \]

where $A$ records the edges and $D$ counts the connections at each node. Its eigenvectors, sometimes called the graph’s modes by analogy with the vibration modes of a drum skin, describe patterns that vary smoothly across well-connected regions of the graph. The smoothest non-trivial one, the Fiedler vector, changes gradually from one side of the graph to the other; higher modes change more quickly. Because these patterns follow the connectivity, they make natural position coordinates.^[05]05Belkin and Niyogi (2003), “Laplacian Eigenmaps for Dimensionality Reduction and Data Representation,” Neural Computation.

Two of these eigenvectors are computed for every PATTERN graph and added to each node’s input through a learned linear map:

\[ h_i^{(0)} = W_x x_i + W_{\lambda}\lambda_i+b_{\lambda}. \]

These coordinates have one quirk: an eigenvector and its exact negative are equally valid, so the solver’s choice of sign is arbitrary. During training each one is randomly flipped for each graph, so the model learns to rely on the pattern of the coordinates rather than on a sign that carries no meaning.

The model in brief

Each layer runs the neighbour-attention above, then passes the result through a small feed-forward network, with a residual connection and a normalization step around each. The default follows the paper and normalizes with BatchNorm (across all nodes in a batch), while the normalization ablation swaps in LayerNorm, which normalizes each node across its own features instead.^[06]06Ioffe and Szegedy (2015), “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift,” ICML 2015. ^[07]07Ba, Kiros, and Hinton (2016), “Layer Normalization.” The two use the same number of parameters but different statistics, which turns out to matter.

Ten such layers, eight attention heads, a hidden width of 80, and a small per-node classifier make up the network, about 523,000 trainable parameters in all. It reuses the paper’s architecture and PATTERN settings but is written from scratch in current PyTorch Geometric, so each ablation is a small, self-contained change.

The benchmark and the four variants

PATTERN is a synthetic benchmark: each example hides a small 20-node pattern inside a larger random graph, and the job is to label, for every node, whether it belongs to that pattern. The node features are random, so the only way to tell pattern nodes apart is from the graph structure, which is exactly what makes it a good test of these design choices. It comes as a fixed set of 10,000 training, 2,000 validation, and 2,000 test graphs.^[08]08Dwivedi et al. (2023), “Benchmarking Graph Neural Networks,” JMLR 24(43).

The graphs average about 119 nodes, and only around 18% of nodes belong to the pattern. Because of that imbalance, results are reported as balanced accuracy, the average of how well each of the two classes is found, where always guessing the majority scores 50%, rather than the misleadingly high number plain accuracy would give.

Four variants isolate the three design choices, changing one at a time from the paper’s design:

All four are trained in exactly the same way (the same optimizer, and a schedule that lowers the learning rate when progress stalls and stops once it bottoms out), so any difference in the results comes from the design choice being tested, not the training setup. Each variant is run a few times to see how much the outcome varies from run to run.

Results: the structure does the work

The chart below shows the test balanced accuracy of all four variants. The bar around each point is how much the result moved from run to run, and hovering a point gives its exact numbers, including how much GPU memory it used.

The paper’s design is the strongest, at about 85%. Taking away the node positions costs a few points and, just as tellingly, makes the result much less consistent from run to run: without positions the model cannot reliably tell which nodes form the pattern, so the graph coordinates are clearly pulling their weight. Switching BatchNorm for LayerNorm costs a little more and produces the widest spread of all: some runs come close to the paper’s design while others fall well short, so the lesson is not that LayerNorm fails but that BatchNorm makes training far more repeatable here.

The biggest change by far comes from full attention. Letting every node attend to every other node, instead of just its neighbours, drops accuracy from the mid-80s to around 58%, barely above the 50% floor, and more than doubles the GPU memory. The edges were doing real work: they tell the model where the graph’s local structure is, and removing that guidance while flooding each node with irrelevant messages is far worse than keeping attention on the graph.

Looking at one graph

An averaged score does not show where the model actually goes wrong. The graph below runs the trained paper-design model on one test graph, the smallest in the test set, chosen before looking at its predictions so its 44 nodes stay legible. It is a single illustrative example, not a stand-in for the overall accuracy.

The prediction view makes the averaged number concrete: four pattern nodes are found, seven background nodes are wrongly flagged, and two pattern nodes are missed. A single graph can be easier or harder than the test set as a whole, which is why the chart above, not any one graph, is the basis for comparing variants.

How the variants train

The final numbers hide how differently the variants get there. The chart below traces each run’s validation accuracy over training; the larger marker on each line is the version actually kept: the epoch with the lowest validation loss, which is the one later measured on the test set.

The paper’s design settles early and at a consistent point across runs. The other two are far less predictable: the best moment can land very early or very late from one run to the next, the same instability the results chart showed. The model can still find good solutions, just far less reliably.

Training continues past that best moment, until the learning rate bottoms out, so keeping the best-validation version rather than the final one matters, since the two are often many epochs apart.

Limitations

This study covers one synthetic benchmark and the node-only version of the Graph Transformer. Each variant was run only a few times, so the numbers describe trends rather than precise measurements, and the findings should not be assumed to hold on tasks where long-range connections or edge features matter.

It is also a focused reimplementation rather than an exact reproduction of the original, and the node positions are computed ahead of time, which does not scale to very large graphs.

Conclusion

Isolating the paper’s design choices one at a time makes it clear what each one is worth. On PATTERN, keeping attention on the graph’s edges, adding positions from the Laplacian, and normalizing with BatchNorm gives the strongest and most consistent result of the four variants. Removing the positions or switching the normalization lowers accuracy and makes training less repeatable, and letting every node attend to every other node is far worse while costing more than twice the memory.

The takeaway is narrow but clear: when the labels depend on where a node sits in the graph, the graph’s own connectivity belongs inside the attention step. Positions from the Laplacian help on top of that, and BatchNorm makes the training more repeatable here.