From 57eb72f30ad318f60c0558bd9b697762b9cd3884 Mon Sep 17 00:00:00 2001 From: Yozen Liu Date: Mon, 15 Jun 2026 10:40:00 -0700 Subject: [PATCH] Add sparse edge attribute bias to graph transformer --- gigl/nn/graph_transformer.py | 147 ++++++++- gigl/transforms/graph_transformer.py | 297 +++++++++++++++++- tests/unit/nn/graph_transformer_test.py | 145 +++++++++ .../unit/transforms/graph_transformer_test.py | 160 ++++++++++ 4 files changed, 743 insertions(+), 6 deletions(-) diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index f9d8b345a..c574ee699 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -376,9 +376,8 @@ class GraphTransformerEncoder(nn.Module): node_type_to_feat_dim_map: Dictionary mapping node types to their input feature dimensions. edge_type_to_feat_dim_map: Dictionary mapping edge types to their - feature dimensions. Accepted for interface conformance with - ``HGT``/``SimpleHGN``; edge features are not used by the - graph transformer. + feature dimensions. Used only when + ``edge_attr_attention_bias_mode="sparse_linear"``. hid_dim: Hidden dimension for transformer layers. All node types are projected to this dimension before processing. out_dim: Output embedding dimension. @@ -432,6 +431,11 @@ 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``. + edge_attr_attention_bias_mode: How to use sampled edge attributes as + attention-logit bias. ``"none"`` preserves the default behavior. + ``"sparse_linear"`` keeps edge attrs sparse, projects each matched + edge attr row to per-head bias, and accumulates it at + ``query=dst, key=src`` for the stored edge 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) @@ -495,6 +499,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, + edge_attr_attention_bias_mode: Literal["none", "sparse_linear"] = "none", feature_embedding_layer_dict: Optional[nn.ModuleDict] = None, pe_integration_mode: Literal["concat", "add"] = "concat", activation: str = "gelu", @@ -509,6 +514,12 @@ def __init__( "pe_integration_mode must be one of {'concat', 'add'}, " f"got '{pe_integration_mode}'" ) + if edge_attr_attention_bias_mode not in {"none", "sparse_linear"}: + raise ValueError( + "edge_attr_attention_bias_mode must be one of " + "{'none', 'sparse_linear'}, " + f"got '{edge_attr_attention_bias_mode}'" + ) self._hid_dim = hid_dim self._out_dim = out_dim @@ -571,6 +582,20 @@ def __init__( self._feature_embedding_layer_dict = feature_embedding_layer_dict self._pe_integration_mode = pe_integration_mode self._num_heads = num_heads + self._edge_type_to_feat_dim_map = { + edge_type: edge_type_to_feat_dim_map[edge_type] + for edge_type in sorted(edge_type_to_feat_dim_map.keys()) + } + self._edge_attr_attention_bias_mode = edge_attr_attention_bias_mode + self._edge_attr_attention_bias_edge_types = ( + [ + edge_type + for edge_type, edge_attr_dim in self._edge_type_to_feat_dim_map.items() + if int(edge_attr_dim) > 0 + ] + if edge_attr_attention_bias_mode == "sparse_linear" + else [] + ) anchor_input_embedding_attr_names = ( set(anchor_based_input_embedding_dict.keys()) if anchor_based_input_embedding_dict is not None @@ -647,6 +672,20 @@ def __init__( bias=False, ) + self._edge_attr_attention_bias_projection_dict = nn.ModuleDict() + for relation_idx, edge_type in enumerate( + self._edge_attr_attention_bias_edge_types + ): + edge_attr_dim = int(self._edge_type_to_feat_dim_map[edge_type]) + if edge_attr_dim <= 0: + continue + projection = nn.Linear(edge_attr_dim, num_heads, bias=True) + nn.init.zeros_(projection.weight) + nn.init.zeros_(projection.bias) + self._edge_attr_attention_bias_projection_dict[str(relation_idx)] = ( + projection + ) + # Transformer encoder layers # Default feedforward ratio: 4.0 for standard activations, 8/3 for XGLU # XGLU's gating mechanism doubles effective parameters, so smaller ratio @@ -801,6 +840,14 @@ def forward( 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, + edge_attr_edge_type_to_feat_dim_map=( + { + edge_type: self._edge_type_to_feat_dim_map[edge_type] + for edge_type in self._edge_attr_attention_bias_edge_types + } + if self._edge_attr_attention_bias_mode == "sparse_linear" + else None + ), ) # Free memory after sequences are built @@ -958,6 +1005,11 @@ def _build_attention_bias( Input shape: (batch, seq, seq, num_pairwise_attrs) After projection: (batch, num_heads, seq, seq) - unique bias per query-key pair. + 4. **Sparse edge-attribute bias** (optional): For each matched sampled + edge attr row, project to per-head bias and accumulate at + ``(batch, head, query_pos=dst, key_pos=src)`` without materializing + dense edge-feature tensors. + Args: valid_mask: Boolean mask of shape (batch_size, seq_len) indicating valid (non-padding) positions. @@ -966,6 +1018,8 @@ def _build_attention_bias( attention_bias_data: Dictionary containing optional PE tensors: - "anchor_bias": (batch, seq, num_anchor_attrs) or None - "pairwise_bias": (batch, seq, seq, num_pairwise_attrs) or None + - "pairwise_edge_attr_indices": dict[int, (n, 3)] or None + - "pairwise_edge_attr_values": dict[int, (n, edge_dim)] or None Returns: Combined attention bias tensor of shape (batch_size, num_heads, seq_len, seq_len) @@ -1029,6 +1083,93 @@ def _build_attention_bias( ) # (batch, num_heads, seq, seq) attn_bias = attn_bias + pairwise_bias + pairwise_edge_attr_indices = attention_bias_data.get( + "pairwise_edge_attr_indices" + ) + pairwise_edge_attr_values = attention_bias_data.get("pairwise_edge_attr_values") + if ( + pairwise_edge_attr_indices is not None + or pairwise_edge_attr_values is not None + ): + if self._edge_attr_attention_bias_mode != "sparse_linear": + raise ValueError( + "Sparse edge-attribute attention-bias payloads require " + "edge_attr_attention_bias_mode='sparse_linear'." + ) + if pairwise_edge_attr_indices is None or pairwise_edge_attr_values is None: + raise ValueError( + "Both pairwise_edge_attr_indices and pairwise_edge_attr_values " + "must be provided together." + ) + if set(pairwise_edge_attr_indices.keys()) != set( + pairwise_edge_attr_values.keys() + ): + raise ValueError( + "pairwise_edge_attr_indices and pairwise_edge_attr_values " + "must contain the same relation indices." + ) + + attn_bias_by_position: Optional[Tensor] = None + for relation_idx in sorted(pairwise_edge_attr_indices.keys()): + edge_attr_indices = pairwise_edge_attr_indices[relation_idx] + edge_attr_values = pairwise_edge_attr_values[relation_idx] + if edge_attr_indices.dim() != 2 or edge_attr_indices.size(-1) != 3: + raise ValueError( + "pairwise_edge_attr_indices values must have shape " + f"(num_edges, 3), got {tuple(edge_attr_indices.shape)}." + ) + if edge_attr_values.dim() != 2: + raise ValueError( + "pairwise_edge_attr_values values must have shape " + f"(num_edges, edge_dim), got {tuple(edge_attr_values.shape)}." + ) + if edge_attr_indices.size(0) != edge_attr_values.size(0): + raise ValueError( + "pairwise_edge_attr_indices and pairwise_edge_attr_values " + "must have the same number of rows for relation " + f"{relation_idx}." + ) + if edge_attr_indices.numel() == 0: + continue + + projection_key = str(relation_idx) + if projection_key not in self._edge_attr_attention_bias_projection_dict: + raise ValueError( + "No sparse edge-attribute attention-bias projection found " + f"for relation index {relation_idx}." + ) + projection = self._edge_attr_attention_bias_projection_dict[ + projection_key + ] + + if attn_bias_by_position is None: + full_attn_bias_shape = ( + batch_size, + self._num_heads, + seq_len, + seq_len, + ) + if tuple(attn_bias.shape) != full_attn_bias_shape: + attn_bias = attn_bias.expand(full_attn_bias_shape).clone() + attn_bias_by_position = attn_bias.permute(0, 2, 3, 1) + + edge_attr_indices = edge_attr_indices.to( + device=device, + dtype=torch.long, + ) + edge_attr_bias = projection( + edge_attr_values.to(device=device, dtype=dtype) + ) + attn_bias_by_position.index_put_( + ( + edge_attr_indices[:, 0], + edge_attr_indices[:, 1], + edge_attr_indices[:, 2], + ), + edge_attr_bias.to(dtype=attn_bias.dtype), + accumulate=True, + ) + return attn_bias def _encode_and_readout( diff --git a/gigl/transforms/graph_transformer.py b/gigl/transforms/graph_transformer.py index 602f95bde..b071a10a7 100644 --- a/gigl/transforms/graph_transformer.py +++ b/gigl/transforms/graph_transformer.py @@ -57,7 +57,7 @@ >>> # attention_bias_data['anchor_bias']: (batch_size, max_seq_len, 1) """ -from typing import Literal, Optional, TypedDict +from typing import Literal, NamedTuple, Optional, TypedDict import torch from torch import Tensor @@ -65,12 +65,16 @@ from torch_geometric.typing import NodeType from torch_geometric.utils import to_torch_sparse_tensor +from gigl.src.common.types.graph_data import EdgeType as GiGLEdgeType + TokenInputData = dict[str, Tensor] class SequenceAuxiliaryData(TypedDict): anchor_bias: Optional[Tensor] pairwise_bias: Optional[Tensor] + pairwise_edge_attr_indices: Optional[dict[int, Tensor]] + pairwise_edge_attr_values: Optional[dict[int, Tensor]] token_input: Optional[TokenInputData] @@ -90,6 +94,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, + edge_attr_edge_type_to_feat_dim_map: Optional[dict[GiGLEdgeType, int]] = None, ) -> tuple[Tensor, Tensor, SequenceAuxiliaryData]: """ Transform a HeteroData object to Graph Transformer sequence input. @@ -131,6 +136,10 @@ 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']. + edge_attr_edge_type_to_feat_dim_map: Optional edge-type feature dims + for sparse edge-attribute attention bias extraction. Positive dims + require ``edge_attr`` on the corresponding edge stores. Zero-dim + edge types are skipped. Returns: (sequences, valid_mask, attention_bias_data), where: @@ -143,6 +152,11 @@ def heterodata_to_graph_transformer_input( ``"anchor_bias"`` shaped ``(batch, seq, num_anchor_attrs)`` or None ``"pairwise_bias"`` shaped ``(batch, seq, seq, num_pairwise_attrs)`` or None + ``"pairwise_edge_attr_indices"`` as a dict mapping sorted + edge-type relation index to ``(num_matches, 3)`` long tensors + containing ``(batch_idx, query_pos, key_pos)`` or None + ``"pairwise_edge_attr_values"`` as a dict mapping sorted + edge-type relation index to raw matched edge attributes or None ``"token_input"`` as a dict mapping attribute name to a ``(batch, seq, 1)`` tensor, or None @@ -204,8 +218,12 @@ def heterodata_to_graph_transformer_input( device = data[anchor_node_type].x.device - # Convert to homogeneous for easier neighborhood extraction - homo_data = data.to_homogeneous() + # Convert to homogeneous for easier neighborhood extraction. K-hop mode does + # not need homogeneous edge attributes, and hetero edge attrs can have + # relation-specific widths, so keep those sparse on the original stores. + homo_data = data.to_homogeneous( + edge_attrs=[] if sequence_construction_method == "khop" else None + ) homo_x = homo_data.x # (total_nodes, feature_dim) num_nodes = homo_data.num_nodes @@ -312,6 +330,27 @@ def heterodata_to_graph_transformer_input( csr_matrices=pairwise_pe_matrices if pairwise_pe_matrices else None, device=device, ) + token_occurrences = None + if edge_attr_edge_type_to_feat_dim_map: + token_occurrences = _build_token_occurrence_index( + node_index_sequences=node_index_sequences, + valid_mask=valid_mask, + num_nodes=num_nodes, + device=device, + ) + ( + pairwise_edge_attr_indices, + pairwise_edge_attr_values, + ) = _lookup_pairwise_edge_attr_payloads( + data=data, + node_index_sequences=node_index_sequences, + valid_mask=valid_mask, + edge_attr_edge_type_to_feat_dim_map=edge_attr_edge_type_to_feat_dim_map, + node_type_offsets=node_type_offsets, + num_nodes=num_nodes, + device=device, + token_occurrences=token_occurrences, + ) anchor_bias_features = _compose_anchor_feature_tensor( anchor_relative_feature_sequences=anchor_relative_feature_sequences, @@ -332,6 +371,8 @@ def heterodata_to_graph_transformer_input( { "anchor_bias": anchor_bias_features, "pairwise_bias": pairwise_feature_sequences, + "pairwise_edge_attr_indices": pairwise_edge_attr_indices, + "pairwise_edge_attr_values": pairwise_edge_attr_values, "token_input": token_input_features, }, ) @@ -875,6 +916,256 @@ def _lookup_pairwise_relative_features( return features +class _TokenOccurrenceIndex(NamedTuple): + batch_indices: Tensor + positions: Tensor + node_indices: Tensor + sorted_node_indices: Tensor + node_sort_perm: Tensor + sorted_batch_node_keys: Tensor + batch_node_sort_perm: Tensor + + +def _build_token_occurrence_index( + node_index_sequences: Tensor, + valid_mask: Tensor, + num_nodes: int, + device: torch.device, +) -> _TokenOccurrenceIndex: + """Index valid sequence token occurrences for sparse edge matching.""" + batch_indices, positions = torch.nonzero(valid_mask, as_tuple=True) + node_indices = node_index_sequences[batch_indices, positions] + + if node_indices.numel() == 0: + empty = torch.empty(0, dtype=torch.long, device=device) + return _TokenOccurrenceIndex( + batch_indices=empty, + positions=empty, + node_indices=empty, + sorted_node_indices=empty, + node_sort_perm=empty, + sorted_batch_node_keys=empty, + batch_node_sort_perm=empty, + ) + + node_sort_perm = torch.argsort(node_indices, stable=True) + sorted_node_indices = node_indices[node_sort_perm] + + batch_node_keys = batch_indices * num_nodes + node_indices + batch_node_sort_perm = torch.argsort(batch_node_keys, stable=True) + sorted_batch_node_keys = batch_node_keys[batch_node_sort_perm] + + return _TokenOccurrenceIndex( + batch_indices=batch_indices, + positions=positions, + node_indices=node_indices, + sorted_node_indices=sorted_node_indices, + node_sort_perm=node_sort_perm, + sorted_batch_node_keys=sorted_batch_node_keys, + batch_node_sort_perm=batch_node_sort_perm, + ) + + +def _match_directed_edges_to_token_pairs( + source_indices: Tensor, + target_indices: Tensor, + token_occurrences: _TokenOccurrenceIndex, + num_nodes: int, + device: torch.device, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Match stored ``source -> target`` edges to ``query=target, key=source``.""" + if source_indices.numel() == 0 or token_occurrences.node_indices.numel() == 0: + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty, empty, empty + + source_indices = source_indices.to(device=device, dtype=torch.long) + target_indices = target_indices.to(device=device, dtype=torch.long) + edge_indices = torch.arange(source_indices.size(0), dtype=torch.long, device=device) + + target_left = torch.searchsorted( + token_occurrences.sorted_node_indices, + target_indices, + right=False, + ) + target_right = torch.searchsorted( + token_occurrences.sorted_node_indices, + target_indices, + right=True, + ) + target_counts = target_right - target_left + has_target = target_counts > 0 + if not has_target.any(): + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty, empty, empty + + target_edge_indices = edge_indices[has_target] + target_counts = target_counts[has_target] + target_left = target_left[has_target] + + target_total = int(target_counts.sum().item()) + target_group_ids = torch.repeat_interleave( + torch.arange(target_edge_indices.size(0), dtype=torch.long, device=device), + target_counts, + ) + target_group_starts = torch.cumsum(target_counts, dim=0) - target_counts + target_offsets = ( + torch.arange(target_total, dtype=torch.long, device=device) + - target_group_starts[target_group_ids] + ) + target_token_perm = token_occurrences.node_sort_perm[ + target_left[target_group_ids] + target_offsets + ] + target_batch_indices = token_occurrences.batch_indices[target_token_perm] + target_query_positions = token_occurrences.positions[target_token_perm] + repeated_target_edge_indices = target_edge_indices[target_group_ids] + + source_batch_node_keys = ( + target_batch_indices * num_nodes + source_indices[repeated_target_edge_indices] + ) + source_left = torch.searchsorted( + token_occurrences.sorted_batch_node_keys, + source_batch_node_keys, + right=False, + ) + source_right = torch.searchsorted( + token_occurrences.sorted_batch_node_keys, + source_batch_node_keys, + right=True, + ) + source_counts = source_right - source_left + has_source = source_counts > 0 + if not has_source.any(): + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty, empty, empty + + target_match_indices = torch.nonzero(has_source, as_tuple=True)[0] + source_counts = source_counts[has_source] + source_left = source_left[has_source] + + source_total = int(source_counts.sum().item()) + source_group_ids = torch.repeat_interleave( + torch.arange(target_match_indices.size(0), dtype=torch.long, device=device), + source_counts, + ) + source_group_starts = torch.cumsum(source_counts, dim=0) - source_counts + source_offsets = ( + torch.arange(source_total, dtype=torch.long, device=device) + - source_group_starts[source_group_ids] + ) + source_token_perm = token_occurrences.batch_node_sort_perm[ + source_left[source_group_ids] + source_offsets + ] + expanded_target_match_indices = target_match_indices[source_group_ids] + + return ( + target_batch_indices[expanded_target_match_indices], + target_query_positions[expanded_target_match_indices], + token_occurrences.positions[source_token_perm], + repeated_target_edge_indices[expanded_target_match_indices], + ) + + +def _lookup_pairwise_edge_attr_payloads( + data: HeteroData, + node_index_sequences: Tensor, + valid_mask: Tensor, + edge_attr_edge_type_to_feat_dim_map: Optional[dict[GiGLEdgeType, int]], + node_type_offsets: dict[NodeType, int], + num_nodes: int, + device: torch.device, + token_occurrences: Optional[_TokenOccurrenceIndex] = None, +) -> tuple[Optional[dict[int, Tensor]], Optional[dict[int, Tensor]]]: + """Return sparse edge-attribute attention-bias payloads by relation index.""" + if not edge_attr_edge_type_to_feat_dim_map: + return None, None + + if token_occurrences is None: + token_occurrences = _build_token_occurrence_index( + node_index_sequences=node_index_sequences, + valid_mask=valid_mask, + num_nodes=num_nodes, + device=device, + ) + + pairwise_edge_attr_indices: dict[int, Tensor] = {} + pairwise_edge_attr_values: dict[int, Tensor] = {} + if token_occurrences.node_indices.numel() == 0: + return pairwise_edge_attr_indices, pairwise_edge_attr_values + + for relation_idx, edge_type in enumerate( + sorted(edge_attr_edge_type_to_feat_dim_map.keys()) + ): + edge_attr_dim = int(edge_attr_edge_type_to_feat_dim_map[edge_type]) + if edge_attr_dim <= 0: + continue + + edge_type_key = edge_type.tuple_repr() + if edge_type_key not in data.edge_types: + continue + edge_store = data[edge_type_key] + edge_index = edge_store.edge_index + if edge_index.numel() == 0: + continue + + if not hasattr(edge_store, "edge_attr") or edge_store.edge_attr is None: + raise ValueError( + "edge_attr_attention_bias_mode='sparse_linear' requires edge_attr " + f"for edge type {edge_type}." + ) + + edge_attr = edge_store.edge_attr + if edge_attr.dim() == 1: + edge_attr = edge_attr.unsqueeze(-1) + elif edge_attr.dim() != 2: + raise ValueError( + f"edge_attr for edge type {edge_type} must be 1D or 2D, " + f"got shape {tuple(edge_attr.shape)}." + ) + if edge_attr.size(0) != edge_index.size(1): + raise ValueError( + f"edge_attr for edge type {edge_type} has {edge_attr.size(0)} rows, " + f"but edge_index has {edge_index.size(1)} edges." + ) + if edge_attr.size(1) != edge_attr_dim: + raise ValueError( + f"edge_attr for edge type {edge_type} has feature dim " + f"{edge_attr.size(1)}, expected {edge_attr_dim}." + ) + + source_indices = ( + edge_index[0].to(device=device, dtype=torch.long) + + node_type_offsets[edge_type.src_node_type] + ) + target_indices = ( + edge_index[1].to(device=device, dtype=torch.long) + + node_type_offsets[edge_type.dst_node_type] + ) + ( + batch_indices, + query_positions, + key_positions, + matched_edge_indices, + ) = _match_directed_edges_to_token_pairs( + source_indices=source_indices, + target_indices=target_indices, + token_occurrences=token_occurrences, + num_nodes=num_nodes, + device=device, + ) + if matched_edge_indices.numel() == 0: + continue + + pairwise_edge_attr_indices[relation_idx] = torch.stack( + [batch_indices, query_positions, key_positions], + dim=-1, + ) + pairwise_edge_attr_values[relation_idx] = edge_attr.to(device)[ + matched_edge_indices + ] + + return pairwise_edge_attr_indices, pairwise_edge_attr_values + + def _get_k_hop_neighbors_sparse( anchor_indices: Tensor, edge_index: Tensor, diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index d0fce10c3..42de888d3 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -274,6 +274,18 @@ def _create_user_graph_with_ppr_edges() -> HeteroData: return data +def _create_user_graph_with_edge_attrs() -> HeteroData: + data = HeteroData() + + data["user"].x = torch.tensor( + [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]] + ) + data["user", "connects", "user"].edge_index = torch.tensor([[0, 1, 1], [1, 2, 2]]) + data["user", "connects", "user"].edge_attr = torch.tensor([[0.5], [1.5], [2.5]]) + + return data + + class TestGraphTransformerEncoderPEModes(TestCase): def setUp(self) -> None: self._node_type = NodeType("user") @@ -411,6 +423,8 @@ def test_attention_bias_features_are_projected_per_head(self) -> None: ] ] ), + "pairwise_edge_attr_indices": None, + "pairwise_edge_attr_values": None, "token_input": None, }, ) @@ -421,6 +435,135 @@ def test_attention_bias_features_are_projected_per_head(self) -> None: self.assertEqual(attn_bias[0, 0, 2, 2].item(), 27.0) self.assertEqual(attn_bias[0, 1, 2, 2].item(), 38.0) + def test_edge_attr_attention_bias_zero_init_matches_baseline(self) -> None: + data = _create_user_graph_with_edge_attrs() + + base_encoder = self._create_encoder( + edge_type_to_feat_dim_map={self._edge_type: 1}, + ) + edge_attr_encoder = self._create_encoder( + edge_type_to_feat_dim_map={self._edge_type: 1}, + edge_attr_attention_bias_mode="sparse_linear", + ) + edge_attr_encoder.load_state_dict(base_encoder.state_dict(), strict=False) + + projection = edge_attr_encoder._edge_attr_attention_bias_projection_dict["0"] + assert isinstance(projection, nn.Linear) + self.assertTrue( + torch.equal(projection.weight, torch.zeros_like(projection.weight)) + ) + self.assertTrue(torch.equal(projection.bias, torch.zeros_like(projection.bias))) + + base_encoder.eval() + edge_attr_encoder.eval() + with torch.no_grad(): + base_embeddings = base_encoder( + data=data, + anchor_node_type=self._node_type, + device=self._device, + ) + edge_attr_embeddings = edge_attr_encoder( + data=data, + anchor_node_type=self._node_type, + device=self._device, + ) + + self.assertTrue( + torch.allclose(base_embeddings, edge_attr_embeddings, atol=1e-6) + ) + + def test_edge_attr_attention_bias_adds_presence_and_feature_biases(self) -> None: + encoder = self._create_encoder( + edge_type_to_feat_dim_map={self._edge_type: 2}, + edge_attr_attention_bias_mode="sparse_linear", + ) + + projection = encoder._edge_attr_attention_bias_projection_dict["0"] + assert isinstance(projection, nn.Linear) + with torch.no_grad(): + projection.weight.copy_(torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + projection.bias.copy_(torch.tensor([0.5, 1.5])) + + attn_bias = encoder._build_attention_bias( + valid_mask=torch.ones((1, 3), dtype=torch.bool), + sequences=torch.zeros((1, 3, 8), dtype=torch.float), + attention_bias_data={ + "anchor_bias": None, + "pairwise_bias": None, + "pairwise_edge_attr_indices": { + 0: torch.tensor([[0, 1, 0], [0, 1, 0], [0, 2, 1]]) + }, + "pairwise_edge_attr_values": { + 0: torch.tensor([[1.0, 1.0], [2.0, 0.0], [0.0, 1.0]]) + }, + "token_input": None, + }, + ) + + expected = torch.zeros((1, 2, 3, 3), dtype=torch.float) + expected[0, :, 1, 0] = torch.tensor([6.0, 16.0]) + expected[0, :, 2, 1] = torch.tensor([2.5, 5.5]) + self.assertTrue(torch.allclose(attn_bias, expected)) + + def test_edge_attr_attention_bias_supports_ppr_sequence_construction(self) -> None: + data = _create_user_graph_with_ppr_edges() + ppr_edge_type = EdgeType(self._node_type, Relation("ppr"), self._node_type) + + base_encoder = self._create_encoder( + edge_type_to_feat_dim_map={ppr_edge_type: 1}, + sequence_construction_method="ppr", + ) + edge_attr_encoder = self._create_encoder( + edge_type_to_feat_dim_map={ppr_edge_type: 1}, + sequence_construction_method="ppr", + edge_attr_attention_bias_mode="sparse_linear", + ) + edge_attr_encoder.load_state_dict(base_encoder.state_dict(), strict=False) + + base_encoder.eval() + edge_attr_encoder.eval() + with torch.no_grad(): + base_embeddings = base_encoder( + data=data, + anchor_node_type=self._node_type, + device=self._device, + ) + edge_attr_embeddings = edge_attr_encoder( + data=data, + anchor_node_type=self._node_type, + device=self._device, + ) + + self.assertTrue( + torch.allclose(base_embeddings, edge_attr_embeddings, atol=1e-6) + ) + + def test_edge_attr_attention_bias_skips_zero_dim_edge_types(self) -> None: + data = HeteroData() + data["user"].x = torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]) + data["user", "connects", "user"].edge_index = torch.tensor([[0], [1]]) + + encoder = self._create_encoder( + edge_type_to_feat_dim_map={self._edge_type: 0}, + edge_attr_attention_bias_mode="sparse_linear", + ) + self.assertEqual(len(encoder._edge_attr_attention_bias_projection_dict), 0) + encoder.eval() + + with torch.no_grad(): + embeddings = encoder( + data=data, + anchor_node_type=self._node_type, + device=self._device, + ) + + self.assertEqual(embeddings.shape, (2, 6)) + self.assertFalse(torch.isnan(embeddings).any()) + + def test_edge_attr_attention_bias_rejects_unknown_mode(self) -> None: + with self.assertRaisesRegex(ValueError, "edge_attr_attention_bias_mode"): + self._create_encoder(edge_attr_attention_bias_mode="dense") + def test_attention_bias_supports_anchor_relative_attrs_and_ppr_weights( self, ) -> None: @@ -447,6 +590,8 @@ def test_attention_bias_supports_anchor_relative_attrs_and_ppr_weights( [[[1.0, 0.5], [2.0, 0.25], [3.0, 0.125]]] ), "pairwise_bias": None, + "pairwise_edge_attr_indices": None, + "pairwise_edge_attr_values": None, "token_input": None, }, ) diff --git a/tests/unit/transforms/graph_transformer_test.py b/tests/unit/transforms/graph_transformer_test.py index 18551014b..d01eeb8ab 100644 --- a/tests/unit/transforms/graph_transformer_test.py +++ b/tests/unit/transforms/graph_transformer_test.py @@ -7,6 +7,7 @@ from absl.testing import absltest from torch_geometric.data import HeteroData +from gigl.src.common.types.graph_data import EdgeType, NodeType, Relation from gigl.transforms.graph_transformer import ( _get_k_hop_neighbors_sparse, heterodata_to_graph_transformer_input, @@ -252,6 +253,8 @@ def test_basic_transform(self): self.assertIsInstance(attention_bias_data, dict) self.assertIn("anchor_bias", attention_bias_data) self.assertIn("pairwise_bias", attention_bias_data) + self.assertIn("pairwise_edge_attr_indices", attention_bias_data) + self.assertIn("pairwise_edge_attr_values", attention_bias_data) def test_attention_mask_validity(self): """Test that attention mask correctly identifies valid positions.""" @@ -305,6 +308,163 @@ def test_anchor_first(self): # First position should be anchor node self.assertTrue(torch.allclose(sequences[0, 0], anchor_feature)) + def test_pairwise_edge_attr_payloads_follow_stored_direction_and_padding(self): + """Test sparse edge attrs use query=dst, key=src and exclude padding.""" + user_node_type = NodeType("user") + likes_edge_type = EdgeType( + user_node_type, + Relation("likes"), + user_node_type, + ) + follows_edge_type = EdgeType( + user_node_type, + Relation("follows"), + user_node_type, + ) + no_feat_edge_type = EdgeType( + user_node_type, + Relation("no_feat"), + user_node_type, + ) + missing_edge_type = EdgeType( + user_node_type, + Relation("missing"), + user_node_type, + ) + + data = HeteroData() + data["user"].x = torch.tensor( + [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], + ) + data[likes_edge_type.tuple_repr()].edge_index = torch.tensor([[0], [1]]) + data[likes_edge_type.tuple_repr()].edge_attr = torch.tensor([[0.25, 1.25]]) + data[follows_edge_type.tuple_repr()].edge_index = torch.tensor( + [[0, 1, 1], [1, 2, 2]], + ) + data[follows_edge_type.tuple_repr()].edge_attr = torch.tensor( + [[2.0], [3.0], [4.0]], + ) + data[no_feat_edge_type.tuple_repr()].edge_index = torch.tensor([[0], [2]]) + + edge_attr_dim_map = { + likes_edge_type: 2, + follows_edge_type: 1, + no_feat_edge_type: 0, + missing_edge_type: 1, + } + _, valid_mask, sequence_auxiliary_data = heterodata_to_graph_transformer_input( + data=data, + batch_size=1, + max_seq_len=4, + anchor_node_type="user", + hop_distance=2, + edge_attr_edge_type_to_feat_dim_map=edge_attr_dim_map, + ) + + self.assertTrue( + torch.equal(valid_mask[0], torch.tensor([True, True, True, False])) + ) + pairwise_edge_attr_indices = sequence_auxiliary_data[ + "pairwise_edge_attr_indices" + ] + pairwise_edge_attr_values = sequence_auxiliary_data["pairwise_edge_attr_values"] + assert pairwise_edge_attr_indices is not None + assert pairwise_edge_attr_values is not None + + sorted_edge_types = sorted(edge_attr_dim_map.keys()) + likes_relation_idx = sorted_edge_types.index(likes_edge_type) + follows_relation_idx = sorted_edge_types.index(follows_edge_type) + no_feat_relation_idx = sorted_edge_types.index(no_feat_edge_type) + missing_relation_idx = sorted_edge_types.index(missing_edge_type) + + self.assertTrue( + torch.equal( + pairwise_edge_attr_indices[likes_relation_idx], + torch.tensor([[0, 1, 0]]), + ) + ) + self.assertTrue( + torch.allclose( + pairwise_edge_attr_values[likes_relation_idx], + torch.tensor([[0.25, 1.25]]), + ) + ) + self.assertTrue( + torch.equal( + pairwise_edge_attr_indices[follows_relation_idx], + torch.tensor([[0, 1, 0], [0, 2, 1], [0, 2, 1]]), + ) + ) + self.assertTrue( + torch.allclose( + pairwise_edge_attr_values[follows_relation_idx], + torch.tensor([[2.0], [3.0], [4.0]]), + ) + ) + self.assertNotIn(no_feat_relation_idx, pairwise_edge_attr_indices) + self.assertNotIn(missing_relation_idx, pairwise_edge_attr_indices) + all_indices = torch.cat(list(pairwise_edge_attr_indices.values()), dim=0) + self.assertFalse((all_indices[:, 1:] == 3).any()) + + def test_pairwise_edge_attr_payloads_missing_edge_attr_raises(self): + user_node_type = NodeType("user") + edge_type = EdgeType(user_node_type, Relation("connects"), user_node_type) + + data = HeteroData() + data["user"].x = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + data[edge_type.tuple_repr()].edge_index = torch.tensor([[0], [1]]) + + with self.assertRaisesRegex(ValueError, "requires edge_attr"): + heterodata_to_graph_transformer_input( + data=data, + batch_size=1, + max_seq_len=2, + anchor_node_type="user", + hop_distance=1, + edge_attr_edge_type_to_feat_dim_map={edge_type: 1}, + ) + + def test_pairwise_edge_attr_payloads_reject_malformed_edge_attrs(self): + user_node_type = NodeType("user") + edge_type = EdgeType(user_node_type, Relation("connects"), user_node_type) + + data = HeteroData() + data["user"].x = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + data[edge_type.tuple_repr()].edge_index = torch.tensor([[0], [1]]) + data[edge_type.tuple_repr()].edge_attr = torch.tensor([[1.0], [2.0]]) + + with self.assertRaisesRegex(ValueError, "has 2 rows"): + heterodata_to_graph_transformer_input( + data=data, + batch_size=1, + max_seq_len=2, + anchor_node_type="user", + hop_distance=1, + edge_attr_edge_type_to_feat_dim_map={edge_type: 1}, + ) + + data[edge_type.tuple_repr()].edge_attr = torch.tensor([[1.0, 2.0]]) + with self.assertRaisesRegex(ValueError, "feature dim 2"): + heterodata_to_graph_transformer_input( + data=data, + batch_size=1, + max_seq_len=2, + anchor_node_type="user", + hop_distance=1, + edge_attr_edge_type_to_feat_dim_map={edge_type: 1}, + ) + + data[edge_type.tuple_repr()].edge_attr = torch.tensor([[[1.0]]]) + with self.assertRaisesRegex(ValueError, "1D or 2D"): + heterodata_to_graph_transformer_input( + data=data, + batch_size=1, + max_seq_len=2, + anchor_node_type="user", + hop_distance=1, + edge_attr_edge_type_to_feat_dim_map={edge_type: 1}, + ) + def test_different_anchor_types(self): """Test with different anchor node types.""" data = create_simple_hetero_data()