Add CLS readout mode to graph transformer

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)

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: