Track: Track1; Team name: LangDiff; Model: AdvDIFFormer

[Mullerio]

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.

Description

This PR adds an implementation of AdvDIFFormer (Advective Diffusion Transformer) from Wu et al., “Supercharging Graph Transformers with Advective Diffusion” https://arxiv.org/abs/2310.06417.

AdvDIFFormer is a PDE inspired graph Transformer derived from the advective diffusion equations. It combines graph-structured propagation with learned non-local interactions.

This PR also adds, as per challenge a example config:

configs/model/graph/advdifformer.yaml

And some unit tests in:

test/nn/backbones/graph/test_advdifformer.py

Testing was run locally:

uv run --no-sync ruff check topobench/nn/backbones/graph/advdifformer.py test/nn/backbones/graph/test_advdifformer.py test/pipeline/test_pipeline.py

uv run --no-sync pytest test/nn/backbones/graph/test_advdifformer.py -q

uv run --no-sync pytest test/pipeline/test_pipeline.py -q

Additional context

The implementation supports both AdvDIFFormer variants: the more scalable AdvDIFFormer-S formulation and the AdvDIFFormer-I (inverse) formulation. For more details, we refer to the original paper.

For AdvDIFFormer-S, the implementation supports batched PyG graphs while keeping global attention graph-local and avoiding building of dense node-to-node attention matrices. Propagation across multiple heads is batched. The normalized graph structure and batch segmentation are prepared once forward and reused across layers. This improved speed in local testing.

Propagation-order projections are accumulated incrementally rather than by materializing a concatenation of all intermediate propagation states. In line with the original implementation, this implemementation supports weighted edges and multi-head propagation. Additionally, it supports optional graph symmetrization, arbitrary PyG batch assignments and layer normalization. These can be configured through the yaml as in the example.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB