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.

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

Divergences from the reference implementation

The core mathematics (attention normalization, gated energy flattening, frame projection, QR steps) matches RiemannGraph/GAUGE. Three deliberate divergences:

1. Similarity scorer. The reference scores each edge with a single Linear(2d -> r) followed by LeakyReLU on the raw embeddings. Here 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. The reference has no residual in the node update. This implementation adds an optional learnable MLP residual phi(z), controlled by phi_hidden_layers; setting it to null recovers the reference behavior exactly. The submitted config enables it.
3. Configurability. Activations are resolved through an activation_dict, and dropout is split into two knobs (feed-forward blocks vs. scorer), where the reference uses fixed activations and a single dropout value.

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

Defaults follow the reference: n_layers=2, n_gated=2, d_embedd=512, r=16, gamma=0.01, tau=1.0, GELU feed-forward / LeakyReLU scorer activations, and dropout 0.1.

Checklist

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

Notice

The code in gauge.py was, to the largest extent, written by hand. 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/.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB