Mask padding tokens during text pooling

src/colipri/model/text.py

@@ -36,7 +36,10 @@ class TextEncoder(nn.Module, ModelMixin):
         output = self.backbone(token_ids, attention_mask=attention_mask)
         embeddings = output["last_hidden_state"]
         if pool:
-            embeddings = self.pooler(embeddings)
+            embeddings = self.pooler(
+                embeddings,
+                attention_mask=attention_mask,
+            )
         if normalize:
             embeddings = F.normalize(embeddings, dim=1)
         return embeddings

src/colipri/pooling.py

@@ -4,6 +4,57 @@ from torch import nn
 
 from .types import TypePooledEmbeddings
 from .types import TypeSequenceEmbeddings
+from .types import TypeTextAttentionMask
+from .types import TypeValidTokenMask
+
+
+def _get_valid_token_mask(
+    sequence: TypeSequenceEmbeddings,
+    attention_mask: TypeTextAttentionMask | None,
+) -> TypeValidTokenMask | None:
+    """Get a validated Boolean mask for valid sequence positions.
+
+    Args:
+        sequence: Sequence embeddings.
+        attention_mask: Mask where nonzero entries identify valid positions.
+
+    Returns:
+        A Boolean mask on the sequence device, or `None` when no mask is provided.
+
+    Raises:
+        ValueError: If the mask shape does not match the sequence or any sequence
+            contains no valid positions.
+    """
+    if attention_mask is None:
+        return None
+
+    expected_shape = sequence.shape[:2]
+    if attention_mask.shape != expected_shape:
+        msg = (
+            f"Expected attention_mask shape {expected_shape},"
+            f" but got {attention_mask.shape}"
+        )
+        raise ValueError(msg)
+
+    valid_mask = attention_mask.to(device=sequence.device, dtype=torch.bool)
+    if not valid_mask.any(dim=1).all():
+        msg = "Each sequence must contain at least one valid token"
+        raise ValueError(msg)
+    return valid_mask
+
+
+def _compute_masked_mean(
+    sequence: TypeSequenceEmbeddings,
+    valid_mask: TypeValidTokenMask | None,
+) -> TypePooledEmbeddings:
+    """Compute the sequence mean while excluding masked positions."""
+    if valid_mask is None:
+        return sequence.mean(dim=1)
+    weights = rearrange(
+        valid_mask,
+        "batch length -> batch length 1",
+    ).to(sequence.dtype)
+    return (sequence * weights).sum(dim=1) / weights.sum(dim=1)
 
 
 class AttentionPool1D(nn.Module):
@@ -11,10 +62,24 @@ class AttentionPool1D(nn.Module):
         super().__init__()
         self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
 
-    def forward(self, x: TypeSequenceEmbeddings) -> TypePooledEmbeddings:
-        query = x.mean(dim=1, keepdim=True)
+    def forward(
+        self,
+        x: TypeSequenceEmbeddings,
+        attention_mask: TypeTextAttentionMask | None = None,
+    ) -> TypePooledEmbeddings:
+        valid_mask = _get_valid_token_mask(x, attention_mask)
+        query = rearrange(
+            _compute_masked_mean(x, valid_mask),
+            "batch embed_dim -> batch 1 embed_dim",
+        )
         key = value = x
-        pooled, _ = self.attn(query, key, value)
+        key_padding_mask = None if valid_mask is None else ~valid_mask
+        pooled, _ = self.attn(
+            query,
+            key,
+            value,
+            key_padding_mask=key_padding_mask,
+        )
         return rearrange(pooled, "batch 1 embed_dim -> batch embed_dim")
 
     def to_dense(self):
@@ -62,13 +127,24 @@ class MultiLearnedQueryAttentionPool1D(AttentionPool1D):
         # 4 Queries instead of 1
         self.query = nn.Parameter(torch.randn(1, 4, embed_dim) / embed_dim**0.5)
 
-    def forward(self, x: TypeSequenceEmbeddings) -> TypePooledEmbeddings:
+    def forward(
+        self,
+        x: TypeSequenceEmbeddings,
+        attention_mask: TypeTextAttentionMask | None = None,
+    ) -> TypePooledEmbeddings:
         """
         x: [B, T, D] — sequence of token embeddings
         returns: [B, D] — pooled representation
         """
-        B, T, D = x.shape
-        query = self.query.expand(B, -1, -1)  # [B, 4, D]
-        pooled, _ = self.attn(query, x, x)  # [B, 4, D]
+        valid_mask = _get_valid_token_mask(x, attention_mask)
+        batch_size = x.shape[0]
+        query = self.query.expand(batch_size, -1, -1)  # [B, 4, D]
+        key_padding_mask = None if valid_mask is None else ~valid_mask
+        pooled, _ = self.attn(
+            query,
+            x,
+            x,
+            key_padding_mask=key_padding_mask,
+        )
         # pooled: [4, B, D_out], want [B, D_out] by pooling over queries (mean)
         return pooled.mean(dim=1)  # [B, D]

src/colipri/types.py

@@ -2,6 +2,7 @@ import os
 from typing import TypeAlias
 
 import torchio as tio
+from jaxtyping import Bool
 from jaxtyping import Float
 from jaxtyping import Int64
 from torch import Tensor
@@ -25,6 +26,7 @@ TypeImagesTensor: TypeAlias = Float[Tensor, "batch 1 x_in y_in z_in"]
 
 TypeTextTokenIds: TypeAlias = Int64[Tensor, "batch text_length"]
 TypeTextAttentionMask: TypeAlias = Int64[Tensor, "batch text_length"]
+TypeValidTokenMask: TypeAlias = Bool[Tensor, "batch sequence_length"]
 
 # Outputs
 TypeSequenceEmbeddings: TypeAlias = Float[Tensor, "batch sequence_length embed_dim"]

tests/test_pooling.py

@@ -0,0 +1,47 @@
+import pytest
+import torch
+
+from colipri.pooling import AttentionPool1D
+from colipri.pooling import MultiLearnedQueryAttentionPool1D
+
+
+@pytest.mark.parametrize(
+    "pool_type",
+    [AttentionPool1D, MultiLearnedQueryAttentionPool1D],
+)
+def test_ignores_masked_tokens(pool_type):
+    """Ignore padded positions when pooling language embeddings."""
+    torch.manual_seed(0)
+    pool = pool_type(embed_dim=8, num_heads=2)
+    valid = torch.randn(2, 3, 8)
+    padded = torch.cat([valid, torch.full((2, 5, 8), 1e6)], dim=1)
+    valid_mask = torch.ones(2, 3, dtype=torch.long)
+    padded_mask = torch.cat(
+        [valid_mask, torch.zeros(2, 5, dtype=torch.long)],
+        dim=1,
+    )
+
+    expected = pool(valid)
+    result = pool(padded, attention_mask=padded_mask)
+
+    torch.testing.assert_close(result, expected)
+
+
+def test_rejects_incorrect_attention_mask_shape():
+    """Reject an attention mask that does not match the input sequence."""
+    pool = AttentionPool1D(embed_dim=8, num_heads=2)
+    sequence = torch.randn(2, 3, 8)
+    attention_mask = torch.ones(2, 4, dtype=torch.long)
+
+    with pytest.raises(ValueError, match="attention_mask shape"):
+        pool(sequence, attention_mask=attention_mask)
+
+
+def test_rejects_fully_masked_sequence():
+    """Reject a sequence without any valid tokens."""
+    pool = AttentionPool1D(embed_dim=8, num_heads=2)
+    sequence = torch.randn(2, 3, 8)
+    attention_mask = torch.tensor([[1, 1, 1], [0, 0, 0]])
+
+    with pytest.raises(ValueError, match="at least one valid token"):
+        pool(sequence, attention_mask=attention_mask)

tests/test_text.py

@@ -0,0 +1,52 @@
+import torch
+from torch import nn
+
+from colipri.model.text import TextEncoder
+from colipri.pooling import AttentionPool1D
+
+
+class _StubBackbone(nn.Module):
+    """Encode token IDs as deterministic embeddings."""
+
+    def __init__(self):
+        """Initialize embeddings with a distinctive padding-token vector."""
+        super().__init__()
+        weights = torch.tensor(
+            [
+                [100.0, -100.0, 50.0, -50.0],
+                [1.0, 0.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0, 0.0],
+            ]
+        )
+        self.embedding = nn.Embedding.from_pretrained(weights, freeze=False)
+
+    def forward(self, token_ids, attention_mask=None):  # noqa: ARG002
+        """Encode token IDs without altering padded positions."""
+        return {"last_hidden_state": self.embedding(token_ids)}
+
+
+def test_text_encoder_ignores_masked_padding():
+    """Preserve the pooled embedding when masked padding is appended."""
+    torch.manual_seed(0)
+    encoder = TextEncoder(
+        backbone=_StubBackbone(),
+        pooler=AttentionPool1D(embed_dim=4, num_heads=2),
+    )
+    valid_ids = torch.tensor([[1, 2, 3]])
+    padded_ids = torch.tensor([[1, 2, 3, 0, 0]])
+    valid_mask = torch.ones_like(valid_ids)
+    padded_mask = torch.tensor([[1, 1, 1, 0, 0]])
+
+    expected = encoder(
+        valid_ids,
+        valid_mask,
+        normalize=False,
+    )
+    result = encoder(
+        padded_ids,
+        padded_mask,
+        normalize=False,
+    )
+
+    torch.testing.assert_close(result, expected)