Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On the Stability of Expressive Positional Encodings for Graphs

Yinan Huang, William Lu, Joshua Robinson, Yu Yang, Muhan Zhang, Stefanie Jegelka, Pan Li

TL;DR

This work tackles the instability and non-uniqueness of Laplacian-based graph positional encodings by introducing Stable and Expressive Positional Encodings (SPE). SPE formalizes a soft, eigenvalue-informed partition of eigenspaces, yielding provable stability and universal expressivity for basis-invariant encodings. Theoretical results bound output sensitivity to Laplacian perturbations and establish domain-generalization guarantees under distribution shifts, while experiments on molecular datasets and DrugOOD show improved generalization and substructure counting capabilities. Overall, SPE offers a principled, scalable approach to enhance graph transformers and GNNs by combining stability with expressive power.

Abstract

Designing effective positional encodings for graphs is key to building powerful graph transformers and enhancing message-passing graph neural networks. Although widespread, using Laplacian eigenvectors as positional encodings faces two fundamental challenges: (1) \emph{Non-uniqueness}: there are many different eigendecompositions of the same Laplacian, and (2) \emph{Instability}: small perturbations to the Laplacian could result in completely different eigenspaces, leading to unpredictable changes in positional encoding. Despite many attempts to address non-uniqueness, most methods overlook stability, leading to poor generalization on unseen graph structures. We identify the cause of instability to be a ``hard partition'' of eigenspaces. Hence, we introduce Stable and Expressive Positional Encodings (SPE), an architecture for processing eigenvectors that uses eigenvalues to ``softly partition'' eigenspaces. SPE is the first architecture that is (1) provably stable, and (2) universally expressive for basis invariant functions whilst respecting all symmetries of eigenvectors. Besides guaranteed stability, we prove that SPE is at least as expressive as existing methods, and highly capable of counting graph structures. Finally, we evaluate the effectiveness of our method on molecular property prediction, and out-of-distribution generalization tasks, finding improved generalization compared to existing positional encoding methods. Our code is available at \url{https://github.com/Graph-COM/SPE}.

On the Stability of Expressive Positional Encodings for Graphs

TL;DR

This work tackles the instability and non-uniqueness of Laplacian-based graph positional encodings by introducing Stable and Expressive Positional Encodings (SPE). SPE formalizes a soft, eigenvalue-informed partition of eigenspaces, yielding provable stability and universal expressivity for basis-invariant encodings. Theoretical results bound output sensitivity to Laplacian perturbations and establish domain-generalization guarantees under distribution shifts, while experiments on molecular datasets and DrugOOD show improved generalization and substructure counting capabilities. Overall, SPE offers a principled, scalable approach to enhance graph transformers and GNNs by combining stability with expressive power.

Abstract

Designing effective positional encodings for graphs is key to building powerful graph transformers and enhancing message-passing graph neural networks. Although widespread, using Laplacian eigenvectors as positional encodings faces two fundamental challenges: (1) \emph{Non-uniqueness}: there are many different eigendecompositions of the same Laplacian, and (2) \emph{Instability}: small perturbations to the Laplacian could result in completely different eigenspaces, leading to unpredictable changes in positional encoding. Despite many attempts to address non-uniqueness, most methods overlook stability, leading to poor generalization on unseen graph structures. We identify the cause of instability to be a ``hard partition'' of eigenspaces. Hence, we introduce Stable and Expressive Positional Encodings (SPE), an architecture for processing eigenvectors that uses eigenvalues to ``softly partition'' eigenspaces. SPE is the first architecture that is (1) provably stable, and (2) universally expressive for basis invariant functions whilst respecting all symmetries of eigenvectors. Besides guaranteed stability, we prove that SPE is at least as expressive as existing methods, and highly capable of counting graph structures. Finally, we evaluate the effectiveness of our method on molecular property prediction, and out-of-distribution generalization tasks, finding improved generalization compared to existing positional encoding methods. Our code is available at \url{https://github.com/Graph-COM/SPE}.
Paper Structure (34 sections, 17 theorems, 69 equations, 4 figures, 8 tables)

This paper contains 34 sections, 17 theorems, 69 equations, 4 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. A Provably Stable and Expressive PE
  4. Stable Positional Encodings
  5. SPE: A positional encoding with guaranteed stability
  6. From stability to out-of-distribution generalization
  7. SPE is a Universal Basis Invariant Architecture
  8. Related Works
  9. Experiments
  10. Small Molecule Property Prediction
  11. Out-Of-Distribution Generalization: Binding Affinity Prediction
  12. Trade-offs between stability, generalization and expressivity
  13. Counting Graph Substructures
  14. Conclusion
  15. Deferred Proofs
  16. ...and 19 more sections

Key Result

Theorem 3.1

Under Assumption assumption:spe, SPE is stable with respect to the input Laplacian: for Laplacians ${\bm{L}},{\bm{L}}^{\prime}$, where the constants are $\alpha_1=2J \sum_{l=1}^mK_{\ell}$, $\alpha_2=4\sqrt{2}J\sum_{l=1}^mM_{\ell}$ and $\alpha_3=J\sum_{l=1}^mK_{\ell}$. Here $M_{\ell}=\sup_{\bm{\lambda}\in[0,2]^d}\left\lVert \phi_{\ell}(\bm{\lambda}) \right\rVert$ and again ${\bm{P}}_*=\mathop{\mat

Figures (4)

  • Figure 1: Illustration of the SPE architecture (first two rows) and its stability property (last row). The input graph is first decomposed into eigenvectors ${\bm{V}}$ and eigenvalues $\bm{\lambda}$. In step 1, a permutation equivariant $\phi_{\ell}$ act on $\bm{\lambda}$ to produce another vector $\phi_{\ell}(\bm{\lambda})$. In step 2, we compute ${\bm{V}}\text{diag}\{\phi_{\ell}(\bm{\lambda})\}{\bm{V}}^{\top}$ for each $\phi_{\ell}$ and concatenate the results into a tensor. This tensor is input into a permutation equivariant network $\rho$ to produce final node positional encodings.
  • Figure 2: Training error, test error and generalization gap v.s. model complexity (stability). In the first row, we directly change the Lipschitz constant of individual MLPs; in the second row, we choose $\phi_{\ell}$ to be piecewise spline functions and change the number of pieces.
  • Figure 3: $\log_e(\text{Test MAE})$ over 3 random seeds on cycle counting task. Lower is better.
  • Figure 4: A illustration of $[\phi_1(\bm{\lambda})]_2$ for $\bm{\lambda}=(\lambda_1, \lambda_2)$. Clearly $\phi_1$ is discontinuous.

Theorems & Definitions (28)

  • Definition 3.1: PE Stability
  • Remark 3.1: Stability implies permutation equivariance
  • Theorem 3.1: Stability of SPE
  • Proposition 3.1
  • Proposition 3.2: Basis Universality
  • Proposition 3.3
  • Proposition 3.4: SPE can count cycles
  • Theorem A.1: Stability of SPE
  • proof
  • Proposition A.1
  • ...and 18 more