Track: Track1; Team name: Sheafu; Model: Neural Sheaf Propagation (NSP) + Holonomy Measurement (analysis notebook)

[Agrover112]

Checklist

• My pull request has a clear and explanatory title.
• My pull request passes the linting checks.
• I added appropriate unit tests and confirmed that the test suite passes.
• My PR follows PEP 8 guidelines.
• My code is documented using NumPy-style docstrings, and the documentation renders correctly.
• I linked the issues and pull requests relevant to this contribution.

Description

This PR implements Neural Sheaf Propagation (NSP), introduced in Surfing on the Neural Sheaf, as a Track 1 (GNN) submission for the TDL Challenge 2026 (model=graph/nsp).

From sheaf diffusion to sheaf propagation

Neural Sheaf Diffusion (NSD) is motivated by the first-order sheaf heat equation:

$$\dot{X}(t)=-\Delta_{\mathcal F(t)}X(t).$$

Heat diffusion progressively damps non-harmonic components of the signal. This is useful for smoothing, but repeated diffusion can also make node representations increasingly similar.

NSP replaces this first-order dynamics with a second-order sheaf wave equation:

$$\ddot{X}(t)=-\Delta_{\mathcal F(t)}X(t).$$

In the ideal linear system, the spectral modes oscillate rather than being monotonically damped. NSP uses this wave-like behaviour to propagate information while retaining non-smooth components that ordinary diffusion may suppress.

The learnable spatial term combines the sheaf Laplacian with left and right feature transformations:

$$\sigma\!\left( \Delta_{\mathcal F(t)} \left(I_n\otimes W_1^{t}\right) X^{t}W_2^{t} \right).$$

This implementation integrates the resulting second-order dynamics with a leapfrog update:

$$X^{t+1} = 2X^{t}-X^{t-1} -h^2\, \sigma\!\left( \Delta_{\mathcal F(t)} \left(I_n\otimes W_1^{t}\right) X^{t}W_2^{t} \right),$$

where:

• (X^{t-1}) and (X^t) are the two states retained by the leapfrog integrator;
• (h) is the configured step_size;
• (\mathcal F(t)) is the learned sheaf at propagation step (t);
• (\Delta_{\mathcal F(t)}) denotes the corresponding sheaf propagation operator;
• (W_1^t) mixes the stalk channels and (W_2^t) mixes the remaining feature channels.

The implementation initializes x_prev and x_curr to the same encoded feature tensor. This corresponds to zero initial discrete velocity.

Important: the underlying linear wave equation is non-dissipative. The complete neural layer also contains learned maps, nonlinear activations, and dropout, so the implemented network should not be described as exactly energy-conserving.

Provenance. The architecture follows the propagation scheme introduced in the paper and reuses TopoBench's existing Neural Sheaf Diffusion infrastructure. This is an independent implementation; see Related.

Configuration

configs/model/graph/nsp.yaml exposes one NSP architecture with the following options:

Option	Default	Effect
d	3	Sets the stalk dimension.
step_size	0.5	Sets the leapfrog step (h). The implementation scales the force term by (h^2).
sheaf_type	diag	Selects diagonal, orthogonal-bundle, or unrestricted restriction maps.
second_linear	false	Adds an extra feature projection, lin_second, before propagation begins.
new_laplacian_each_step	true	Recomputes the learned sheaf operator at every step. Set it to false to reuse one fixed operator.

Restriction-map families

sheaf_type	Restriction maps	Constraint	Evaluation status
diag	Diagonal maps	d >= 1	Default and benchmarked
bundle	Orthogonal maps in O(d)	d > 1	Implemented and unit-tested
general	Unrestricted d × d maps	d > 1	Implemented and unit-tested

The diagonal configuration is used for the reported challenge results.

What is contributed

Area	File	Contribution
Backbone	topobench/nn/backbones/graph/nsp.py	Adds NSPEncoder and adapts sheaf wave propagation to TopoBench's forward(x, edge_index) -> [N, hidden_dim] interface.
Propagation models	topobench/nn/backbones/graph/nsd_utils/inductive_discrete_models.py	Adds diagonal, bundle, and general NSP variants using the leapfrog recurrence.
Configuration	configs/model/graph/nsp.yaml	Defines the complete TBModel composition used by both challenge tasks.
Unit tests	test/nn/backbones/graph/test_nsp.py	Tests the encoder, configuration options, map families, gradients, and numerical finiteness.
Pipeline coverage	test/pipeline/test_pipeline.py	Adds end-to-end training coverage for graph/nsp on graph/MUTAG.

Reused TopoBench components

No changes were required to the following components:

• AllCellFeatureEncoder encodes the input features;
• GNNWrapper integrates the backbone into the model pipeline;
• MLPReadout converts the node embeddings produced by the backbone into task predictions.

Implementation and faithfulness notes

Leapfrog step size

The paper lists a step-size hyperparameter, but its displayed propagation equation does not explicitly place an (h^2) factor on the nonlinear force term. This implementation follows the standard central-difference discretization:

$$\frac{X^{t+1}-2X^t+X^{t-1}}{h^2} \approx \ddot{X}(t).$$

Substituting the wave equation gives the implemented recurrence with step_size**2:

$$X^{t+1}-2X^t+X^{t-1} = -h^2\, \sigma\!\left( \Delta_{\mathcal F(t)} \left(I_n\otimes W_1^t\right) X^tW_2^t \right).$$

The (h^2) coefficient is therefore an explicit numerical-discretization choice in this implementation rather than a claim that the paper writes the coefficient in exactly this form.

Dynamic and fixed sheaf geometry

The submitted configuration uses dynamic geometry by default:

new_laplacian_each_step: true

With this setting, the restriction maps and sheaf operator are recomputed from the current features at every propagation step. Setting the flag to false builds the operator once from the initially encoded features and reuses it throughout the network. In fixed mode, only one sheaf learner is retained, which avoids unused parameters.

No diffusion residual coefficient

The NSP classes remove NSD's learnable per-layer diffusion coefficients, commonly represented by (\varepsilon). Those coefficients belong to the first-order residual diffusion update and are not used by the second-order leapfrog recurrence.

Positional features

The paper also studies Laplacian eigenvectors as optional positional features. They are not computed inside NSPEncoder because TopoBench already provides LapPE, which can be applied as a separate input transform.

Framework adaptations

• The backbone returns node embeddings; MLPReadout produces the final task outputs.
• The input graph is converted to an undirected graph with to_undirected, matching the existing TopoBench graph-backbone convention.
• The diagonal and general variants reuse their corresponding TopoBench sheaf-Laplacian builders. The bundle variant uses the normalized connection-Laplacian builder.

Testing

test/nn/backbones/graph/test_nsp.py covers:

• default and custom initialization;
• output shape and finite values, with and without a batch vector;
• stalk dimension d=1 for the diagonal model;
• automatic backbone discovery;
• all three restriction-map families and their dimension constraints;
• the effect of step_size on the magnitude of a deep propagation response;
• numerical finiteness for a 16-layer model at the default step size;
• creation of lin_second when second_linear=true;
• one sheaf learner per layer for dynamic geometry and a single learner for fixed geometry;
• removal of NSD's unused epsilons parameters;
• gradient flow to every registered trainable parameter under the tested flag combinations.

graph/nsp is also included in test/pipeline/test_pipeline.py for end-to-end training on graph/MUTAG. The new files pass the configured ruff and numpydoc checks.

Evaluation

2026_tdl_challenge/run_evaluation.ipynb was run with the default diagonal restriction maps. The resulting file contains 72 runs across seeds 42, 43, and 44:

2026_tdl_challenge/outputs/nsp/results.json

The in-distribution means over the 12 GraphUniverse settings are:

Task	Metric	Mean
Community detection	Accuracy, higher is better	0.382
Triangle counting	MSE per triangle, lower is better	0.92

The full per-setting, per-seed, and out-of-distribution results are stored in results.json. The bundle and general variants are implemented and unit-tested but were not included in this benchmark run.

results.json provenance

On main, the notebook's committed expected_hash (f87b2c…) does not match the hash computed from its committed cells (3c1d78…). As a result, the notebook's integrity check raises a ValueError before evaluation begins.

For the reported run, expected_hash was updated locally to the hash of the existing notebook cells. The evaluation grid and utils.py were not changed. A separate maintainer issue can be used to refresh the stale committed hash.

Reproducing the results

# Unit tests
pytest test/nn/backbones/graph/test_nsp.py

# End-to-end pipeline test
pytest test/pipeline/test_pipeline.py

# Single training run
python -m topobench model=graph/nsp dataset=graph/MUTAG

Issue

As per the research questions of track-1, NSP was contributed to test whether second-order wave-based sheaf propagation can preserve structural and high-frequency information better than the first-order smoothing dynamics used by Neural Sheaf Diffusion.

Related

• TDL Challenge 2026, Track 1
• Surfing on the Neural Sheaf
• Neural Sheaf Diffusion

This PR is an independent NSP implementation built on TopoBench's NSD utilities. Its defining implementation choices are the explicit (h^2) leapfrog scaling, support for dynamic and fixed sheaf geometry, and diagonal, orthogonal-bundle, and general restriction-map variants.

Supplementary analysis: Holonomy Measurement (Team Sheafu)

Exploratory analysis notebook, under 2026_tdl_challenge/holonomy_audit/. The required NSP evaluation artifacts (run_evaluation.ipynb, results.json) are unchanged; the analysis adds no dependency to the NSP model. Its folder README.md follows.

Do sheaf networks use their geometry?

This directory contains supplementary analysis for Team Sheafu's Neural Sheaf Propagation (NSP) submission to the Topological Deep Learning Challenge 2026, Track 1. The analysis opens trained sheaf neural networks and looks at the geometry they learn around triangles.

Question and findings

The analysis asks two things: does a trained sheaf network actually develop cycle geometry, and do its predictions come to depend on it?

In these experiments the pattern is consistent. Training a model to count triangles makes its SO(2) maps rotate features around each triangle about 40x more than at initialisation; training the same architecture to detect communities leaves those maps nearly flat. The geometry is recruited, and only for the task that should need it. With enough data the model also starts to rely on it: flatten the learned connection afterwards and the error grows sharply.

What the geometry never does is count. A plain ridge regression on a handful of degree statistics still beats every sheaf model, so the networks do not even match that simple shortcut. And once we remove the shortcut, using regular graphs where every node has the same degree so that reading the wiring is the only way to tell graphs apart, no model does better than predicting the average, at any data scale and however large the twist grows. The geometry gets recruited and becomes load-bearing, but it never turns into a working cycle counter.

The companion paper is Measuring Triangle Holonomy in Sheaf Neural Networks (arXiv:2607.19514).

Files

holonomy_audit/
├── README.md
├── holonomy_audit.ipynb          the analysis notebook
├── build_holonomy_notebook.py    builds the notebook
├── holonomy_experiment.py        main sweep
├── holonomy_lobotomy.py          flattening intervention
├── holonomy_regular.py           regular-graph counting
├── degree_baselines.py           degree-shortcut baseline
├── verify_diagnostic.py          measurement-tool checks
├── scaling_experiment.py         data-scaling sweep
├── scaling_baselines.py          baselines at each size
├── bigtest_eval.py               200-graph re-check
└── *_results.json  (11 files)    saved outputs the notebook reads

File	Purpose
holonomy_audit.ipynb	The analysis itself: the notebook you read.
build_holonomy_notebook.py	Builds holonomy_audit.ipynb. The notebook's text and code live here as a plain Python script.
holonomy_experiment.py	The main sweep. Trains the small sheaf models (both NSP and NSD dynamics, across the three map families, on both tasks, over many seeds) and records the twist, gain, and flip each one learns.
holonomy_lobotomy.py	The intervention. Takes a trained model, deletes its learned geometry by setting every restriction map back to the identity, and re-measures it, to test whether the model actually depends on that geometry.
holonomy_regular.py	Runs triangle counting on fixed-degree (regular) graphs, where every node has the same degree, so the degree shortcut is gone and only the wiring distinguishes one graph from another.
degree_baselines.py	Fits the shortcut itself: simple ridge regressions that predict triangle counts from structural summaries like node count, edge count, and degree statistics. This is the bar the sheaf models are measured against.
verify_diagnostic.py	Checks that the measurement tool is trustworthy, not the models: gauge invariance, numerical stability, and that flattening a connection really does zero out the readouts.
scaling_experiment.py	Repeats the counting experiment at training-set sizes from 30 up to 3000 graphs, to see how the findings change as the model gets more data.
scaling_baselines.py	Fits the constant predictor and the degree-statistics baseline at each of those training-set sizes, so the models can be compared to the shortcut at every scale.
bigtest_eval.py	Re-checks the main scaling results on 200 freshly generated test graphs instead of the original 8, to make sure the conclusions are not an artefact of a small test set.
*_results.json (11 files)	The saved outputs of the scripts above. The notebook reads these committed results and recomputes its figures and statistics.

The measurement tool itself lives with the library it audits, outside this directory: topobench/nn/backbones/graph/nsd_utils/sheaf_holonomy.py and holonomy_capture.py, with 26 unit tests under test/nn/backbones/graph/.

Reproducing the notebook

From this directory, with the TopoBench Python 3.11 environment active, run:

jupyter nbconvert --to notebook --execute --inplace holonomy_audit.ipynb

The notebook does no training and needs no GPU. Drawing the figures and computing the statistics needs numpy, pandas, scipy, and matplotlib. A few cells re-run the small checks that confirm the holonomy measurement tool is behaving correctly; those import torch and torch_geometric and run fine on CPU. Regenerating the JSON files from scratch is optional and is not needed to read or re-render the notebook.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB