Sheaf-MPNN: A Library for Sheaf Neural Networks as Message Passing

Alessio Borgi; Mario Severino; Luke Braithwaite; Steve Azzolin; Rémi Bourgerie; Olga Zaghen; Emanuele Mule; Riccardo Ali

cs.LGNov 3, 2026

Sheaf-MPNN: A Library for Sheaf Neural Networks as Message Passing

Authors: Alessio Borgi, Mario Severino, Luke Braithwaite, Steve Azzolin, Rémi Bourgerie, Olga Zaghen, Emanuele Mule, Riccardo Ali,

Organizations: Department of Computer Science and Technology, University of Cambridge, Cambridge, UK · Department of Computer, Control and Management Engineering, Sapienza University of Rome, Rome, Italy · Department of Information Engineering and Computer Science, University of Trento, Trento, Italy · KTH Royal Institute of Technology, Stockholm, Sweden ·

Abstract

Standard Graph Neural Networks (GNNs) implicitly assume that neighboring nodes share the same feature space, which contributes to oversmoothing and poor performance on heterophilic graphs. Sheaf Neural Networks (SNNs) remove this assumption by giving each node and edge its own vector space and learning restriction maps that transform features as they move between neighbors. These maps define a sheaf Laplacian that drives diffusion over the graph, but existing implementations often assemble it explicitly as a global sparse block matrix, which adds memory and computational overhead and can hinder scalability. In this paper, we introduce Sheaf-MPNN, a PyTorch Geometric library that reformulates sheaf diffusion as local incidence-wise transport and aggregation, eliminating the need to construct the global Laplacian. Sheaf-MPNN implements fifteen families of SNNs under a unified message-passing framework, each supporting both homogeneous and heterogeneous graphs, and directed and undirected settings, with diagonal, general, orthogonal, and low-rank restriction maps and attention-based predictors. We further provide a unified benchmarking suite covering node- and graph-level tasks, heterogeneous graph learning, and computational efficiency, enabling systematic and reproducible comparisons across architectures. We establish theoretical equivalence between these message-passing formulations and their original counterparts and identify the architectural mechanisms needed to preserve their propagation dynamics. This unified perspective provides a principled foundation for implementing, comparing, and extending sheaf-based architectures within standard graph learning frameworks. Source code is available at https://github.com/alessioborgi/sheaf-mpnn.

Figures & tables

Panel of Figure 1 — Figure 1: Sheaf Message Passing Neural Network (SheafMPNN). Neighbourhood update at node $0$ (red), with neighbouring nodes in blue. (Left) A standard MPNN computes messages $m_{0i}$ (green) from neighbouring node states and aggregates them at node $0$ . (Right) A SheafMPNN transports both endpoint representations into the shared edge stalk $\mathcal{F}(e_{0i})$ (boxes) through restriction maps (dashed grey arrows). The function $\psi$ computes their interaction, and the resulting message is pulled back through $\F_{0\trianglelefteq e_{0i}}^{\top}$ and aggregated at node $0$ (solid black arrows).

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Module	Responsibility
sheaf_mpnn.base_conv	Shared stalk transforms, map-predictor inputs, and propagation interfaces.
sheaf_mpnn.<family>	Architecture-specific convolutions and end-to-end models.
sheaf_mpnn.utils	Normalization, orthogonal parameterizations, connection maps, and training utilities.
sheaf_mpnn.sheaf_pool	Graph-level pooling of stalk-valued representations.
hetsheaf_mpnn	Type-conditioned layers and model wrappers.
exp.config , exp.registries	Typed experiment settings and factories for models, datasets, and presets.

Appendix

Table 1: Sheaf-MPNN Main Modules.

Family	Model type	Map and model options
Neural Sheaf Diffusion	nsd	Six
Polynomial NSD	poly_nsd	Six
Connection-Laplacian NSD	conn_nsd	Precomputed maps
Deep NSD	dnsd	Six
Neural Sheaf Propagation	nsp	Six
Sheaf adjacency	sheaf_adj	Six

Appendix

Table 2: Homogeneous Architecture Families. The table lists the 15 homogeneous families and their corresponding model types and restriction-map options. In the last column, “six” denotes the six NSD restriction-map variants of Table 3 , and “four” the diagonal, general, orthogonal, and low-rank maps.

Variant	Restriction parameterization	Outputs
DIAGONAL	$\operatorname{diag}(w)$ , with $w\in(-1,1)^{d}$	$d$
GENERAL	Dense matrix with entries in $(-1,1)$	$d^{2}$
GENERAL_ATTENTION	$I-\operatorname{softmax}(P)$ , with row-wise softmax	$d^{2}$
ORTHOGONAL	Cayley or Householder construction	$\binom{d}{2}$ or $d^{2}$
ORTHOGONAL_ATTENTION	Cayley transform of the skew part of $I-\operatorname{softmax}(P)$	$d^{2}$
LOW_RANK	$AB^{\top}$ , with $A,B\in\mathbb{R}^{d\times r}$	$2dr$

Appendix

Table 3: NSD Restriction-map Variants . The last column gives the number of predictor outputs per restriction map, excluding fixed channels and edge weights. Low-rank maps need fewer outputs, but their products

S_{ue}

and

C_{uve}

are still stored as dense

d\times d

blocks.

Option	Effect
variant , rank , orth_strategy	Select the restriction-map variant and its parameterization.
alpha , learn_alpha , add_eps	Set the diffusion step and enable the learnable skip scaling.
use_edge_weights	Enable learned symmetric edge weights for orthogonal maps.
add_lp , add_hp	Append a fixed channel with signless ( $D+A$ ) or standard ( $D-A$ ) graph-Laplacian blocks, respectively.
sparse_learner	Sum the paired contexts over the stalk axis, reducing the predictor width from $2d_{\mathrm{tot}}f$ to $2f$ for flattened stalk inputs.
second_linear	Add a second encoder layer.

Appendix

Table 4: Options of NSDModel beyond the core dimensions in_channels , out_channels , stalk_dim , hidden_dim , and num_layers .

Sheaf Neural Networks

Sheaf neural networks generalize graph learning with local vector spaces and edge maps that model how information is compared and transported.

See all 30 papers

On the Necessity of Learnable Sheaf Laplacians
Mar 5, 2026Ferran Hernandez Caralt, Mar Gonzàlez i Català, Adrián Bazaga +1GNN Oversmoothing Sheaf Neural Networks
Department of Computer Science and Technology, University of Cambridge, Cambridge, United Kingdom · Microsoft
Hierarchical Pooling for Sheaf Neural Networks
Jun 18, 2026Dionisia Naddeo, Carlo Abate, Pietro Liò +2Graph Neural Networks Graph Coarsening
UiT The Arctic University of Norway · UiT The Arctic University of Norway NORCE Norwegian Research Centre AS
Gaussian Sheaf Neural Networks
May 20, 2026André Ribeiro, Ana Luiza Tenório, Tiago da Silva +1Graph Neural Networks Sheaf Neural Networks
Getulio Vargas Foundation · MBZUAI

Mar 5, 2026J/K move · Enter open · S save

On the Necessity of Learnable Sheaf Laplacians

Ferran Hernandez Caralt, Mar Gonzàlez i Català, Adrián Bazaga, Pietro Liò

Department of Computer Science and Technology, University of Cambridge, Cambridge, United Kingdom · Microsoft

GNN Oversmoothing Sheaf Neural Networks Graph Neural Networks Over-Smoothing in GNNs

Open paper AskPDF arXiv

Sheaf Neural Networks (SNNs) were introduced as an extension of Graph Convolutional Networks to address oversmoothing on heterophilous graphs by attaching a sheaf to the input graph and replacing the adjacency-based operator with a sheaf Laplacian defined by (learnable) restriction maps. Prior work motivates this design through theoretical properties of sheaf diffusion and the kernel of the sheaf Laplacian, suggesting that suitable non-identity restriction maps can avoid representations converging to constants across connected components. Since oversmoothing can also be mitigated through residual connections and normalization, we revisit a trivial sheaf construction to ask whether the additional complexity of learning restriction maps is necessary. We introduce an Identity Sheaf Network baseline, where all restriction maps are fixed to the identity, and use it to ablate the empirical improvements reported by sheaf-learning architectures. Across five popular heterophilic benchmarks, the identity baseline achieves comparable performance to a range of SNN variants. Finally, we introduce the Rayleigh quotient as a normalized measure for comparing oversmoothing across models and show that, in trained networks, the behavior predicted by the diffusion-based analysis of SNNs is not reflected empirically. In particular, Identity Sheaf Networks do not appear to suffer more significant oversmoothing than their SNN counterparts.

Abstract

Sheaf Neural Networks

Explore similar work

On the Necessity of Learnable Sheaf Laplacians

Hierarchical Pooling for Sheaf Neural Networks

Gaussian Sheaf Neural Networks