Track: Track1; Team name: Lettlini; Model: GAUGE (Riemannian Graph Foundation Model)

[lettlini]

Summary

This submission implements GAUGE, the gauge-equivariant graph backbone from "Are Common Substructures Transferable: Riemannian Graph Foundation Model with Neural Vector Bundles" (reference implementation: RiemannGraph/GAUGE). Each node is assigned a local orthonormal frame (a gauge) spanning an r-dimensional subspace of the d-dimensional embedding space. These frames are smoothed across the graph and then used to update the node features. The backbone is integrated into TopoBench with a dedicated wrapper and a Hydra config.

Method

The backbone follows the paper's equations directly and is organized into three message-passing components, stacked over n_layers:

• LocalCoordinatesLayer (eqs. 2-4): projects node embeddings into r subspaces, aggregates a neighborhood-smoothed reconstruction under per-subspace attention, and applies a QR decomposition to obtain a per-node orthonormal frame.
• GatedFlatteningLayer (eqs. 5-8): aligns each node's frame with its neighbors' via a trace-based gating weight, blends with coefficient gamma, and re-orthonormalizes with a second QR. Applied n_gated times per layer.
• NodeUpdateLayer (eqs. 9-10): projects each embedding onto its frame (Q^T Q z), maps it through a bias-free linear, and aggregates over the neighborhood.

Dirichlet-energy regularization

The paper's Dirichlet-energy loss (smoothness of frame-projected embeddings across edges) (cf. paper equation (13)) is added to every task's loss via TopoBench's loss-composition mechanism: a loss: node under the backbone in gauge.yaml makes TBLoss add λ · L_Dir automatically, with λ = 0.1 and a configurable mean/sum neighbor reduction.

Files

Component	File
Backbone (all layers + GaugeModel)	topobench/nn/backbones/graph/gauge.py
Wrapper	topobench/nn/wrappers/graph/gauge_wrapper.py
Config	configs/model/graph/gauge.yaml
Unit tests	test/nn/backbones/graph/test_gauge.py
Pipeline test entry	test/pipeline/test_pipeline.py
Dirichlet loss	topobench/loss/model/DirichletLoss.py
Dirichlet loss tests	test/loss/test_dirichlet_loss.py
Benchmark results	2026_tdl_challenge/outputs/2026-07-26_13-33-37/results.json

Divergences from the reference implementation

Where the paper and its reference implementation differ, this implementation aims to follow the equations and formulae as stated in the paper rather than mirror the reference implementation. The core mathematics (attention normalization, gated energy flattening, frame projection, QR steps) matches RiemannGraph/GAUGE. There are, however, two important and deliberate divergences:

1. Similarity scorer (cf. paper eq. 2 & 3). The reference implementation scores each edge with a single Linear(2d -> r) followed by LeakyReLU on the raw embeddings. Here, as indicated in the paper by the $(h)$-superscript, the scorer is a per-subspace MLP applied to the projected multi-path embeddings, with the activation applied to the final score before the softmax (matching the reference's placement). Its activation and dropout are independently configurable (f_sim_act, f_sim_dropout).

2. Optional residual in the node update (cf. paper eq. 10). The reference implementation has no residual ( $\varphi(z)$ ) in the node update. This implementation adds an optional learnable MLP residual $\varphi(z)$, controlled by phi_hidden_layers; setting it to null recovers the reference behavior exactly. The submitted config enables it.

The wrapper's residual connection is disabled (it is not part of the paper), which decouples the backbone embedding width d_embedd from the feature-encoder width.

Configuration defaults

Due to the high computational complexity introduced by the per-subspace MLPs ($f^{(h)}$) and learnable residual $\varphi(z)$ the default model configuration from the paper had to be adjusted as follows:

• in_channels=64 (instead of 128)
• n_layers=2
• n_gated=2
• d_embedd=128 (instead of 512)
• r=16
• gamma=0.01
• tau=1.0
• GELU feed-forward / LeakyReLU scorer activations
• dropout 0.1.
• lamb=0.1 (Weight of the loss produced by DirichletLoss: $\mathcal{L}_\text{total} = \mathcal{L}_\text{task} + \lambda \cdot \mathcal{L}_\text{Dirichlet}$, cf. eq. 14)

Details

Per-subspace layers: MultiHeadLinear and MultiHeadFF

MultiHeadLinear applies r independent linear maps in parallel (one per subspace/head) to an input of shape [N, r, in], giving [N, r, out]. The per-head weights live in one stacked parameter applied via a batched einsum ("roi,Nri->Nro"), so each subspace has its own weights with no sharing. MultiHeadFF stacks these into a per-subspace MLP with activation and dropout between layers.

They implement the similarity scorer f of equation (3), which needs one score per subspace.

(Non-)Influence of $\lambda$

This quick comparison of the results found in the two different versions of results.json

• old, without DirichletLoss
• current, with DirichletLoss and lamb=0.1

shows only very minute differences between the performance metrics. It would be interesting to analyze whether a more thorough sweep of $\lambda$-values would produce more pronounced differences and whether it positively or negatively impacts task performance in these benchmarks.

It has to be noted, however, that this supervised setting is not really how the GaugeModel was intended to be used. A separate pre-training stage might be necessary to fully unlock the potential of this architecture. Still, it will be very interesting to see how it compares to the other novel algorithms.

Notice / Disclaimer

The code in gauge.py was almost entirely written by hand from scratch. The unit tests as well as the docstrings were generated with Claude Code and verified. Inspiration was taken from the reference implementation at https://github.com/RiemannGraph/GAUGE/.

Checklist

• Descriptive PR title following the challenge convention
• Linting passes
• Unit tests added and passing
• NumPy-style docstrings
• Relevant issue linked

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB