Sheaf Graph Neural Networks via PAC-Bayes Spectral Optimization

TL;DR

SGPC introduces a unified framework for learning cellular-sheaf GNNs with PAC-Bayes calibration, combining a Wasserstein-Entropic Lift for restriction maps, SVR diffusion, and adaptive frequency mixing to robustly handle heterophily. The approach yields a spectrum-aware, bound-guided objective that comes with convergence guarantees, monotone spectral-gap growth, and risk-variance contraction, enabling end-to-end training with linear-time complexity. Theoretical contributions establish a tight PAC-Bayes generalization bound for cellular-sheaf GNNs and quantify diffusion stability via the spectral gap. Empirically, SGPC achieves state-of-the-art results across nine benchmarks with calibrated uncertainty intervals, demonstrating strong robustness to over-smoothing and heterophily while remaining scalable to large graphs.

Abstract

Over-smoothing in Graph Neural Networks (GNNs) causes collapse in distinct node features, particularly on heterophilic graphs where adjacent nodes often have dissimilar labels. Although sheaf neural networks partially mitigate this problem, they typically rely on static or heavily parameterized sheaf structures that hinder generalization and scalability. Existing sheaf-based models either predefine restriction maps or introduce excessive complexity, yet fail to provide rigorous stability guarantees. In this paper, we introduce a novel scheme called SGPC (Sheaf GNNs with PAC-Bayes Calibration), a unified architecture that combines cellular-sheaf message passing with several mechanisms, including optimal transport-based lifting, variance-reduced diffusion, and PAC-Bayes spectral regularization for robust semi-supervised node classification. We establish performance bounds theoretically and demonstrate that end-to-end training in linear computational complexity can achieve the resulting bound-aware objective. Experiments on nine homophilic and heterophilic benchmarks show that SGPC outperforms state-of-the-art spectral and sheaf-based GNNs while providing certified confidence intervals on unseen nodes. The code and proofs are in https://github.com/ChoiYoonHyuk/SGPC.

Figures (5)

• Figure 1: In the small-gap regime (left), two nodes are connected by a weak edge, so the Laplacian spectrum shows only a narrow separation between the first eigenvalue (0) and the second one ($\lambda_2$). After $\lambda_2$-gap expansion (right), the edge becomes strong and smoothly color-graded with a wide spectrum, illustrating the enlarged spectral gap
• Figure 2: (a) A Jordan-Kinderlehrer-Otto (JKO) step refines the initial Sinkhorn plan $P_{0}$ under the feature-cost matrix $C_{\mathrm{feat}}$, producing a globally stable coupling $P_{*}$; (b) The coupling $P_{ *}$ is turned into restriction maps $R_{ij}$, which in turn define the sheaf Laplacian $L_{\mathcal{F}}$; (c) Stochastic variance-reduced diffusion $D_{\xi}$ with adaptive frequency mixing $A_\theta$ yields node-level predictions $\hat{y}$; (d) A $\beta$-Dirichlet posterior calibrates edge uncertainty, while an optimizer enlarges the spectral gap $\lambda_{2}$
• Figure 3: (RQ2) Ablation and generalization study
• Figure 4: (RQ3) Hyperparameter analysis on loss function
• Figure 5: (RQ4) Over‑smoothing analysis on Cora dataset, where the metrics are Accuracy ($\uparrow$) and NRS ($\downarrow$)

Theorems & Definitions (14)

• Theorem 1: PAC-Bayes Sheaf Generalization Bound
• Theorem 2: CG convergence with sparsifier
• Theorem 3: Wolfe-controlled gap ascent
• Lemma 1: Variance reduction
• Theorem 4: Risk-Variance Contraction
• Lemma 2: Algorithmic stability bound
• Theorem 5: PAC-Bayes population risk
• proof
• proof
• proof