Track: Track1; Team name: SweetLesson; Model: Polynormer

[AaravG42]

Checklist

• My pull request has a clear and explanatory title.
• My pull request passes the Linting test.
• I added appropriate unit tests and I made sure the code passes all unit tests.
• My PR follows PEP8 guidelines.
• My code is properly documented, using numpy docs conventions, and I made sure the documentation renders properly.
• I linked to issues and PRs that are relevant to this PR.

Description

Track 1 (GNNs) submission — Team SweetLesson — Model: Polynormer.

This PR integrates Polynormer into TopoBench:

Chenhui Deng, Zichao Yue, Zhiru Zhang. Polynormer: Polynomial-Expressive Graph Transformer in Linear Time. ICLR 2024. arXiv:2403.01232 · official code: cornell-zhang/polynormer.

Polynormer learns a high-degree equivariant polynomial on the node features whose coefficients are produced by attention, composed of a local GAT-style equivariant attention (Eq. 7) followed by a global linear (kernel) attention in O(N·d²) time (Eq. 6 & 8). Stacking L local layers reaches degree-2^L expressivity (Thm. 3.3).

What's added

• Backbone — topobench/nn/backbones/graph/polynormer.py: Polynormer and PolynormerAttention. Faithful to the official model.py; every block's docstring cites the corresponding paper equation and the reference implementation.
• Config — configs/model/graph/polynormer.yaml: reuses GNNWrapper + NoReadOut, so a single config serves both challenge tasks (node-level community detection and graph-level triangle counting).
• Unit tests — test/nn/backbones/graph/test_polynormer.py: 100% line coverage of the backbone, including a batch-isolation correctness test (below).
• Pipeline test — graph/polynormer added to test/pipeline/test_pipeline.py.
• Tutorial — tutorials/tutorial_polynormer.ipynb: walks through the local→global structure, the batch-aware property, and an end-to-end MUTAG run.
• Results — 2026_tdl_challenge/outputs/<study>/results.json from the official GraphUniverse grid. (The notebook also renders in-distribution and OOD heatmaps; per the repo's .gitignore, only results.json under outputs/ is committed.)

Adaptations to TopoBench (correctness notes)

1. Batch-aware global attention. The reference runs on a single graph; TopoBench feeds mini-batches of disjoint graphs. The global linear attention is therefore made batch-aware: the kernel sums σ(K)ᵀV and Σ σ(Kᵢ) are accumulated per graph segment (via torch_geometric.utils.scatter over the batch vector), so nodes never attend across graph boundaries — essential for graph-level triangle counting. This reduces exactly to Eq. 6/8 for a single graph and preserves the O(N·d²) linear-in-N complexity. A unit test verifies a graph's embeddings are identical whether it is run alone or inside a batch.
2. Embeddings, not class logits. The paper's pred_local/pred_global task heads are replaced by a single output projection to node embeddings; the TopoBench readout produces the final logits.
3. Joint local+global training. The reference toggles a _global flag mid-training (local warm-up → global). TopoBench runs a single Lightning loop, so the two modules are trained jointly (global_layers=0 recovers the faithful local-only variant).

Complexity / scalability

Local layers are O(E·d) (sparse GAT); the global module is O(N·d²) time and O(N·d²) memory (linear in N), versus O(N²·d) for dense attention — the paper's central efficiency claim, preserved here. The benchmarked config (in=hidden=out=64, heads=1, local_layers=3, global_layers=2) has 76,160 trainable parameters.

Results summary (GraphUniverse grid, 72 runs over 3 seeds)

• Community detection (node, accuracy): 0.30–0.73 in-distribution across the 12 settings — well above the multi-community random baseline, and higher under homophily, as expected.
• Triangle counting (graph, normalized MSE / total triangles): 0.012–3.76 in-distribution; all finite.
• Full per-setting, per-seed, and OOD (each model evaluated on the 11 other settings) results are in results.json. The notebook additionally renders in-distribution and OOD heatmaps locally (not committed — the repo's .gitignore keeps only results.json under outputs/).
• Note: the optional wandb_config timing/param fields are empty because the grid was run with WANDB_MODE=offline (no W&B account); utils.py reads those from online run-* dirs only. The model's parameter count is reported above.

Note on results.json generation

On upstream main, 2026_tdl_challenge/run_evaluation.ipynb carries a self-integrity hash (expected_hash = f87b2c…) that does not match the hash of its own shipped cells (hash_remaining_cells(...) → 3c1d78…), so the guard cell raises ValueError before any work — for the unmodified notebook. To produce the required artifact without modifying the notebook or utils.py, I invoked the notebook's own backend directly — run_challenge_grid(...) + save_challenge_artifacts(...) from 2026_tdl_challenge/utils.py — which runs the identical pipeline. Happy to open a separate issue/PR to refresh the stored hash for maintainers.

Issue

Submission to the TDL Challenge 2026, Track 1 (GNNs). No existing PR implements Polynormer.

Additional context

Tested with the project environment (Python 3.11, torch 2.3.0+cu121). Unit tests, the pipeline test, the tutorial notebook, and a 1-epoch GraphUniverse sanity over all 12 settings × both tasks all pass locally.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB