Random Walk Positional Encodings

Intuition First

Imagine a random walker starting at node \(v\). After \(k\) steps, what is the probability it is back at \(v\)? If \(v\) sits in a tight clique there are many short routes home. If \(v\) sits on a long path there are far fewer, and the walker tends to drift.

By stacking the return probabilities at steps \(1, 2, \dots, K\) into a vector, we get a structural fingerprint for \(v\). Odd-length returns are especially informative: they are possible only if \(v\) lies on an odd cycle, so \(P^3[v,v] > 0\) exactly when \(v\) is in a triangle. A bipartite neighbourhood, by contrast, gives zero at every odd step. The fingerprint tells the model what local structure surrounds \(v\) — with none of the sign ambiguity of Laplacian eigenvectors.

The Idea

The random walk transition matrix is \(P = D^{-1}A\) — row-stochastic, so each row sums to 1. Then \(P^k[u,v]\) is the probability that a walk started at \(u\) is at \(v\) after exactly \(k\) steps.

The diagonal entries \(P^k[v,v]\) — the probability of being back at \(v\) after \(k\) steps — encode \(v\)’s local structural role:

• Nodes in dense cliques have high return probabilities, because there are many short closed walks
• Nodes on long paths return less often, because the walk has more room to drift
• A non-zero \(P^k[v,v]\) at odd \(k\) certifies that \(v\) lies on a closed walk of that odd length

Why RWPE Works as a Structural Encoding

Each entry answers: how likely is a \(k\)-step walk from \(v\) to be back at \(v\)? Note that this counts closed walks, not simple cycles — a walk may revisit nodes and backtrack — so the readings are suggestive rather than exact subgraph counts:

• Odd cycles. \(P^3[v,v] > 0\) if and only if \(v\) is in a triangle: the only closed 3-walks are triangles. This is a genuine iff. It does not extend cleanly: \(P^4[v,v] > 0\) for any node with a neighbour, since \(v \to u \to v \to u \to v\) is a closed 4-walk, so a positive \(P^4\) does not imply a 4-cycle.
• Bipartiteness. If \(v\)’s connected component is bipartite, \(P^k[v,v] = 0\) for every odd \(k\).
• Local density. Higher clustering around \(v\) raises the early entries.

RWPE is a structural encoding: it describes \(v\)’s local role. Two nodes with identical local topology get identical RWPE — intended behaviour, and the reason it transfers across graphs.

Comparison with LapPE

Computing RWPE

Only the diagonal is needed, but there is no shortcut that gets it from sparse matrix–vector products alone: recovering \(\mathrm{diag}(P^k)\) exactly requires propagating all \(N\) basis vectors, so the standard implementation carries a dense \(N \times N\) matrix through the recursion \(P^k = P\,P^{k-1}\):

P = D_inv @ A          # row-stochastic: P[v,:] sums to 1
Pk = torch.eye(N)      # P^0 = I
rwpe = []
for k in range(1, K+1):
    Pk = Pk @ P
    rwpe.append(Pk.diagonal())  # return probabilities at step k
pe = torch.stack(rwpe, dim=-1)  # [N, K]

Cost, stated honestly. Each step multiplies a dense \(N \times N\) matrix by the sparse \(P\): \(O(N\lvert E\rvert)\) time, giving \(O(K N \lvert E\rvert)\) overall, with \(O(N^2)\) memory. It is not \(O(K\lvert E\rvert)\) — that would be the cost of propagating a single vector, which does not yield the diagonal. This is cheap and completely reliable on the small graphs RWPE is normally used for (molecules, where \(N\) is tens of atoms), and it becomes the bottleneck well before LapPE would. For large graphs the practical options are stochastic estimators of the diagonal, or restricting to a \(K\)-hop subgraph around each node.

Worked Example: Triangle vs. Star Centre

Triangle node (\(K_3\), every node degree 2). Here \(P = \begin{pmatrix}0 & 0.5 & 0.5\\ 0.5 & 0 & 0.5\\ 0.5 & 0.5 & 0\end{pmatrix}\), and the diagonal entries are:

• \(P^1[v,v] = 0\) — no self-loop, so the walk must leave
• \(P^2[v,v] = 0.5\) — go to either neighbour (\(\tfrac12\) each) and come straight back (\(\tfrac12\) each): \(2 \times \tfrac12 \times \tfrac12\)
• \(P^3[v,v] = 0.25\) — the two triangle traversals, clockwise and anticlockwise, each of probability \(\tfrac18\)

Star centre (\(K_{1,4}\), centre of degree 4, leaves of degree 1). Every walk from the centre goes to a leaf, and from a leaf the only move is back:

• \[P^1[v,v] = 0\]
• \(P^2[v,v] = 1.0\) — whichever leaf is chosen, the return is forced
• \(P^3[v,v] = 0\) — a star is bipartite, so no odd return is possible

The profiles \([0,\, 0.5,\, 0.25,\, \dots]\) and \([0,\, 1.0,\, 0,\, \dots]\) differ immediately, so the model can tell a clique member from a star hub — a distinction 1-WL with uniform features cannot make in a regular graph, and one that plain message passing does not supply.

RWPE in Practice

GPS (Rampášek et al., 2022) evaluated both and found RWPE a strong default:

• No sign or basis ambiguity, so training is more stable and no invariantisation machinery is needed
• Directly encodes the cycle and clique structure that molecular tasks care about

The trade-off is exactly the one the encoding is designed around: RWPE tells you what a node is like, not where it is. On tasks where structural role dominates — molecular property prediction, graph classification — that is usually what is wanted, and RWPE is often the better choice. On tasks needing global position it is the wrong tool, and no value of \(K\) fixes that.

Summary

RWPE is the practical, ambiguity-free alternative to LapPE. It encodes local structural fingerprints — triangle membership, bipartiteness, local density — through repeated multiplication by the walk matrix, with no eigensolver and nothing to disambiguate. The cost is quadratic in \(N\) for the exact diagonal, which is fine at molecule scale and prohibitive at network scale. When global position matters, use LapPE; when local structural role matters, use RWPE; when both matter, use both, as GPS does.

References

• Dwivedi, V. P., Lim, A. T., Beaini, D., & Lió, P. (2021). Graph Neural Networks with Learnable Structural and Positional Representations. ICLR 2022.
• Rampasek, L., Galkin, M., Dwivedi, V. P., Lim, A. T., Wolf, G., & Beaini, D. (2022). Recipe for a General, Powerful, Scalable Graph Transformer. NeurIPS 2022 (GPS — uses RWPE as default positional encoding).