Shortest-Path Encodings for Graph Transformers
Published:
Intuition First
In a standard Transformer, every token can attend to every other token — but the attention score is purely based on content similarity. For graphs, two distant nodes might have very similar features yet share no direct structural relationship. SPD encoding adds a “distance penalty” to attention: nodes far apart in the graph should attend less strongly, regardless of feature similarity.
Think of it loosely like gravity — attraction weakens with distance. But the analogy only goes so far: gravity decays as a fixed \(1/d^2\), whereas here the model learns one free scalar per distance value and is under no obligation to make the profile decreasing at all.
Why Shortest Paths?
In Graphormer (Ying et al., 2021), every pair of nodes can attend to each other. But without structural information, the model has no way to know that nodes 1 hop apart should interact more strongly than nodes 10 hops apart.
Shortest-path distance (SPD) encoding injects this directly as an attention bias:
\(\varphi\) is a lookup table of learned scalars, one per distance value \(0, 1, 2, \dots, d_{\max}\), plus one for “no path”. In Graphormer it is learned separately per attention head, so different heads can specialise to different ranges.
What It Captures
- Self (\(\mathrm{dist} = 0\)): its own entry
- Immediate neighbours (\(\mathrm{dist} = 1\)): typically the largest bias
- 2-hop neighbours (\(\mathrm{dist} = 2\)): moderate
- Distant nodes: small or negative
- Disconnected (\(\mathrm{dist} = \infty\)): a dedicated “no path” entry
Two things are worth being precise about. First, nothing constrains \(\varphi\) to decrease with distance — it is a free table of scalars, and the model may well learn a non-monotone profile if the task rewards it. Second, it is a bias on the logits, not a mask: a large content match \(q_i^{\top}k_j\) can still outweigh a strongly negative \(\varphi\), so distant nodes remain reachable. That is precisely what distinguishes it from a hard \(k\)-hop restriction.
All-Pairs Shortest Paths: Computation Cost
On an unweighted graph, BFS from every node gives all-pairs distances in \(O(N(N + \lvert E\rvert))\) time. The output is an \(N \times N\) integer matrix, so memory is \(\Theta(N^2)\) regardless of how sparse the graph is — and in a Transformer that matrix must sit alongside the attention matrix, which is also \(N \times N\).
For small graphs (\(N < 1000\)) this is cheap and computed once per graph as a preprocessing step. At \(N > 10^5\) the memory alone rules it out, well before the time cost does. This is why SPD encoding is used mainly for molecules (\(N\) in the tens) and small protein or structure graphs.
SPD vs LapPE vs RWPE
| Encoding | Type | Captures | Cost |
|---|---|---|---|
| LapPE | Node PE | Low-frequency global position | \(O(k\lvert E\rvert T)\), \(T\) gap-dependent |
| RWPE | Node PE | Local closed-walk structure | \(O(K N \lvert E\rvert)\) time, \(O(N^2)\) memory |
| SPD | Pairwise bias | Exact graph metric between every pair | \(O(N(N+\lvert E\rvert))\) time, \(O(N^2)\) memory |
SPD is a pairwise encoding — a property of a pair, not of a node. It cannot be concatenated to node features; it has to enter through the attention mechanism as a bias. That is a genuine architectural constraint: SPD is unavailable to a plain message-passing GNN, which has no pairwise scoring step to attach it to.
Beyond SPD: Distance Encoding and Anchor Sets
Two related ideas turn pairwise distances into something node-level, and they are often conflated.
Distance Encoding (Li et al., 2020) measures distance from every node to the target node set of the prediction — the node or node pair being classified. For link prediction on \((u,v)\), each node \(w\) is featurised by its distances to \(u\) and to \(v\). The encoding is therefore task-relative, recomputed per query, and this is exactly what gives it provable power beyond 1-WL: it breaks the symmetry between nodes that are structurally equivalent but differently placed relative to the target.
Anchor sets (P-GNN, You et al., 2019) take the other route: sample random anchor sets \(S_1, \dots, S_k\) once, and featurise each node by its distance to each anchor:
This is a fixed node-level encoding, cheaper than full APSP, and it is genuinely positional — it depends on the random anchor draw, so it is not a function of the graph alone and two runs give different encodings.
The distinction matters: DE is task-relative and permutation-equivariant; anchor distances are position-like but carry an arbitrary random choice, much as Laplacian eigenvectors carry an arbitrary sign.
Summary
SPD encoding injects graph metric structure straight into Graph Transformer attention. It is simple, interpretable, and effective on small graphs. Its ceiling is memory, not cleverness: the \(N \times N\) distance matrix scales with the attention matrix, so SPD is viable exactly where full attention already is. For larger graphs, use RWPE or LapPE, which encode structure per node rather than per pair — accepting that neither gives the model an exact distance between two specified nodes.
References
- Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., & Liu, T.-Y. (2021). Do Transformers Really Perform Bad for Graph Representation?. NeurIPS 2021 (Graphormer — introduces SPD and edge-distance encodings).
- Li, P., Wang, Y., Wang, H., & Leskovec, J. (2020). Distance Encoding: Design Provably More Powerful Graph Neural Networks for Structural Representation Learning. NeurIPS 2020.
- You, J., Ying, R., & Leskovec, J. (2019). Position-aware Graph Neural Networks. ICML 2019 (P-GNN — introduces random anchor sets).