Track: Track1; Team name: TripleA; Model: Graph Inductive Bias Transformer (GRIT)

[amiiiza]

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. (refer to comment below)
• My PR follows PEP8 guidelines. (refer to comment below)
• 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

This PR adds GRIT (Graph Inductive Bias Transformer) to TopoBench, follwing the paperÖ

Ma, Lin, Lim, Romero-Soriano, Dokania, Coates, Torr, Lim.
"Graph Inductive Biases in Transformers without Message Passing." ICML 2023.
Paper: https://arxiv.org/abs/2305.17589 · Official code: https://github.com/LiamMa/GRIT

The model

GRIT is a graph transformer that injects graph inductive biases without any message-passing modules. It differs from both classical MPNNs and hybrid graph transformers in three ways:

1. RRWP positional encodings (Eq. 1). Each node pair is encoded by its Relative Random Walk Probabilities $\mathbf{P}{i,j} = [\mathbf{I}, \mathbf{M}, \dots, \mathbf{M}^{K-1}]{i,j}$ with $\mathbf{M}=\mathbf{D}^{-1}\mathbf{A}$. These are learned (updated in every layer), and provably more expressive than shortest-path distances (Prop. 3.2). An MLP on RRWP can recover SPD and general propagation matrices (Prop. 3.1).
2. Attention conditioned on node-pair representations (Eq. 2–3). Every layer jointly updates node representations and pair representations, with a signed-square-root stabilizer and an edge-enhanced value aggregation, so the attention pattern itself can learn to emulate $k$-hop propagation.
3. Degree scaler + BatchNorm (Eq. 5, Prop. 3.3). attention is degree-invariant, therefore GRIT re-injects degree information with adaptive scalers and uses BatchNorm because LayerNorm would cancel the degree signal.

Attention is global (all intra-graph node pairs) by default, but convertable to the sparse RRWP support with pad_to_full_graph: false. The default configuration has ~190K parameters (hidden 64, 4 layers, 8 heads, $K=8$).

Files added

Model

File	Description
topobench/nn/backbones/graph/grit.py	The backbone: GRITAttention, GRITTransformerLayer (attention + degree scaler + BatchNorm + FFN), GRITBackbone (RRWP node/pair encoders + layer stack), and the full_edge_index utility.
topobench/transforms/data_manipulations/rrwp_positional_encodings.py	AddRRWP data-manipulation transform + compute_rrwp: precomputes node-level (rrwp) and sparse relative (rrwp_index/rrwp_val) encodings plus log_deg per graph.
topobench/nn/wrappers/graph/grit_wrapper.py	GRITWrapper: thin wrapper forwarding connectivity, RRWP attributes and log-degrees from the batch to the backbone. Returns rank-0 cell embeddings.

Configs

File	Description
configs/model/graph/grit.yaml	Full TBModel composition (AllCellFeatureEncoder → GRITBackbone → GRITWrapper → NoReadOut).
configs/transforms/data_manipulations/rrwp_positional_encodings.yaml	Config for the AddRRWP transform.
configs/transforms/model_defaults/grit.yaml	Automatically applies RRWP precomputation whenever model=graph/grit is selected, interpolates walk_length from the backbone config so transform and model can never drift apart.

Tests

File	Description
test/nn/backbones/graph/test_grit.py	Attention/layer/backbone tests: shapes, batching-vs-individual equivalence, precomputed-vs-on-the-fly RRWP equivalence, error handling, gradients, determinism, Hydra composition & instantiation.
test/transforms/data_manipulations/test_AddRRWP.py	RRWP correctness against explicit $\mathbf{M}^k$ matrix powers, identity channel, sparsity, isolated nodes, empty graphs, batching offsets, registration, determinism.
test/nn/wrappers/graph/test_grit_wrapper.py	Wrapper forward with and without precomputed RRWP.

Modified

File	Description
test/pipeline/test_pipeline.py	Set DATASET = "graph/MUTAG", MODELS = ["graph/grit"], full end-to-end training pipeline test.

Testing

• 42 new unit tests, all passing, deterministic; 100% line coverage on all three contributed source files.
• Pipeline test (graph/grit on graph/MUTAG) passes end-to-end through the real Hydra → PreProcessor → Lightning stack.

Evaluation results (official notebook, results.json included)

2026_tdl_challenge/run_evaluation.ipynb was run unmodified with MODEL_CONFIG = "graph/grit": 12 GraphUniverse settings × 3 seeds = 36 runs per task.

Task	Metric	Result	Params	Avg epoch time
Community detection (node, inductive)	avg test accuracy	0.5207	190,145	8.87 s
Triangle counting (graph, inductive)	avg test MSE / total triangles	0.0693	188,910	8.40 s

Limitations

• Pair representations scale as $O(n^2)$ with pad_to_full_graph: true (as in the paper). Set it to false for sparse RRWP-support attention on larger graphs.
• Datasets with their own dataset_defaults transform config use the on-the-fly RRWP fallback, which is slower than the cached precomputation.

[cursor]

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.