Track: Track1; Team name: howyadoin; Model: SheafAN

[dario-loi]

Track

Track 1 — Graph Neural Networks (GNNs)

Team Name

howyadoin

• Dario Loi — Northeastern University
• Gabriele Onorato — Northeastern University

Model

Sheaf Attention Networks (SheafAN)

Status

Ready for review

Summary

This PR contributes a TopoBench-native implementation of Sheaf Attention
Networks (SheafAN) by Barbero et al. (NeurIPS 2022 Workshop on Symmetry
and Geometry in Neural Representations), a generalisation of GAT to
cellular sheaves. Each edge carries both a GAT-style scalar attention
coefficient and a learned d x d restriction map (orthogonal by default);
the per-layer update is the attention-weighted sheaf-adjacency message
pass from equation (5) of the paper, with an optional residual variant
(equation 6 / Res-SheafAN) toggled via the residual flag.

The implementation reuses NSD's sheaf-learner, orthogonal parameterization
and sparse index plumbing where possible, and adds three new pieces
under nsd_utils/: a GAT-style attention module, sheaf-adjacency
builders (one per restriction-map family), and three inductive SheafAN
models mirroring NSD's diag / bundle / general variants. The bundle
variant with O(d) restriction maps is the paper-faithful setting and is
the default in the Hydra config.

Implementation

• Adjacency builders mirroring NSD's Laplacian builders, scaled by
  per-direction attention coefficients with identity self-loops
  (topobench/nn/backbones/graph/nsd_utils/adjacency_builders.py).
• Multi-head GAT-style attention with row-stochastic softmax over
  source neighbours, mirroring torch_geometric.nn.GATConv's
  head/dim flexibility (nsd_utils/san_attention.py).
• Three inductive SheafAN model variants — diag / bundle / general
  — with optional Res-SheafAN residual update
  (nsd_utils/inductive_attention_models.py).
• SANEncoder wrapper plugged into the TBModel pipeline, mirroring
  NSDEncoder's public API (topobench/nn/backbones/graph/san.py).
• Hydra config (configs/model/graph/san.yaml) referencing
  SANEncoder with paper-style defaults
  (sheaf_type=bundle, d=3, num_heads=4).
• Unit tests covering all variants, multi-head attention,
  adjacency builders and residual mode (70 tests,
  test/nn/backbones/graph/test_san.py); coverage on the four new
  files is 100% / 100% / 100% / 98% (well above the 93%
  threshold).
• graph/san added to test/pipeline/test_pipeline.py so the
  end-to-end pipeline is exercised in CI.
• Challenge notebook (2026_tdl_challenge/run_evaluation.ipynb)
  pointed at graph/san and run end-to-end over the full
  GraphUniverse grid (12 settings × 3 seeds × 2 experiments = 72
  runs); the auto-generated results.json is included.

Reference

Barbero, Bodnar, Sáez de Ocáriz Borde, Liò, "Sheaf Attention Networks,"
NeurIPS 2022 Workshop on Symmetry and Geometry in Neural Representations
(Extended Abstract Track).

• OpenReview: https://openreview.net/forum?id=LIDvgVjpkZr
• PDF: https://openreview.net/pdf?id=LIDvgVjpkZr

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB