Track: Track2; Team name: TopoCoop; Model: Cellular Transformer - #371
Open
magoflaco wants to merge 6 commits into
Open
Track: Track2; Team name: TopoCoop; Model: Cellular Transformer#371magoflaco wants to merge 6 commits into
magoflaco wants to merge 6 commits into
Conversation
Implements the Cellular Transformer (Ballester et al., 2024, arXiv:2405.14094) as a cell-domain backbone with: - Sparse pairwise cellular attention (Eq. (2)) over the neighborhood structure of the complex: rank-0 upper adjacency, rank-1/2 lower adjacencies and cross-rank incidences (Eq. (3), Section 4.2.1) - Prenorm transformer layers with per-rank LayerNorm, residual aggregation over attention routes and feed-forward blocks (Appendix A) - ConcatPE preprocessing (Eq. (4)) with random walk positional encodings (RWPe, Appendix C.1) and a zero-PE ablation variant - CellularTransformerWrapper and Hydra configs for the CT(A^s, RWPe) and CT(A^s, zero) variants - Unit tests with 100% coverage of the new modules and pipeline test No public reference implementation exists; the architecture is implemented from the paper equations. Dense attention variants are omitted as they do not scale under block-diagonal batching (modification allowed to respect computational requirements). Co-authored-by: marcoschmmath <171987457+marcoschmmath@users.noreply.github.com>
Produced by 2026_tdl_challenge/run_evaluation.ipynb with MODEL_CONFIG=cell/cellular_transformer on a Kaggle P100 GPU (72 runs: 12 grid settings x 3 seeds x 2 tasks, full OOD evaluation, graphs lifted to cell complexes via the cycle lifting; 8.9h wall time). Co-authored-by: marcoschmmath <171987457+marcoschmmath@users.noreply.github.com>
- Document the ConcatPE interpretation (single joint projection per rank, vs SumPE) and the on-the-fly RWPe computation rationale - Add invariance tests: masked-softmax normalization (attention output is a convex combination of value vectors) and block-diagonal batching equivalence (no cross-complex attention leakage) - Drop an unnecessary tensor clone in the layer residual path - Remove the unused pe_steps entry from the zero-PE variant config Co-authored-by: marcoschmmath <171987457+marcoschmmath@users.noreply.github.com>
|
Bugbot is not enabled for this team, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
levtelyatnikov
added
the
track-2-tnn
…e matmuls Issue: CellularTransformer.forward() called random_walk_pe() on every single forward pass — every batch, every epoch — even though adjacency_0, coadjacency_1 and coadjacency_2 are fixed per complex. Each call performs pe_steps sequential sparse-sparse matrix products; by step 6–8 the intermediates are near-dense, making this a major hot-path cost. Fix: move RWPE computation into a one-time data transform (CellRandomWalkPE) that runs during PreProcessor caching. The computed tensors (rwpe_0, rwpe_1, rwpe_2) are attached to each Data object, cached to disk by TopoBench's existing preprocessor, and batched alongside the complexes. The backbone now reads them directly instead of recomputing. Changes: 1. New transform: topobench/transforms/data_manipulations/cell_random_walk_pe.py 2. Hydra config: configs/transforms/data_manipulations/cell_random_walk_pe.yaml 3. Register in experiment: configs/experiment/hopse_g_gnn_cell.yaml 4. Thread through wrapper: CellularTransformerWrapper extracts rwpe_* and passes them to the backbone 5. Read in backbone: CellularTransformer.forward uses precomputed PE when available, with an on-the-fly fallback for backward compatibility Impact: RWPE is now computed once per graph (pre-training) instead of once per batch per epoch. This is likely the single biggest wall-clock win for the cellular transformer backbone.
Corrects the previous commit so the optimization actually takes effect in the TopoBench pipeline: - Move CellRandomWalkPE from transforms/feature_liftings/ to transforms/data_manipulations/ following the manipulation conventions (BaseTransform.forward, auto-registration in TRANSFORMS) - Rewrite its config with the transform_name/transform_type schema consumed by DataTransform (the previous _target_-style config was never picked up by the pipeline) - Wire the transform for this model via configs/transforms/model_defaults/cellular_transformer.yaml (lifting + manipulation, SANN pattern) so get_default_transform applies it automatically, including in the challenge evaluation; pe_steps stays in sync with the backbone via config interpolation - Revert the unrelated configs/experiment/hopse_g_gnn_cell.yaml edit (out of scope for this PR and incompatible with the HOPSE backbone) - Add tests: CellRandomWalkPE unit tests and a backbone equivalence test pinning precomputed == on-the-fly encodings (100% coverage of backbone, wrapper and transform) - Apply repo formatting (ruff) and file hygiene fixes Co-authored-by: marcoschmmath <171987457+marcoschmmath@users.noreply.github.com>
Re-run of 2026_tdl_challenge/run_evaluation.ipynb on a Kaggle P100 with MODEL_CONFIG=cell/cellular_transformer at the current pipeline (CellRandomWalkPE precompute transform active). Scientific metrics are statistically identical to the previous run (max per-run accuracy delta 0.004, GPU nondeterminism); computational-complexity metrics now reflect the shipped code: 5.0 s/epoch mean (was 7.9), 6.86 h total (was 8.93 h), 465,883 parameters. Co-authored-by: marcoschmmath <171987457+marcoschmmath@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Model
This PR implements the Cellular Transformer (CT) as a Track 2 (TNN) backbone for the TDL Challenge 2026.
The Cellular Transformer generalizes transformers to cell complexes: cells of ranks 0/1/2 exchange information through pairwise cellular attention (Eq. (2)) restricted to the neighborhood structure of the complex — upper adjacency within rank 0, lower adjacencies within ranks 1–2, and non-signed incidences across adjacent ranks (Eq. (3), Section 4.2.1) — with prenorm transformer layers (Appendix A) and topological positional encodings combined via ConcatPE (Eq. (4)).
There is no public reference implementation of this paper; the architecture was implemented from the paper's equations. All docstrings reference the specific equations/sections they implement.
What is included
topobench/nn/backbones/cell/cellular_transformer.pytopobench/nn/wrappers/cell/cellular_transformer_wrapper.pyconfigs/model/cell/cellular_transformer.yamlconfigs/model/cell/cellular_transformer_zero_pe.yamltest/nn/backbones/cell/test_cellular_transformer.pytest/pipeline/test_pipeline.py2026_tdl_challenge/results.jsonImplementation notes (modifications allowed to respect TopoBench computational requirements):
AllCellFeatureEncoderand standard readout (global add pooling over rank-0 signals, matching the paper's experimental setup).Evaluation
2026_tdl_challenge/run_evaluation.ipynbwas run on a Kaggle P100 GPU withMODEL_CONFIG = "cell/cellular_transformer"(72 runs: 12 grid settings × 3 seeds × 2 tasks, full OOD evaluation, graphs lifted to cell complexes via the cycle lifting; 6.9 h wall time; 465,883 parameters, ≈5.0 s/epoch). The generatedresults.jsonis committed.Highlights: community-detection accuracy averages 0.64 under high homophily and 0.28 under severe heterophily (20 classes). On triangle counting — the task probing higher-order structure — the CT reaches a mean MSE/triangles of 0.51, roughly half of what we measured for a strong Track 1 GNN baseline (Co-GNN, 0.93) under the identical pipeline: direct evidence that rank-2 cells give the model explicit access to the cycle structure the task asks about.
Note on the notebook integrity check: current
mainships a staleexpected_hashthat fails on a clean clone (#330). Only theexpected_hashcell (above the# UNIQUE_HASH_MARKER, i.e. in the user-editable region) was updated to the SHA-256 of the untouched protected cells; all cells from the marker onward are byte-identical tomain.Tests
test/nn/backbones/cell/test_cellular_transformer.py: 11 tests covering the backbone (both PE variants, empty rank-2 complexes, gradient flow through all attention routes, masked-attention semantics, RWPe correctness on a known graph, row normalization, invalid-argument handling); 100% line coverage of both new modules.test/pipeline/test_pipeline.py: end-to-end training ongraph/MUTAGof the upstream models plus both CT variants (with graph→cell cycle lifting).Result figures
In-distribution test performance across the 12 GraphUniverse grid settings (mean ± std over 3 seeds):
Community detection (accuracy):
Triangle counting (MSE / total triangles):
OOD delta plots (OOD − in-distribution), by training homophily level