diff --git a/gigl/nn/graph_transformer.py b/gigl/nn/graph_transformer.py index f9d8b345a..470b8a17b 100644 --- a/gigl/nn/graph_transformer.py +++ b/gigl/nn/graph_transformer.py @@ -1,6 +1,6 @@ """Graph Transformer encoder for heterogeneous graphs. -Adapted from RelGT's LocalModule (https://github.com/snap-stanford/relgt). +Adapted from RelGT's LocalModule. Converts heterogeneous graph data into fixed-length sequences via ``heterodata_to_graph_transformer_input``, processes through a stack of pre-norm transformer encoder layers, then produces per-node embeddings via @@ -28,6 +28,8 @@ heterodata_to_graph_transformer_input, ) +_GRAPH_TRANSFORMER_READOUT_MODES = {"anchor_attention", "cls"} + def _get_node_type_positional_encodings( data: torch_geometric.data.hetero_data.HeteroData, @@ -402,6 +404,10 @@ class GraphTransformerEncoder(nn.Module): attention_dropout_rate: Dropout probability for attention weights. should_l2_normalize_embedding_layer_output: Whether to L2 normalize output embeddings. + readout_mode: Sequence readout strategy. ``"anchor_attention"`` keeps + the historical anchor plus learned neighbor aggregation behavior, + ``"cls"`` prepends a learned sequence token and returns it after the + transformer stack. pe_attr_names: List of node-level positional encoding attribute names. In ``"concat"`` mode these are concatenated to sequence features. In ``"add"`` mode they are projected to ``hid_dim`` and added to @@ -490,6 +496,7 @@ def __init__( dropout_rate: float = 0.1, attention_dropout_rate: float = 0.0, should_l2_normalize_embedding_layer_output: bool = False, + readout_mode: Literal["anchor_attention", "cls"] = "anchor_attention", pe_attr_names: Optional[list[str]] = None, anchor_based_attention_bias_attr_names: Optional[list[str]] = None, anchor_based_input_attr_names: Optional[list[str]] = None, @@ -531,6 +538,11 @@ def __init__( "{None, 'sinusoidal'}, " f"got '{sequence_positional_encoding_type}'" ) + if readout_mode not in _GRAPH_TRANSFORMER_READOUT_MODES: + raise ValueError( + f"readout_mode must be one of {sorted(_GRAPH_TRANSFORMER_READOUT_MODES)}, " + f"got '{readout_mode}'" + ) if ( sequence_construction_method == "khop" and sequence_positional_encoding_type is not None @@ -558,6 +570,7 @@ def __init__( ) self._sequence_construction_method = sequence_construction_method self._sequence_positional_encoding_type = sequence_positional_encoding_type + self._readout_mode = readout_mode self._should_l2_normalize_embedding_layer_output = ( should_l2_normalize_embedding_layer_output ) @@ -605,6 +618,11 @@ def __init__( None, persistent=False, ) + if self._readout_mode == "cls": + self._cls_token = nn.Parameter(torch.zeros(1, 1, hid_dim)) + nn.init.normal_(self._cls_token, mean=0.0, std=0.02) + else: + self.register_parameter("_cls_token", None) # Per-node-type input projection to hid_dim (like HGT's lin_dict) self._node_projection_dict = nn.ModuleDict( @@ -793,7 +811,11 @@ def forward( ) = heterodata_to_graph_transformer_input( data=projected_data, batch_size=num_anchor_nodes, - max_seq_len=self._max_seq_len, + max_seq_len=( + self._max_seq_len - 1 + if self._readout_mode == "cls" + else self._max_seq_len + ), anchor_node_type=anchor_node_type, anchor_node_ids=anchor_node_ids, hop_distance=self._hop_distance, @@ -820,6 +842,15 @@ def forward( valid_mask=valid_mask, ) + if self._readout_mode == "cls": + sequences, valid_mask, sequence_auxiliary_data = ( + self._prepend_cls_token_to_sequence( + sequences=sequences, + valid_mask=valid_mask, + sequence_auxiliary_data=sequence_auxiliary_data, + ) + ) + sequence_positional_encoding = self._get_sequence_positional_encoding( valid_mask=valid_mask, sequences=sequences, @@ -845,6 +876,72 @@ def forward( return embeddings + def _prepend_cls_token_to_sequence( + self, + sequences: Tensor, + valid_mask: Tensor, + sequence_auxiliary_data: SequenceAuxiliaryData, + ) -> tuple[Tensor, Tensor, SequenceAuxiliaryData]: + """Prepend a learned CLS token and zero-valued auxiliary features.""" + if self._cls_token is None: + raise ValueError("CLS token is not initialized.") + + batch_size = sequences.size(0) + cls_token = self._cls_token.to( + device=sequences.device, + dtype=sequences.dtype, + ).expand(batch_size, -1, -1) + sequences = torch.cat([cls_token, sequences], dim=1) + + cls_valid_mask = torch.ones( + (batch_size, 1), + dtype=valid_mask.dtype, + device=valid_mask.device, + ) + valid_mask = torch.cat([cls_valid_mask, valid_mask], dim=1) + + anchor_bias = sequence_auxiliary_data["anchor_bias"] + if anchor_bias is not None: + cls_anchor_bias = torch.zeros( + batch_size, + 1, + anchor_bias.size(-1), + dtype=anchor_bias.dtype, + device=anchor_bias.device, + ) + anchor_bias = torch.cat([cls_anchor_bias, anchor_bias], dim=1) + + pairwise_bias = sequence_auxiliary_data["pairwise_bias"] + if pairwise_bias is not None: + cls_key_bias = torch.zeros( + batch_size, + pairwise_bias.size(1), + 1, + pairwise_bias.size(-1), + dtype=pairwise_bias.dtype, + device=pairwise_bias.device, + ) + pairwise_bias = torch.cat([cls_key_bias, pairwise_bias], dim=2) + cls_query_bias = torch.zeros( + batch_size, + 1, + pairwise_bias.size(2), + pairwise_bias.size(-1), + dtype=pairwise_bias.dtype, + device=pairwise_bias.device, + ) + pairwise_bias = torch.cat([cls_query_bias, pairwise_bias], dim=1) + + return ( + sequences, + valid_mask, + { + "anchor_bias": anchor_bias, + "pairwise_bias": pairwise_bias, + "token_input": sequence_auxiliary_data["token_input"], + }, + ) + def _get_sequence_positional_encoding( self, valid_mask: Tensor, @@ -1031,61 +1128,63 @@ def _build_attention_bias( return attn_bias - def _encode_and_readout( + def _readout_from_encoded_sequences( self, - sequences: Tensor, + x: Tensor, valid_mask: Tensor, - attn_bias: Optional[Tensor] = None, ) -> Tensor: - """Process sequences through transformer layers and attention readout. + if self._readout_mode == "cls": + return x[:, 0, :] - Args: - sequences: Input tensor of shape ``(batch_size, max_seq_len, hid_dim)``. - valid_mask: Boolean mask of shape ``(batch_size, max_seq_len)``. - attn_bias: Optional additive attention bias broadcastable to - ``(batch_size, num_heads, seq, seq)``. + anchor = x[:, 0, :].unsqueeze(1) - Returns: - Output embeddings of shape ``(batch_size, hid_dim)``. - """ - x = sequences * valid_mask.unsqueeze(-1).to(sequences.dtype) - - for encoder_layer in self._encoder_layers: - x = encoder_layer(x, attn_bias=attn_bias, valid_mask=valid_mask) - - x = self._final_norm(x) - x = x * valid_mask.unsqueeze(-1).to(x.dtype) - - # Readout: anchor (position 0) + attention-weighted neighbor aggregation - anchor = x[:, 0, :].unsqueeze(1) # (batch, 1, hid_dim) - neighbors = x[:, 1:, :] # (batch, seq-1, hid_dim) + # Historical readout: anchor token plus attention-weighted neighbors. + neighbors = x[:, 1:, :] neighbor_valid_mask = valid_mask[:, 1:] seq_minus_one = neighbors.size(1) if seq_minus_one == 0: return anchor.squeeze(1) - # Expand anchor to match neighbor dimension for concatenation anchor_expanded = anchor.expand(-1, seq_minus_one, -1) - - # Compute attention scores over neighbors readout_scores = self._readout_attention( torch.cat([anchor_expanded, neighbors], dim=-1) - ) # (batch, seq-1, 1) + ) readout_scores = readout_scores.masked_fill( ~neighbor_valid_mask.unsqueeze(-1), torch.finfo(readout_scores.dtype).min, ) - readout_weights = F.softmax(readout_scores, dim=1) # (batch, seq-1, 1) + readout_weights = F.softmax(readout_scores, dim=1) readout_weights = torch.nan_to_num(readout_weights, nan=0.0) readout_weights = readout_weights * neighbor_valid_mask.unsqueeze(-1).to( readout_weights.dtype ) + neighbor_aggregation = (neighbors * readout_weights).sum(dim=1, keepdim=True) + return (anchor + neighbor_aggregation).squeeze(1) + + def _encode_and_readout( + self, + sequences: Tensor, + valid_mask: Tensor, + attn_bias: Optional[Tensor] = None, + ) -> Tensor: + """Process sequences through transformer layers and attention readout. - neighbor_aggregation = (neighbors * readout_weights).sum( - dim=1, keepdim=True - ) # (batch, 1, hid_dim) + Args: + sequences: Input tensor of shape ``(batch_size, max_seq_len, hid_dim)``. + valid_mask: Boolean mask of shape ``(batch_size, max_seq_len)``. + attn_bias: Optional additive attention bias broadcastable to + ``(batch_size, num_heads, seq, seq)``. - output = (anchor + neighbor_aggregation).squeeze(1) # (batch, hid_dim) + Returns: + Output embeddings of shape ``(batch_size, hid_dim)``. + """ + x = sequences * valid_mask.unsqueeze(-1).to(sequences.dtype) + + for encoder_layer in self._encoder_layers: + x = encoder_layer(x, attn_bias=attn_bias, valid_mask=valid_mask) + + x = self._final_norm(x) + x = x * valid_mask.unsqueeze(-1).to(x.dtype) - return output + return self._readout_from_encoded_sequences(x=x, valid_mask=valid_mask) diff --git a/tests/unit/nn/graph_transformer_test.py b/tests/unit/nn/graph_transformer_test.py index d0fce10c3..ded25a3ee 100644 --- a/tests/unit/nn/graph_transformer_test.py +++ b/tests/unit/nn/graph_transformer_test.py @@ -571,6 +571,30 @@ def test_forward_supports_sinusoidal_sequence_position_encoding_in_ppr_mode( ) ) + def test_forward_supports_cls_readout_in_ppr_mode(self) -> None: + data = _create_user_graph_with_ppr_edges() + ppr_edge_type = EdgeType(self._node_type, Relation("ppr"), self._node_type) + + encoder = self._create_encoder( + edge_type_to_feat_dim_map={ppr_edge_type: 0}, + sequence_construction_method="ppr", + sequence_positional_encoding_type="sinusoidal", + readout_mode="cls", + anchor_based_attention_bias_attr_names=["ppr_weight"], + ) + encoder.eval() + + with torch.no_grad(): + embeddings = encoder( + data=data, + anchor_node_type=self._node_type, + device=self._device, + ) + + self.assertIsNotNone(encoder._cls_token) + self.assertEqual(embeddings.shape, (3, 6)) + self.assertFalse(torch.isnan(embeddings).any()) + def test_forward_supports_anchor_relative_and_ppr_token_input_features( self, ) -> None: