Fix FLAVA bugs #530 and #533 (attention masks and CLS tokens)

tests/models/flava/test_flava.py

@@ -192,10 +192,14 @@ def flava(
         image_encoder,
         text_encoder,
     ):
+        class DummyMMEncoder(nn.Module):
+            def forward(self, hidden_states, attention_mask=None):
+                return hidden_states
+
         flava_model = FLAVAModel(
             image_encoder=image_encoder,
             text_encoder=text_encoder,
-            mm_encoder=nn.Identity(),
+            mm_encoder=DummyMMEncoder(),
             image_to_mm_projection=nn.Identity(),
             text_to_mm_projection=nn.Identity(),
             text_projection=nn.Identity(),
@@ -343,3 +347,49 @@ def test_forward_image(self, image_encoder, flava, inputs):
         assert_expected(
             actual.projected_image_embeddings, expected_image.last_hidden_state[:, 0, :]
         )
+
+    def test_attention_mask_affects_output(self):
+        from torchmultimodal.models.flava.model import flava_multimodal_encoder
+        from torchmultimodal.utils.attention import get_extended_attention_mask
+
+        model = flava_multimodal_encoder(
+            hidden_size=4, num_attention_heads=1, num_hidden_layers=1, intermediate_size=4
+        )
+        model.eval()
+
+        torch.manual_seed(42)
+        B, seq_len = 2, 8  # simulates fused image+text sequence
+        hidden_states = torch.randn(B, seq_len, 4)  # non-trivial random inputs
+
+        # Mask that blocks last 2 positions (simulating PAD tokens)
+        mask_with_pad = torch.ones(B, seq_len, dtype=torch.long)
+        mask_with_pad[:, -2:] = 0
+        mask_no_pad = torch.ones(B, seq_len, dtype=torch.long)
+
+        with torch.no_grad():
+            out_masked   = model(hidden_states, 
+                                 attention_mask=get_extended_attention_mask(mask_with_pad))
+            out_unmasked = model(hidden_states, 
+                                 attention_mask=get_extended_attention_mask(mask_no_pad))
+
+        assert not torch.allclose(
+            out_masked.last_hidden_state,
+            out_unmasked.last_hidden_state,
+        ), "Masking pad positions should change encoder output"
+
+    def test_flava_transformer_without_embeddings_attention_mask(self):
+        from torchmultimodal.models.flava.model import flava_multimodal_encoder
+        model = flava_multimodal_encoder(
+            hidden_size=2,
+            num_attention_heads=1,
+            num_hidden_layers=1,
+            intermediate_size=2,
+        )
+        batch_size = 2
+        seq_len = 5
+        hidden_states = torch.randn(batch_size, seq_len, 2)
+        attention_mask = torch.ones(batch_size, seq_len)
+        from torchmultimodal.utils.attention import get_extended_attention_mask
+        attention_mask = get_extended_attention_mask(attention_mask)
+        out = model(hidden_states, attention_mask=attention_mask)
+        assert out.last_hidden_state is not None

torchmultimodal/models/flava/model.py

@@ -29,6 +29,7 @@
     FLAVAPretrainingLossOutput,
     Pooler,
 )
+from torchmultimodal.utils.attention import get_extended_attention_mask
 from torchmultimodal.utils.common import load_module_from_url, ModelOutput
 from typing_extensions import Literal
 
@@ -132,7 +133,16 @@ def forward(
         text_masked: Optional[Tensor] = None,
         required_embedding: Optional[EMBEDDING_OPTIONS] = None,
         skip_unmasked_mm_encoder: bool = True,
+        text_pad_mask: Optional[Tensor] = None,
+        text_masked_pad_mask: Optional[Tensor] = None,
     ) -> FLAVAOutput:
+        if text_pad_mask is None and text is not None:
+            pad_token_id = getattr(self.text_encoder.embeddings, "pad_token_id", 0) if hasattr(self.text_encoder, "embeddings") else 0
+            text_pad_mask = (text != pad_token_id).long()
+        if text_masked_pad_mask is None and text_masked is not None:
+            pad_token_id = getattr(self.text_encoder.embeddings, "pad_token_id", 0) if hasattr(self.text_encoder, "embeddings") else 0
+            text_masked_pad_mask = (text_masked != pad_token_id).long()
+
         if required_embedding is None:
             if image is not None and text is not None:
                 required_embedding = "mm"
@@ -160,7 +170,7 @@ def forward(
             text,
             required_embedding,
             ["text", "mm"],
-            partial(self.encode_text, projection=True),
+            partial(self.encode_text, text_mask=text_pad_mask, projection=True),
         )
         if len(text_encoding_out) == 2:
             text_outputs, projected_text_embeddings = (
@@ -182,7 +192,7 @@ def forward(
             text_masked,
             required_embedding,
             ["text", "mm"],
-            self.encode_text,
+            partial(self.encode_text, text_mask=text_masked_pad_mask),
         )
         assert type(text_masked_outputs) == TransformerOutput
 
@@ -193,6 +203,16 @@ def forward(
             # Take last hidden state and not the last_hidden_state because
             # for flava we want the hidden state without final layernorm.
             if not skip_unmasked_mm_encoder:
+                mm_attention_mask = None
+                if text_pad_mask is not None and image_outputs.hidden_states:
+                    image_seq_len = image_outputs.hidden_states[-1].shape[1]
+                    image_mask = torch.ones(
+                        (text_pad_mask.shape[0], image_seq_len),
+                        device=text_pad_mask.device,
+                        dtype=text_pad_mask.dtype
+                    )
+                    mm_attention_mask = torch.cat([image_mask, text_pad_mask], dim=1)
+
                 # Unmasked multimodal embedding is not currently used by any of the FLAVA losses.
                 multimodal_outputs = self.encode_mm(
                     (
@@ -205,7 +225,19 @@ def forward(
                         if text_outputs.hidden_states  # type: ignore
                         else None
                     ),
+                    attention_mask=mm_attention_mask,
                 )
+
+            mm_masked_attention_mask = None
+            if text_masked_pad_mask is not None and image_masked_outputs.hidden_states:
+                image_seq_len = image_masked_outputs.hidden_states[-1].shape[1]
+                image_mask = torch.ones(
+                    (text_masked_pad_mask.shape[0], image_seq_len),
+                    device=text_masked_pad_mask.device,
+                    dtype=text_masked_pad_mask.dtype
+                )
+                mm_masked_attention_mask = torch.cat([image_mask, text_masked_pad_mask], dim=1)
+
             multimodal_masked_outputs = self.encode_mm(
                 (
                     image_masked_outputs.hidden_states[-1]
@@ -217,6 +249,7 @@ def forward(
                     if text_masked_outputs.hidden_states
                     else None
                 ),
+                attention_mask=mm_masked_attention_mask,
             )
 
         return FLAVAOutput(
@@ -286,6 +319,7 @@ def encode_mm(
         self,
         image_embedding: Tensor,
         text_embedding: Tensor,
+        attention_mask: Optional[Tensor] = None,
     ) -> TransformerOutput:
         if image_embedding is None or text_embedding is None:
             # Since nothing is passed, it might be case without
@@ -295,7 +329,9 @@ def encode_mm(
         image_embedding = self.image_to_mm_projection(image_embedding)
         text_embedding = self.text_to_mm_projection(text_embedding)
         fused_state = torch.cat([image_embedding, text_embedding], dim=1)
-        return self.mm_encoder(fused_state)
+        if attention_mask is not None:
+            attention_mask = get_extended_attention_mask(attention_mask)
+        return self.mm_encoder(fused_state, attention_mask=attention_mask)
 
 
 class FLAVAForPreTraining(nn.Module):
@@ -342,6 +378,8 @@ def forward(
         skip_unmasked_mm_encoder: bool = True,
         itm_labels: Optional[Tensor] = None,
         mlm_labels: Optional[Tensor] = None,
+        text_pad_mask: Optional[Tensor] = None,
+        text_masked_pad_mask: Optional[Tensor] = None,
     ) -> FLAVAPretrainingLossOutput:
         image_labels = None
         if image_for_codebook is not None:
@@ -356,6 +394,8 @@ def forward(
             text_masked=text_masked,
             required_embedding=required_embedding,
             skip_unmasked_mm_encoder=skip_unmasked_mm_encoder,
+            text_pad_mask=text_pad_mask,
+            text_masked_pad_mask=text_masked_pad_mask,
         )
 
         return self.loss(
@@ -397,13 +437,15 @@ def forward(
         required_embedding: Optional[EMBEDDING_OPTIONS] = None,
         labels: Optional[Tensor] = None,
         cls_index: int = 0,
+        text_pad_mask: Optional[Tensor] = None,
     ) -> FLAVAForClassificationOutput:
         flava_output: FLAVAOutput = self.model(
             image=image,
             text=text,
             required_embedding=required_embedding,
             # Don't skip the encoder for classification
             skip_unmasked_mm_encoder=False,
+            text_pad_mask=text_pad_mask,
         )
 
         hidden_state: Optional[Tensor] = None
@@ -741,4 +783,4 @@ def get_codebook_probs(self, images: Tensor) -> Tensor:
         return nn.Softmax(dim=1)(z_logits)
 
     def forward(self, img_seq_prob: Tensor) -> Tensor:
-        return self.get_codebook_indices(img_seq_prob)
+        return self.get_codebook_indices(img_seq_prob)

torchmultimodal/models/flava/transformer.py

@@ -49,13 +49,29 @@ def forward(
         hidden_states: Optional[Tensor] = None,
         attention_mask: Optional[Tensor] = None,
     ) -> TransformerOutput:
+        """
+        Args:
+            hidden_states (Tensor, optional): Input tensor of shape [batch, seq_len, hidden_size]
+            attention_mask (Tensor, optional): Mask to be applied to self-attention inputs. 
+                If provided, it should be pre-expanded via `get_extended_attention_mask` to 
+                shape (batch_size, 1, 1, seq_len). The internal cls_token expansion is handled automatically.
+        """
         if hidden_states is None:
             raise ValueError("You have to specify hidden_states")
 
         if self.cls_token is not None:
             batch_size = hidden_states.shape[0]
             cls_tokens = self.cls_token.expand(batch_size, -1, -1)
             hidden_states = torch.cat((cls_tokens, hidden_states), dim=1)
+            if attention_mask is not None:
+                cls_mask_shape = list(attention_mask.shape)
+                cls_mask_shape[-1] = 1
+                cls_mask = torch.ones(
+                    cls_mask_shape,
+                    device=attention_mask.device,
+                    dtype=attention_mask.dtype,
+                )
+                attention_mask = torch.cat((cls_mask, attention_mask), dim=-1)
 
         encoder_output = self.encoder(
             hidden_states,
@@ -307,4 +323,4 @@ def init_transformer_weights(module: nn.Module, initializer_range: float) -> Non
             module.weight.data[module.padding_idx].zero_()
     elif isinstance(module, nn.LayerNorm):
         module.bias.data.zero_()
-        module.weight.data.fill_(1.0)
+        module.weight.data.fill_(1.0)