Add incoming sampling direction for graph transformer input contruction & anchor based PE (#671)

Author: yliu2-sc

gigl/nn/graph_transformer.py

@@ -437,6 +437,12 @@ class GraphTransformerEncoder(nn.Module):
         pairwise_attention_bias_attr_names: List of pairwise feature names used
             as additive attention bias. These must correspond to sparse
             graph-level attributes on ``data``.
+        sampling_direction: Direction used for sequence token construction.
+            ``"out"`` preserves the existing k-hop reachability expansion.
+            ``"in"`` expands over reversed edges and is supported only
+            when ``sequence_construction_method="khop"``. Directed relative
+            encodings such as ``"hop_distance"`` should be computed with the
+            same direction.
         feature_embedding_layer_dict: Optional ModuleDict mapping node types to
             feature embedding layers. If provided, these are applied to node
             features before node projection. (default: None)
@@ -504,6 +510,7 @@ def __init__(
         anchor_based_input_attr_names: Optional[list[str]] = None,
         anchor_based_input_embedding_dict: Optional[nn.ModuleDict] = None,
         pairwise_attention_bias_attr_names: Optional[list[str]] = None,
+        sampling_direction: Literal["in", "out"] = "out",
         feature_embedding_layer_dict: Optional[nn.ModuleDict] = None,
         pe_integration_mode: Literal["concat", "add"] = "concat",
         activation: str = "gelu",
@@ -535,6 +542,16 @@ def __init__(
                 "sequence_construction_method must be one of {'khop', 'ppr'}, "
                 f"got '{sequence_construction_method}'"
             )
+        if sampling_direction not in {"in", "out"}:
+            raise ValueError(
+                "sampling_direction must be one of {'in', 'out'}, "
+                f"got '{sampling_direction}'"
+            )
+        if sequence_construction_method == "ppr" and sampling_direction != "out":
+            raise ValueError(
+                "sequence_construction_method='ppr' supports only "
+                "sampling_direction='out'."
+            )
         if sequence_positional_encoding_type is not None:
             sequence_positional_encoding_type = (
                 sequence_positional_encoding_type.lower()
@@ -573,6 +590,7 @@ def __init__(
                 "sequence_construction_method='ppr'."
             )
         self._sequence_construction_method = sequence_construction_method
+        self._sampling_direction = sampling_direction
         self._sequence_positional_encoding_type = sequence_positional_encoding_type
         self._should_l2_normalize_embedding_layer_output = (
             should_l2_normalize_embedding_layer_output
@@ -816,6 +834,7 @@ def forward(
             anchor_node_ids=anchor_node_ids,
             hop_distance=self._hop_distance,
             sequence_construction_method=self._sequence_construction_method,
+            sampling_direction=self._sampling_direction,
             anchor_based_attention_bias_attr_names=self._anchor_based_attention_bias_attr_names,
             anchor_based_input_attr_names=self._anchor_based_input_attr_names,
             pairwise_attention_bias_attr_names=self._pairwise_attention_bias_attr_names,

gigl/transforms/add_positional_encodings.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Literal, Optional
 
 import torch
 from torch_geometric.data import HeteroData
@@ -253,17 +253,30 @@ class AddHeteroHopDistanceEncoding(BaseTransform):
         is_undirected (bool, optional): If set to :obj:`True`, the graph is
             assumed to be undirected for distance computation.
             (default: :obj:`False`)
+        sampling_direction (str, optional): Direction used for directed
+            distance computation. ``"out"`` preserves existing shortest paths
+            over graph edges, while ``"in"`` computes shortest paths over
+            reversed graph edges. Use ``"in"`` to align hop-distance relative
+            encodings with incoming Graph Transformer sampling.
+            (default: :obj:`"out"`)
     """
 
     def __init__(
         self,
         h_max: int,
         attr_name: Optional[str] = "hop_distance",
         is_undirected: bool = False,
+        sampling_direction: Literal["in", "out"] = "out",
     ) -> None:
+        if sampling_direction not in {"in", "out"}:
+            raise ValueError(
+                "sampling_direction must be one of {'in', 'out'}, "
+                f"got '{sampling_direction}'."
+            )
         self.h_max = h_max
         self.attr_name = attr_name
         self.is_undirected = is_undirected
+        self.sampling_direction = sampling_direction
 
     def forward(self, data: HeteroData) -> HeteroData:
         assert isinstance(data, HeteroData), (
@@ -274,6 +287,8 @@ def forward(self, data: HeteroData) -> HeteroData:
         # Convert to homogeneous to compute shortest paths
         homo_data = data.to_homogeneous()
         edge_index = homo_data.edge_index
+        if self.sampling_direction == "in":
+            edge_index = edge_index.flip(0)
         num_nodes = homo_data.num_nodes
         num_edges = edge_index.size(1)
 
@@ -442,4 +457,8 @@ def forward(self, data: HeteroData) -> HeteroData:
         return data
 
     def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(h_max={self.h_max})"
+        return (
+            f"{self.__class__.__name__}("
+            f"h_max={self.h_max}, sampling_direction='{self.sampling_direction}'"
+            ")"
+        )

gigl/transforms/graph_transformer.py

@@ -90,6 +90,7 @@ def heterodata_to_graph_transformer_input(
     anchor_based_attention_bias_attr_names: Optional[list[str]] = None,
     anchor_based_input_attr_names: Optional[list[str]] = None,
     pairwise_attention_bias_attr_names: Optional[list[str]] = None,
+    sampling_direction: Literal["in", "out"] = "out",
 ) -> tuple[Tensor, Tensor, SequenceAuxiliaryData]:
     """
     Transform a HeteroData object to Graph Transformer sequence input.
@@ -131,6 +132,12 @@ def heterodata_to_graph_transformer_input(
         pairwise_attention_bias_attr_names: List of pairwise feature names used
             as attention bias. These must correspond to sparse graph-level
             attributes on ``data``. Example: ['pairwise_distance'].
+        sampling_direction: Direction used for token construction.
+            ``"out"`` preserves the existing k-hop reachability expansion.
+            ``"in"`` expands over reversed edges and is supported only
+            with ``sequence_construction_method="khop"``. Directed relative
+            encodings such as ``"hop_distance"`` should be computed with the
+            same direction.
 
     Returns:
         (sequences, valid_mask, attention_bias_data), where:
@@ -190,6 +197,17 @@ def heterodata_to_graph_transformer_input(
             "be used as pairwise attention bias."
         )
 
+    if sampling_direction not in {"in", "out"}:
+        raise ValueError(
+            "sampling_direction must be one of {'in', 'out'}, "
+            f"got '{sampling_direction}'."
+        )
+
+    if sequence_construction_method == "ppr" and sampling_direction != "out":
+        raise ValueError(
+            "sequence_construction_method='ppr' supports only sampling_direction='out'."
+        )
+
     if (
         PPR_WEIGHT_FEATURE_NAME in anchor_bias_attr_names + anchor_input_attr_names
         and sequence_construction_method != "ppr"
@@ -233,6 +251,8 @@ def heterodata_to_graph_transformer_input(
     ppr_weight_sequences: Optional[Tensor] = None
     if sequence_construction_method == "khop":
         homo_edge_index = homo_data.edge_index  # (2, num_edges)
+        if sampling_direction == "in":
+            homo_edge_index = homo_edge_index.flip(0)
         # Use sparse matrix operations for efficient k-hop neighbor extraction
         # Returns: (batch_size, num_nodes) sparse matrix where non-zero entries are reachable
         reachable = _get_k_hop_neighbors_sparse(

tests/unit/nn/graph_transformer_test.py

@@ -1,6 +1,6 @@
 """Tests for GraphTransformerEncoder."""
 
-from typing import cast
+from typing import Literal, cast
 
 import torch
 import torch.nn as nn
@@ -238,10 +238,24 @@ def test_deterministic_eval_mode(self) -> None:
 
         self.assertTrue(torch.allclose(result_1, result_2))
 
+    def test_forward_with_in_sampling_direction(self) -> None:
+        """Test forward pass with incoming k-hop sampling."""
+        data = _create_simple_hetero_data()
+        encoder = self._create_encoder(sampling_direction="in")
+        encoder.eval()
+
     def test_readout_mode_rejects_invalid_value(self) -> None:
         """Test that unsupported readout modes fail fast."""
         with self.assertRaisesRegex(ValueError, "readout_mode"):
-            self._create_encoder(readout_mode="neighbor_average")
+            self._create_encoder(
+                readout_mode=cast(
+                    Literal[
+                        "anchor_neighbor_attention",
+                        "anchor_only",
+                    ],
+                    "mean_pool",
+                )
+            )
 
     def test_forward_with_anchor_only_readout(self) -> None:
         """Test public forward with anchor-only readout."""
@@ -261,6 +275,17 @@ def test_forward_with_anchor_only_readout(self) -> None:
         self.assertEqual(embeddings.shape, (3, self._out_dim))
         self.assertFalse(torch.isnan(embeddings).any())
 
+    def test_sampling_direction_rejects_invalid_value(self) -> None:
+        with self.assertRaisesRegex(ValueError, "sampling_direction"):
+            self._create_encoder(sampling_direction="sideways")
+
+    def test_in_sampling_direction_requires_khop(self) -> None:
+        with self.assertRaisesRegex(ValueError, "supports only"):
+            self._create_encoder(
+                sequence_construction_method="ppr",
+                sampling_direction="in",
+            )
+
     def test_anchor_only_readout_returns_anchor_token(self) -> None:
         """Test anchor-only readout returns the post-norm anchor token."""
         encoder = self._create_encoder(

tests/unit/transforms/add_positional_encodings_test.py

@@ -1,3 +1,5 @@
+from typing import Literal, cast
+
 import torch
 from absl.testing import absltest
 from torch_geometric.data import HeteroData
@@ -50,6 +52,19 @@ def create_empty_hetero_data() -> HeteroData:
     return data
 
 
+def create_directed_chain_data() -> HeteroData:
+    """Create a directed chain 0 -> 1 -> 2 for direction tests."""
+    data = HeteroData()
+    data["user"].x = torch.randn(3, 4)
+    data["user", "to", "user"].edge_index = torch.tensor(
+        [
+            [0, 1],
+            [1, 2],
+        ]
+    )
+    return data
+
+
 class TestAddHeteroRandomWalkEncodings(TestCase):
     """Tests for AddHeteroRandomWalkEncodings (consolidated PE and SE in single pass)."""
 
@@ -267,6 +282,45 @@ def test_forward_undirected(self):
         self.assertTrue(result.hop_distance.is_sparse_csr)
         self.assertEqual(result.hop_distance.shape, (5, 5))
 
+    def test_forward_sampling_direction_defaults_to_out(self):
+        """Out hop distances preserve existing directed reachability."""
+        data = create_directed_chain_data()
+        transform = AddHeteroHopDistanceEncoding(h_max=2)
+
+        result = transform(data)
+        dense = result.hop_distance.to_dense()
+
+        self.assertEqual(dense[0, 1].item(), 1.0)
+        self.assertEqual(dense[0, 2].item(), 2.0)
+        self.assertEqual(dense[2, 1].item(), 0.0)
+        self.assertEqual(dense[2, 0].item(), 0.0)
+
+    def test_forward_sampling_direction_in_reverses_reachability(self):
+        """In hop distances are computed over reversed graph edges."""
+        data = create_directed_chain_data()
+        transform = AddHeteroHopDistanceEncoding(
+            h_max=2,
+            sampling_direction="in",
+        )
+
+        result = transform(data)
+        dense = result.hop_distance.to_dense()
+
+        self.assertEqual(dense[2, 1].item(), 1.0)
+        self.assertEqual(dense[2, 0].item(), 2.0)
+        self.assertEqual(dense[0, 1].item(), 0.0)
+        self.assertEqual(dense[0, 2].item(), 0.0)
+
+    def test_sampling_direction_rejects_invalid_value(self):
+        with self.assertRaisesRegex(ValueError, "sampling_direction"):
+            AddHeteroHopDistanceEncoding(
+                h_max=2,
+                sampling_direction=cast(
+                    Literal["in", "out"],
+                    "sideways",
+                ),
+            )
+
     def test_forward_empty_graph(self):
         """Test forward pass with empty graph."""
         data = create_empty_hetero_data()
@@ -283,7 +337,21 @@ def test_forward_empty_graph(self):
     def test_repr(self):
         """Test string representation."""
         transform = AddHeteroHopDistanceEncoding(h_max=5)
-        self.assertEqual(repr(transform), "AddHeteroHopDistanceEncoding(h_max=5)")
+        self.assertEqual(
+            repr(transform),
+            "AddHeteroHopDistanceEncoding(h_max=5, sampling_direction='out')",
+        )
+
+    def test_repr_in_sampling_direction(self):
+        """Test string representation with non-default direction."""
+        transform = AddHeteroHopDistanceEncoding(
+            h_max=5,
+            sampling_direction="in",
+        )
+        self.assertEqual(
+            repr(transform),
+            "AddHeteroHopDistanceEncoding(h_max=5, sampling_direction='in')",
+        )
 
 
 if __name__ == "__main__":