Polynomial Neural Sheaf Diffusion: A Spectral Filtering Approach on Cellular Sheaves

TL;DR

PolyNSD introduces a spectral polynomial diffusion method for sheaf neural networks, achieving state-of-the-art results with improved stability, efficiency, and scalability on various graph benchmarks.

Contribution

It proposes a novel polynomial spectral filtering approach for sheaf diffusion that reduces complexity and enhances stability, outperforming existing methods.

Findings

Achieves state-of-the-art results on homophilic and heterophilic benchmarks.

Decouples performance from stalk dimension, reducing runtime and memory.

Provides explicit K-hop receptive field in a single layer.

Abstract

Sheaf Neural Networks equip graph structures with a cellular sheaf: a geometric structure which assigns local vector spaces (stalks) and a linear learnable restriction/transport maps to nodes and edges, yielding an edge-aware inductive bias that handles heterophily and limits oversmoothing. However, common Neural Sheaf Diffusion implementations rely on SVD-based sheaf normalization and dense per-edge restriction maps, which scale with stalk dimension, require frequent Laplacian rebuilds, and yield brittle gradients. To address these limitations, we introduce Polynomial Neural Sheaf Diffusion (PolyNSD), a new sheaf diffusion approach whose propagation operator is a degree-K polynomial in a normalised sheaf Laplacian, evaluated via a stable three-term recurrence on a spectrally rescaled operator. This provides an explicit K-hop receptive field in a single layer (independently of the stalk dimension), with a trainable spectral response obtained as a convex mixture of K+1 orthogonal polynomial basis responses. PolyNSD enforces stability via convex mixtures, spectral rescaling, and residual/gated paths, reaching new state-of-the-art results on both homophilic and heterophilic benchmarks, inverting the Neural Sheaf Diffusion trend by obtaining these results with just diagonal restriction maps, decoupling performance from large stalk dimension, while reducing runtime and memory requirements.