From e7b03b8d7800e0a9e60756b0044afa1a6fcac616 Mon Sep 17 00:00:00 2001 From: "Vishal S. Pandey" Date: Wed, 8 Apr 2026 19:50:26 +0530 Subject: [PATCH] Fix FLAVA bugs #530 and #533 (attention masks and CLS tokens) This commit fixes two open bugs in the FLAVA model implementation: - Issue #533: Fixed a bug in `FLAVATransformerWithoutEmbeddings` where prepending the CLS token to `hidden_states` caused sequence length mismatches because `attention_mask` was not correspondingly padded. It now dynamically concatenates an active mask token to `attention_mask` matching the CLS token in a dimension-agnostic way (`dim=-1`). Also updated the docstring to clarify that `attention_mask` needs to be pre-expanded to `(B, 1, 1, seq_len)`. - Issue #530: The multimodal encoder `FLAVAModel.encode_mm` was previously ignoring padded text tokens. The text padding mask is now dynamically generated from input text (or can be passed directly), combined with an all-ones image mask, and passed to the multimodal cross-attention encoder. This parameter is properly threaded through both `FLAVAForPreTraining` and `FLAVAForClassification`. - Added a robust regression test (`test_attention_mask_affects_output`) directly targeting `FLAVATransformerWithoutEmbeddings` to verify that padded tokens properly alter the final encoder output. --- tests/models/flava/test_flava.py | 52 ++++++++++++++++++++- torchmultimodal/models/flava/model.py | 50 ++++++++++++++++++-- torchmultimodal/models/flava/transformer.py | 18 ++++++- 3 files changed, 114 insertions(+), 6 deletions(-) diff --git a/tests/models/flava/test_flava.py b/tests/models/flava/test_flava.py index 4080a846..6a6097c4 100644 --- a/tests/models/flava/test_flava.py +++ b/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 \ No newline at end of file diff --git a/torchmultimodal/models/flava/model.py b/torchmultimodal/models/flava/model.py index ccf51250..4af7fc87 100644 --- a/torchmultimodal/models/flava/model.py +++ b/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,6 +437,7 @@ 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, @@ -404,6 +445,7 @@ def forward( 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) \ No newline at end of file diff --git a/torchmultimodal/models/flava/transformer.py b/torchmultimodal/models/flava/transformer.py index b5a273ff..8ce0f555 100644 --- a/torchmultimodal/models/flava/transformer.py +++ b/torchmultimodal/models/flava/transformer.py @@ -49,6 +49,13 @@ 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") @@ -56,6 +63,15 @@ def forward( 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) \ No newline at end of file