Add per-edge-type value/message transform to Graph Transformer encoder

[zfan3-sc]

Summary

Adds an optional relation-aware message channel to GraphTransformerEncoderLayer, run in parallel with self-attention, controlled by a new relation_message_mode argument. This is the message-passing analog of HGT's per-relation value transform (v_rel) / per-edge-type GATv2 weights, applied sparsely over sampled graph edges.

Three modes:

• "none" (default): behavior identical to the existing layer.
• "edge_type_linear": for every sampled directed edge, adds a per-edge-type linear transform of the source token, mean-aggregated per (target, relation), into the attention residual.
• "edge_type_attention": replaces the mean with a per-head softmax over each target's relation neighbors, scored HGT-style from the layer's own query/key projections with per-(relation, head) transforms and priors.

Key properties:

• Compute scales with the number of relation edges, not seq^2; never materializes an (seq, seq) tensor, so the scaled-dot-product attention path and its memory profile are unaffected.
• Message weights are zero-initialized, so startup behavior exactly matches the plain layer.
• Encoder forward now produces relation indices when either relation attention or relation messages are enabled (previously gated only on relation_attention_mode).

Changes

• gigl/nn/graph_transformer.py: new relation_message_mode arg, parameters, init/reset logic, and _compute_relation_messages / _compute_relation_attention_messages.
• tests/unit/nn/graph_transformer_test.py: coverage for zero-init parity, mean aggregation, uniform-score equivalence to mean mode, score concentration, num_relations validation, and forward smoke tests.

Test plan

• pytest tests/unit/nn/graph_transformer_test.py
• Confirm existing Graph Transformer trainer/inferencer configs (default relation_message_mode="none") are unaffected.

Made with Cursor