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

1 Introduction

• Sheaf neural networks assign local vector spaces to nodes and edges and use restriction maps to transport node representations into shared edge spaces for comparison.This lets neighboring representations remain consistent after transport even when their original features differ.Source: §1 · p. 2
• Explicitly assembling the global sheaf Laplacian adds memory and computational overhead, especially for large graphs or higher-dimensional stalks.The paper proposes local message passing as an alternative, while noting that its practical efficiency still requires empirical evaluation.Source: §1 · p. 2
• Sheaf-MPNN implements 15 families of SNNs with homogeneous and heterogeneous variants, directed and undirected settings, and multiple restriction-map types.The library also includes attention-based map predictors and a unified benchmarking suite spanning more than 20 datasets and node-, graph-, heterogeneous-graph, and efficiency tasks.Source: §1 · p. 3
• The framework gives explicit local message-passing formulations for 15 architecture configurations and establishes theoretical equivalence to the original models.Augmented node states, type-conditioned restriction maps, and directed aggregation rules are identified as mechanisms needed to preserve propagation dynamics.Source: §1 · p. 3

2 Sheaf Message Passing in a Nutshell

• SheafMPNN transports endpoint representations into a shared edge stalk, computes their interaction, then pulls the message back to aggregate at the receiving node.This local view treats each edge as comparing endpoint signals after transport.Source: §2 · p. 3, Figure 1 · p. 4
• Heterogeneous graphs use type-conditioned restriction-map predictors, while directed graphs use incoming-arc aggregation.Directed propagation is explicit because an undirected edge’s two restriction maps are transposes rather than independent directional operators.Source: §2 · p. 4
• Theorem 1 represents Neural Sheaf Diffusion, standard MPNNs, SheafAN variants, sheaf adjacency propagation, and Neural Sheaf Propagation within the framework.Neural Sheaf Propagation uses an augmented sheaf with node stalks R2d and edge stalks Rd.Source: §2 · p. 4
• The framework extends to diagonal, general, orthogonal, low-rank, and attention-generated restrictions and to additional sheaf variants.The paper reports matching normalization and updates for the listed representations.Source: §2 · pp. 4–5
• Local transport and aggregation evaluate the sheaf Laplacian without assembling its global nd × nd matrix representation.For materialized messages, the stated cost is O(|E|d2f) operations and O(|E|d2 + |E|df + ndf) memory, excluding map prediction, channel mixing, and block normalization.Source: §2 · p. 5

3 Library Design and Usage

• Sheaf-MPNN provides reusable sheaf message-passing layers and end-to-end model wrappers through a unified interface built on PyTorch and PyTorch Geometric.The example uses orthogonal Neural Sheaf Diffusion for node classification.Source: §3 · p. 5
• The sheaf command-line interface supports reproducible experiments across more than 20 datasets and Optuna-based hyperparameter optimization.It also enables computational efficiency comparisons against reference implementations.Source: §3 · p. 5

4 Related Software

• Existing sheaf implementations primarily target individual architectures and often assemble the explicit Laplacian, whereas Sheaf-MPNN offers reusable local message passing across a broader family.The paper contrasts this with general-purpose frameworks lacking native cellular-sheaf support and TopoX’s support for higher-order domains, including operator-based sheaf models.Source: §4 · p. 5

• The library is publicly available under CC-BY 4.0, with documentation and tutorial notebooks.Appendix C provides implementation details, supported architectures, and experimental configurations.Source: §5 · p. 5

Appendix A. Primer on Cellular Sheaves and Sheaf Neural Networks

• The primer places Sheaf-MPNN within cellular sheaf theory and the broader literature on sheaf learning.It points readers to fuller treatments of cellular sheaf theory, sheaf-Laplacian spectral theory, and sheaf learning surveys.Source: Appendix A · p. 7

A.1 From Message Passing to Sheaves

• Shared-space neighbor averaging can oversmooth features and erase class-separating signals on heterophilic graphs; cellular sheaves instead use node- and edge-specific spaces and learnable transport maps.The maps determine how signals are compared across each edge.Source: §A.1 · p. 7
• A cellular sheaf assigns vector spaces to vertices and edges, with linear restriction maps transporting vertex data into edge spaces for comparison.Agreement is evaluated after transport, rather than by comparing raw neighboring features.Source: §A.1 · pp. 7–8
• The sheaf Laplacian aggregates transported edge disagreements; its zero-disagreement signals are global sections, and diffusion converges toward them.For the constant sheaf, global sections are constant signals, whereas general sheaves can permit different neighboring features that agree after transport.Source: §A.1 · p. 8
• Neural Sheaf Diffusion learns incidence maps from endpoint features and was shown to separate classes in some heterophilic settings where graph diffusion cannot.The maps may be diagonal, orthogonal, or general matrices, and Sheaf-MPNN evaluates the update without materializing the sheaf Laplacian.Source: §A.1 · p. 8

A.2 Sheaf Neural Network Literature

• Early work established sheaf-Laplacian theory, learning, and neural models; later work explored sheaf learning methods, while recent benchmark findings suggest learned-map value depends on the task.The cited studies report that identity sheaves can be competitive on standard heterophily benchmarks, and that SNNs do not surpass the strongest baselines on an inductive benchmark.Source: §A.2 · pp. 8–9
• Architectures vary signal dynamics through attention, higher-order recurrences, polynomial filters, adjacency propagation, and long-range or equilibrium approaches.Other work extends sheaves to heterogeneous, directed, temporal, hierarchical, hypergraph, and non-Euclidean settings, where restriction maps must respect stalk structure.Source: §A.2 · pp. 9–10

• The release uses (Du + Id)^−1/2 normalization, and NSP supports per-block learnable skip rescalings through add_eps.For NSP, the two rescalings lie in [−1, 1]d.Source: Appendix B · pp. 10, 13
• Theorem 1 places Neural Sheaf Diffusion, standard MPNNs, SheafAN variants, sheaf adjacency propagation, and Neural Sheaf Propagation within the SheafMPNN blueprint.The blueprint permits local message-passing evaluation without constructing the global sheaf Laplacian.Source: Appendix B · p. 10
• Neural Sheaf Diffusion uses normalized restriction maps and difference messages; identity restrictions and transforms recover general MPNNs.The NSD update includes a dimension-wise rescaling and diffusion step size.Source: Appendix B · pp. 10–11
• SheafAN locally normalizes attention scores, while its residual variant adds a skip connection; sheaf adjacency propagation uses normalized maps and a residual update.The adjacency formulation absorbs normalization into the restriction maps.Source: Appendix B · pp. 10–12
• Neural Sheaf Propagation is represented by augmenting each node state with its history block, allowing the current block to reproduce its second-order update exactly.The augmented node stalks are R2d and edge stalks Rd; the restriction maps are non-square.Source: Appendix B · p. 12

B.1 Polynomial Neural Sheaf Diffusion

• PolyNSD evaluates a degree-K polynomial filter through the Chebyshev recurrence, replacing one NSD Laplacian application without eigendecomposition.The iterates use T0 = z, T1 = b∆F z, and Tk+1 = 2 b∆F Tk − Tk−1; the spectrum is rescaled from [0, λmax], with λmax = 2 for symmetric normalization.Source: §B.1 · p. 13
• A degree-K PolyNSD layer equals K SheafMPNN steps on an augmented sheaf, with identity weights, difference messages, sum aggregation, and node-local updates.The augmented state carries the iterates and accumulator; the readout is node-wise and needs no further propagation.Source: §B.1 · p. 14
• For K = 1 and λmax = 2, a parameter choice reduces PolyNSD to the original NSD update with unit step size.The choice αhp = θ1 − θ0 cancels the first term, and doubling W(ℓ) absorbs the remaining factor of 1/2.Source: §B.1 · p. 14
• The equivalence extends to all six NSD restriction-map variants because the proof requires only linear restriction maps.For orthogonal maps, the degree matrix is deg(u) Id, so normalization reduces to a scalar per node.Source: §B.1 · p. 14
• The degree-K layer costs O(K|E|d2f) operations and O(|E|d2 + nd2 + ndf) memory, excluding map prediction.Each step has the same O(|E|d2f) arithmetic as sparse multiplication, but the local method avoids assembling and storing the global nd × nd operator.Source: §B.1 · p. 15

B.2 SheafPool as Graph-Level Sheaf Message Passing

• Before pooling, SheafPool whitens each node stalk and aligns it to a learnable anchor by orthogonal Procrustes.The virtual node aggregates score-weighted aligned stalks, updates a learnable graph token, and applies the readout.Source: §B.2 · pp. 15–16
• SheafPool is one SheafMPNN layer on a pooling graph, followed by a node-wise readout at each virtual node.Each graph adds a receive-only virtual node; incoming arcs carry identity restrictions, and real nodes receive no messages.Source: §B.2 · pp. 15–16
• The readout is permutation invariant when all node-wise choices are shared and deterministic, but general linear basis invariance is not guaranteed.Regularized whitening can change magnitude under basis rescaling; exact unregularized whitening requires full row rank, and orthogonal invariance additionally needs a unique Procrustes solution or equivariant tie-breaking.Source: §B.2 · p. 16

B.3 Heterogeneous NSD and NSP as SheafMPNNs

• Type-conditioned Het-NSD and Het-NSP fit the SheafMPNN blueprint; Het-NSP stores recurrence history in an augmented state, and both reduce to homogeneous models with one node and relation type.The map predictors can use different context strategies and combinations of node features, node types, edge types, and ensembles.Source: §B.3 · pp. 17–18

B.4 Directed Extension

• Directed sheaf layers fit the blueprint when messages are computed independently per arc and incoming messages use permutation-invariant aggregation, including asymmetric predictors and messages.Source and target incidences may use different predictors, and endpoint order may affect the message function.Source: §B.4 · p. 18

B.5 Cooperative Sheaf Neural Networks; B.6 Joint Sheaf Neural Networks

• Cooperative SNNs use separate in- and out-maps for sending and receiving, while joint SNNs fit the framework when their edge-state process yields linear maps and feature updates use sheaf diffusion.A recurrent joint process can be represented by augmenting the state with its edge state.Source: §B.5 · p. 19, §B.6 · p. 19

B.7 Bundle Neural Networks

• With a truncated Taylor heat-kernel approximation, BuNN layers become repeated local SheafMPNN steps; exact spectral evaluation instead requires eigendecomposition.Flatness and orthogonality constrain the map family but do not change the local computation.Source: §B.7 · p. 20

B.8 Bayesian Sheaf Neural Networks

• Conditional on each sampled latent sheaf, a BSNN layer is a SheafMPNN evaluation, and its posterior-predictive output averages those evaluations across samples.Sampling changes map learning and gradient estimation, not the local computation after a map draw.Source: §B.8 · p. 20

C.1 Package Layout; C.2 Model Families and Restriction Maps

• Three packages separate homogeneous convolutions, heterogeneous counterparts, and experiment tooling; the registry offers 15 homogeneous families with type-conditioned counterparts and family-specific map variants.Convolution modules can be used in PyTorch Geometric models without adopting the supplied training loop.Source: §C.1 · p. 21, §C.2 · pp. 21–22, Table 2 · p. 22

C.3 Anatomy of an NSD Layer

• The NSD layer precomputes per-edge products and sums pulled-back messages to apply the normalized sheaf Laplacian without forming its global block matrix.Training adds jitter for general and low-rank maps and clips their normalized products; the equation is exact at evaluation when clipping is inactive.Source: §C.3 · p. 23

C.4 End-to-End Models and Direct Layer Use

• NSDModel encodes node features into stalk states, passes current states and flattened context to diffusion layers, and decodes the final flattened state; convolutions can also be called directly with supplied context.The direct-use example notes that hidden_dim is retained for interface compatibility and adds no hidden layer to the linear restriction predictor.Source: §C.4 · pp. 23–25

C.5 Experiment Tooling

• The CLI supports homogeneous node classification, while separate entry points handle heterogeneous node classification and graph classification.The runners cover citation, heterophily, OGB, heterogeneous HGB, and TU graph-classification datasets, among others.Source: §C.5 · p. 25
• Typed configurations specify data, models, training, hardware, and logging; dataset split protocols and edge symmetrization or directed settings determine the experiment setup.The sweep tool uses Optuna’s TPE sampler and median pruner, and re-evaluates the best configuration over ten seeds or folds.Source: §C.5 · p. 25
• The benchmark tool submits SLURM runs and reports convergence, epoch time, peak GPU memory, and inference latency.The example manifest compares diagonal, orthogonal, and general NSD against the original implementation.Source: §C.5 · pp. 25–26

C.6 Development and Reproducibility

• The project pins its Python, PyTorch, and Lightning versions and runs unit, integration, and benchmark integration tests on distinct schedules.The unit tests run on pushes and pull requests, integration tests daily, and benchmark integration tests weekly.Source: §C.6 · p. 26
• Reproducing an experiment requires the code version, environment, data preprocessing, model and map variant, splits and seeds, tuning budget, and checkpoint-selection rule.Numerical comparison with a reference additionally requires matched weights, normalization, random choices, and error tolerances.Source: §C.6 · p. 26
• The paper reports neither test coverage nor numerical parity with reference implementations.Ten tutorial notebooks accompany the code, including hypergraph and directed-hypergraph topics that are not registered model types.Source: §C.6 · p. 26