Graph Neural Networks with Learnable Structural and Positional Representations

In plain words

Graph neural networks (GNNs) learn from data shaped like networks, such as molecules, where atoms are nodes and bonds are edges. Most of them update each node by looking only at its neighbours. Two atoms with the same neighbourhood then look identical to the network, even when they sit in different places in the molecule. A common fix is to give every node a position-like tag. Fixed tags have known limits, so this paper gives each node a small vector of position values that the network learns alongside its normal features and refines at every layer. The starting values come from a simple random-walk calculation. The authors show that this helps several different GNN designs, and the largest gains are on molecular benchmarks.

TL;DR

• Adds a learnable positional vector per node, kept separate from the structural features and updated in every layer (LSPE).
• Initial positions are random-walk return probabilities (RWPE) or Laplacian eigenvectors (LapPE); RWPE has no sign ambiguity.
• Applied to two sparse GNNs (GatedGCN, PNA) and two fully connected Transformer GNNs (SAN, GraphiT).
• Author-reported: GatedGCN-LSPE test MAE 0.090 on ZINC against 0.251 without positional encoding tab. 1. On the two OGB datasets the authors report results comparable to SOTA, not clear wins §4.2 tab. 2.
• Code is public (MIT). The Laplacian-loss term is implemented with a detached tensor, so it does not train as written in the paper (Paper vs code).

Problem

• Message-passing GNNs (MP-GNNs) build node states from local neighbourhoods. They cannot tell apart two nodes with the same 1-hop structure, and their expressivity is bounded by the 1-WL test §1 §2.
• Stacking layers can spread information over more hops, but suffers from over-squashing for distant nodes §1.
• Higher-order WL-style GNNs are more expensive; they are not linear in the number of nodes §1 §2.
• Existing positional encodings have limits: index encodings need permutation training; Laplacian eigenvectors have sign ambiguity (2^k sign patterns for k vectors) and can be unstable under eigenvalue multiplicities; anchor-set embeddings do not generalise inductively §2.
• Existing MP-GNN-PE models concatenate a fixed PE to the input features only, so the positional representation cannot adapt during training §3.1.
• Fully connected Transformer GNNs use non-learnable PEs, or learned PEs built on Laplacian eigenvectors and therefore inherit the sign ambiguity §2.

Method

The paper keeps two streams per node: structural features h and positional features p. Both are updated at every layer. The structural update reads the concatenation [h; p]; the positional update is its own message-passing function §3.1 eq. 7 eq. 8 eq. 9. Positional features are not used as fixed inputs.

1. Standard MP-GNN layer. Node and edge updates are functions of neighbours and edge features §3.1 eq. 1 eq. 2. Transformers are a special case of MP-GNNs with fully connected graphs and no edge features §3.1.
2. Baseline MP-GNN-PE. The PE is concatenated with the input node features and passed through one linear map; the layers do not update it §3.1 eq. 3 eq. 4 eq. 5 eq. 6.
3. LSPE layer. Each layer updates h, e and p §3.1 eq. 7 eq. 8 eq. 9. The p update uses tanh so that coordinates can be positive or negative §3.1.
4. Input projection. The initial PE p_i^PE (length k) is mapped to the model width d by a learnable linear layer §3.1 eq. 11.
5. Initial PE. Two choices §3.1:
   • LapPE: the eigenvectors U_i1 .. U_ik of the Laplacian §3.1. Sign ambiguity needs random sign flips during training §3.1.
   • RWPE: the return probabilities RW_ii, RW_ii^2, ..., RW_ii^k with RW = A D^-1 §3.1 eq. 10. Only the diagonal of each power is kept, unlike Distance Encoding, which uses pairwise return probabilities §4.3 app. A.1. RWPE is precomputed before training app. A.
6. Positional loss (optional). A Laplacian-eigenvector term on the final positional features p^L §3.1 eq. 12. The final p^L is required to be centred and of unit norm per column §3.1. The Laplacian is Δ = I - D^-1/2 A D^-1/2 app. B.2.
7. Readout. The final positional features are concatenated with the structural features before pooling; Table 3 reports this as [h^L, p^L] §4.3 tab. 3.

GatedGCN-LSPE (the instance with the best reported numbers)

The authors write the gated convolution with the concatenated stream as §3.2 eq. 13 eq. 14 eq. 15 eq. 16. The gate for edge (i, j) is a sigmoid of a linear function of both endpoints and the edge, normalised over the incoming edges of i §3.2.

# GatedGCN-LSPE layer  [§3.2] [eq. 13-16]
# Input:  h (n, d), p (n, d), e (E, d); src[E], dst[E] are edge endpoints
# Output: h, p, e for the next layer

def gatedgcn_lspe_layer(h, p, e, src, dst):
    # Gate logit per edge: B1 h_i + B2 h_j + B3 e_ij  [§3.2]
    gate_logit = B1(h)[dst] + B2(h)[src] + B3(e)    # (E, d)
    # Sigmoid gate, normalised over the incoming edges of each node  [§3.2]
    # eps is not stated in the paper
    # Code uses 1e-6  [code: layers/gatedgcn_lspe_layer.py#L51]
    gate = sigmoid(gate_logit)
    gate = gate / (sum_over_dst(gate) + eps)       # (E, d)
    # Structural update: A1 [h; p] plus gated sum of A2 [h_j; p_j]  [eq. 14]
    hp = concat(h, p)                              # (n, 2d)
    h_new = A1(hp) + sum_over_dst(gate * A2(hp)[src])
    # Code only: scale by graph-size normaliser before BN
    # [code: layers/gatedgcn_lspe_layer.py#L104]
    h_new = h_new * node_norm
    # Edge update: residual + ReLU(BN(gate logit))  [eq. 15]
    e_new = e + relu(batch_norm(gate_logit))
    # Positional update: C1 p_i plus gated sum of C2 p_j  [eq. 16]
    # tanh and residual follow in the next lines
    p_new = C1(p) + sum_over_dst(gate * C2(p)[src])
    p_out = p + tanh(p_new)                        # no BatchNorm on p  [§3.2]
    h_out = h + relu(batch_norm(h_new))            # residual  [eq. 14]
    return h_out, p_out, e_new

Positional head (RWPE only in the code)

The code applies the following to the final positional features only when the initial PE is RWPE gatedgcn_net.py.

# Final positional features before readout  [§3.1]
# [code: nets/ZINC_graph_regression/gatedgcn_net.py#L67-L115]
# Input:  h (n, d), p (n, d), graph ids per node
# Output: node features (n, d) for pooling

def positional_head(h, p, graph_ids):
    # Linear map from model width to k  [code: gatedgcn_net.py#L67]
    p = linear_p_out(p)                            # (n, k)
    # Centre each column within each graph
    # [§3.1] [code: gatedgcn_net.py#L98-L103]
    p = p - mean_per_graph(p, graph_ids)
    # Scale each column of each graph to unit L2 norm
    # [§3.1] [code: gatedgcn_net.py#L105-L111]
    p = p / norm_per_graph(p, graph_ids)
    # Fuse structure and position with a linear map
    # [code: gatedgcn_net.py#L115]
    return linear_Whp(concat(h, p))                # (n, d)

Initial PEs

# Random-walk PE (RWPE), precomputed once per graph  [§3.1] [eq. 10]
# Input:  adjacency (n, n), degrees (n,), k steps
# Output: pe (n, k)

def rwpe(adjacency, degrees, k):
    # Random-walk operator RW = A D^-1; degrees clipped to at least 1
    # [code: data/molecules.py#L182-L184]
    walk = adjacency @ diag(1 / max(degrees, 1))   # (n, n)
    # Return probability of each node to itself after 1 .. k steps  [eq. 10]
    power = walk
    columns = [diagonal(power)]
    for _ in range(k - 1):
        power = power @ walk
        columns.append(diagonal(power))
    return stack(columns, axis=1)                  # (n, k)

# Laplacian-eigenvector PE (LapPE)  [§3.1]
# [code: data/molecules.py#L153-L162]
# Input:  adjacency (n, n), degrees (n,), k
# Output: pe (n, k)

def lap_pe(adjacency, degrees, k):
    # Symmetric normalised Laplacian  [app. B.2]
    d_half = diag(1 / sqrt(max(degrees, 1)))
    laplacian = identity(n) - d_half @ adjacency @ d_half
    # Eigenvectors sorted by ascending eigenvalue
    values, vectors = eig(laplacian)
    # The code drops the first (trivial) eigenvector; the paper does not say so
    return vectors[:, 1 : k + 1]                   # (n, k)

# Training-time sign flip for LapPE only  [§3.1]
# [code: train/train_ZINC_graph_regression.py#L35-L37]
# Input:  pe (n_total, k)

def flip_signs(pe):
    # One random sign per column, shared by every node in the batch
    signs = where(uniform(k) >= 0.5, 1.0, -1.0)    # (k,)
    return pe * signs

Laplacian-eigenvector loss (optional)

The paper defines the total loss as a task loss plus alpha times a Laplacian term over the final positional features §3.1 eq. 12. The Laplacian term has a trace (smoothness) part and a Frobenius orthogonality part with weight lambda §3.1.

\mathcal{L} = \mathcal{L}_{\text{task}}\big(h^L, p^L\big) + \alpha\,\mathcal{L}_{\text{LapEig}}\big(p^L\big) \qquad \text{[eq. 12]}

\mathcal{L}_{\text{LapEig}}(p) = \frac{1}{k}\operatorname{tr}\big(p^\top \Delta p\big) + \frac{\lambda}{k}\,\big\|p^\top p - I_k\big\|_F^2 \qquad \text{[§3.1]}

The code computes the same two terms with a different normalisation, and the orthogonality term is detached. See Paper vs code.

SAN-LSPE and GraphiT-LSPE (Transformer GNNs)

• SAN-LSPE. The attention over all node pairs is split into existing and non-existing edges, weighted by a balance scalar gamma app. C.2.1. The authors make gamma learnable and clip it to [0, 1] app. C.2.1. The positional stream has its own query, key and value projections and a tanh update app. C.2.1. The code clips gamma to [0, 1] in the attention layer san_gt_lspe_layer.py.
• GraphiT-LSPE. Attention weights are multiplied by a diffusion kernel K = (I - beta Δ)^p, with separate parameters for edges and non-edges app. C.2.2. The authors give beta in [0.25, 0.50] app. C.2.2. The ZINC config uses beta = 0.25 and p = 16 GraphiT_ZINC_LSPE.json molecules.py.
• PNA-LSPE applies the same two-stream update with PNA's principal aggregators app. C.1.2. Its layer code was not compared line by line (see Paper vs code).

Hyperparameters

Results

Author-reported. Setup: ZINC (12K subset, MAE, lower is better, about 500k parameters); OGBG-MOLTOX21 (ROC-AUC, higher is better); OGBG-MOLPCBA (AP, higher is better); all scores are means over 4 seeds §4.1 tab. 1. Parameter counts for OGB models are larger, following OGB leaderboard practice §4.1.

Headline numbers:

• GatedGCN-LSPE with PosLoss: test MAE 0.090 ± 0.001 on ZINC, against 0.251 ± 0.009 without PE. The paper reports a 64.14% gain §4.2 tab. 1.
• PNA-LSPE: test MAE 0.095 ± 0.002 on ZINC, against 0.141 ± 0.004 for PNA; a 32.62% gain §4.2 tab. 1.
• Against SOTA on ZINC (MAE 0.139 for SAN, 0.122 for Graphormer), GatedGCN-LSPE improves by 35.25% and 26.23% §4.2 tab. 2.
• On MOLTOX21, GatedGCN-LSPE reaches test ROC-AUC 0.7754 against 0.7757 for GIN-VN, i.e. comparable §4.2 tab. 2.
• On MOLPCBA, PNA-LSPE reaches AP 0.2840 against 0.2838 for PNA and 0.2947 for PHC-GNN §4.2 tab. 2.
• Non-molecular graphs (GatedGCN): accuracy rises on IMDB-B, IMDB-M and CIFAR10 app. D tab. 4. The IMDB gains are smaller than the standard deviations tab. 4.

Table 1: All instances (ZINC, OGBG-MOLTOX21, OGBG-MOLPCBA) tab. 1

Timing columns (Epoch/Total time) are dropped here because they depend on hardware. The paper bolds the best results; those marks are omitted here.

ZINC (test MAE, lower is better)

OGBG-MOLTOX21 (test ROC-AUC, higher is better)

OGBG-MOLPCBA (test AP, higher is better)

Table 2: Comparison with baselines and SOTA (best LSPE per dataset) tab. 2

ZINC: models with about 500k parameters. OGBG-MOLTOX21 and OGBG-MOLPCBA baselines are taken from the OGB leaderboards, and their parameter counts differ tab. 2 caption.

The paper excludes Graphormer's MOLPCBA score, which uses 3.8M pre-training graphs, from this comparison §4.2.

Table 3: Ablation on ZINC with GatedGCN tab. 3

The paper's text says that injecting RWPE at the final layer instead of the input performs poorly §4.3. No row in Table 3 tests that injection alone; the closest row, concatenating RWPE at the output without LSPE, gives 0.249 tab. 3.

Table 4: Non-molecular graphs (GatedGCN, accuracy) tab. 4

Baseline CIFAR10 values are taken from the benchmarking paper (†) tab. 4 caption. IMDB uses 10-fold cross-validation per seed tab. 4 caption.

Table 5: Hyperparameters per model tab. 5 app. E

Columns follow the paper. beta and p apply to GraphiT only. alpha and lambda apply with PosLoss only. Dashes mean not applicable.

Paper vs code

Code: https://github.com/vijaydwivedi75/gnn-lspe at 6b38443677befd4726cc0e5c20e9c88073fbc2ed (MIT, LICENSE file). The repository is linked from the abstract's footnote. I checked the ZINC GatedGCN path (model, RWPE, LapPE, loss, configs), the GraphiT kernel parameters and SAN's gamma clamp. I did not compare the PNA layer, the SAN and GraphiT layer bodies, or the non-molecular runs line by line.

• Laplacian loss: orthogonality term has no gradient. The code converts each graph's positional features to NumPy with .detach() before computing the Frobenius term, then wraps the result in a new torch.tensor. The term adds a constant to the loss value and does not change the gradient gatedgcn_net.py. The paper's Eq. 12 treats it as a differentiable term §3.1 eq. 12. Of the two Laplacian-loss terms, only the trace term gives gradients.
• Laplacian loss: normalisation differs. The paper's Loss_LapEig has a 1/k factor only §3.1. The code divides the sum by k × B × n, where B is the number of graphs and n the number of nodes in the batch gatedgcn_net.py. So the effective alpha depends on batch size and graph size. The paper does not state this.
• Laplacian matrix. The code uses I - D^-1/2 A D^-1/2 gatedgcn_net.py, the same matrix as app. B.2.
• Positional head adds layers not in the paper text. The code applies a linear map d → k to p before centring and normalising gatedgcn_net.py and a linear fusion [h; p] → d before pooling gatedgcn_net.py gatedgcn_net.py. Table 3 writes the fusion as concatenation [h^L, p^L] tab. 3.
• Centring and normalisation apply to RWPE only. The code runs them only when pe_init == 'rand_walk' gatedgcn_net.py. The paper states the constraint for p^L generally §3.1.
• LapPE + LSPE is not in the code path. For pe_init == 'lap_pe' the forward pass adds the LapPE to h and sets p = None gatedgcn_net.py. The configs directory has no LapPE + LSPE config. The LapPE + LSPE rows of tab. 3 are therefore not reproducible from the repository as I read it.
• Extra h scaling. The code multiplies h by a graph-size normaliser snorm_n before batch norm gatedgcn_lspe_layer.py. The paper does not mention this step.
• Gate epsilon. The code uses 1e-6 in the gate denominator gatedgcn_lspe_layer.py. The paper writes epsilon without a value §3.2.
• LapPE drops the trivial eigenvector. The code keeps eigenvectors 1 .. k after sorting by eigenvalue and drops index 0 molecules.py. The paper does not say this.
• Sign flips. The code draws one random sign per column for the whole mini-batch at each training step train_ZINC_graph_regression.py. The paper says only "random sign flipping during training" §3.1.
• GraphiT naming. The code's gamma is GraphiT's diffusion parameter beta GraphiT_ZINC_LSPE.json. The SAN balance scalar is also called gamma but is a different quantity SAN_ZINC_LSPE.json.
• SAN gamma. The code clamps gamma to [0, 1] in the attention layer san_gt_lspe_layer.py, as app. C.2.1 states. The initial value 0.1 is set in the config SAN_ZINC_LSPE.json and is not stated in the paper.

Pitfalls

• Abstract range is inconsistent with the results. The abstract says the gains range "from 1.79% up to 64.14%". The MOLTOX21 PNA-LSPE gain in the text is 0.79% (0.761 vs 0.755) §4.2 tab. 1. The 1.79% figure is for MOLPCBA §4.2. Do not use 1.79% as the lower bound.
• The Laplacian loss gain does not come from the orthogonality term in the code. The text says the loss moves the best ZINC score from 0.093 to 0.090 §4.2. In the code, the orthogonality part is constant (see Paper vs code). In the code, any gain from the loss can only come from the trace term and the normalisation. An implementer who wants the paper's loss must decide which one to build and say so.
• Positional features are not unique per node. The authors report that most nodes in ZINC molecules get unique RWPE vectors for k ≥ 24 §3.1. Two example molecules have 100% and 71.43% unique RWPEs §3.1. Appendix A.3 reports that 30 validation graphs with 24 nodes have 21 unique RWPEs, 87.5% of nodes app. A.3. The authors call RWPE uniqueness an approximation app. A.3.
• RWPE and symmetry. For strongly regular graphs (for example CSL), all nodes get the same RWPE, but distinct graph classes get distinct values §3.1 app. A.1.
• LapPE sign ambiguity. The paper's explanation is 2^k sign patterns for k eigenvectors §2. The code's sign flip is one sign per column for the whole batch train_ZINC_graph_regression.py.
• Parameter counts change across rows. LSPE adds parameters in the layers. For GatedGCN on ZINC, 504309 → 522870 tab. 1. For SAN, 501314 → 588066 tab. 1. The paper targets about 500k parameters on ZINC §4.1.
• OGB baselines are not matched on parameters. Table 2 notes that the OGB baselines have different parameter counts tab. 2 caption. Read the MOLTOX21 and MOLPCBA gains as comparable, not as wins §4.2.
• Non-molecular gains are within noise on IMDB. IMDB-B goes from 66.05 ± 6.63 to 70.03 ± 5.15, and IMDB-M from 45.77 ± 4.91 to 46.47 ± 4.00 tab. 4. Only CIFAR10 has a gap larger than its spread tab. 4.
• Cost of Transformer variants. The sparse GNNs train faster than the Transformer GNNs on the same hardware §4.2; the epoch timings are in tab. 1. Full attention scales quadratically with the number of nodes §2 §4.2.
• RWPE precomputation cost. The paper states that RWPE is precomputed before training app. A. The cost of the k matrix products per graph is not stated. The code performs them with sparse matrices once per graph at data load molecules.py.

Lineage

• Builds on: GatedGCN, arXiv:1711.07553 (no card); PNA (Corso et al., NeurIPS 2020; no arXiv id in the source); SAN, arXiv:2106.03893 (no card); GraphiT, arXiv:2106.05667 (no card).
• Related positional encoding: Distance Encoding (Li et al., NeurIPS 2020; no arXiv id in the source). The paper says RWPE is close to it, but uses only the diagonal and larger k than the k = 3 or 4 used there app. A.1 §4.3.
• Compared against: Graphormer, arXiv:2106.05234; the benchmarking protocol, arXiv:2003.00982 (no card).
• Superseded by: not known from this card.