IDEFICS Cross Attention: Text tokens appearing before images still attend to image embeddings
Author: momergul
Pull request: no

### System Info

- `transformers` version: 4.33.1
- Platform: Linux-5.4.0-153-generic-x86_64-with-glibc2.31
- Python version: 3.9.18
- Huggingface_hub version: 0.17.1
- Safetensors version: 0.3.3
- Accelerate version: 0.23.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

1: Run the following code snippet altered from `examples/idefics/inference.py` in the notebooks repo. 

```                                                                                                                                                                                                           
import torch
from transformers import IdeficsForVisionText2Text, AutoProcessor

device = "cuda" if torch.cuda.is_available() else "cpu"

checkpoint = "HuggingFaceM4/idefics-9b"
model = IdeficsForVisionText2Text.from_pretrained(checkpoint, torch_dtype=torch.bfloat16).to(device)
processor = AutoProcessor.from_pretrained(checkpoint, use_auth_token=False)
model.eval()

url = "https://hips.hearstapps.com/hmg-prod/images/cute-photos-of-cats-in-grass-1593184777.jpg"
image = processor.image_processor.fetch_images(url)
prompts = [
    [
        "User:",
        image,
        "Describe this image.\nAssistant: An image of two kittens in grass.",
    ],
]


inputs = processor(prompts, return_tensors="pt").to(device)
logits = model(**inputs)['logits']
```

2: During the model forward pass, inspect hidden states in Line 912 of `models/idefics/modeling_idefics.py` 

### Expected behavior

Hello! I believe there is a bug in how cross attention is performed within `IdeficsGatedCrossAttentionLayer` in `models/idefics/modeling_idefics.py` for text tokens appearing before any images are given to the model. As IDEFICS is autoregressive, the hidden state for a text token appearing before any image is observed should not be changed after cross attention. During the forward pass in the code snippet I provided, I expect the following behavior immediately after Line 911 of `models/idefics/modeling_idefics.py`:

Expected behavior:
`torch.all(residual[0, 0:4] == hidden_states[0, 0:4])` evaluates to `True`

Observed behavior:
`torch.all(residual[0, 0:4] == hidden_states[0, 0:4])` evaluates to `False`

I believe this is due to how the attention mask is applied. For the first 4 tokens which appear before any image, all values of `image_attention_mask` are set to the smallest possible value. This results in the attention weights during the call to `nn.functional.scaled_dot_product_attention` in Line 692 to each be equal to each other. This in turn means that these four text tokens appearing before any image each attend to the image embeddings. 

Is my understanding correct here? I would greatly appreciate it if you could look into this.




Comments:
What do you think @leot13 @VictorSanh ?

Thank you for noticing! It's not easy to detect. We are aware but did training this way. In practice that means the few first tokens with no image are attending to every image instead of none of them, so there's a small information leak.
To fix this, we could apply the image_attention_mask on the output of the cross-attention as a gating mechanism. The image attention mask has shape [bsz, num_tokens, num_images] so we would need to use a gating mechanism along the lines of:
`residuals + self.act_cross_attn(self.alpha_cross_attn) * image_attention_mask.sum(dim=2).unsqueeze(-1) *  cross_attention_hidden_states `
However, it's not certain that the performance would transfer perfectly since this is a different setup from the training one. We would probably need to re-evaluate them on some benchmarks to make sure inference in this setup is fine. Most likely it will be. At least for the instruct ones since we do some finetuning on ultrachat, a text-only dataset for which we zero-out the cross-attentions.

Thanks for the response! I was able to notice only because I began receiving NaNs in the outputs of the cross attention layer for tokens appearing before images while doing QLoRA finetuning. How were you able to avoid this during training? During cross attention, if `(Q @ K.transpose(-2, -1) / math.sqrt(Q.size(-1)))` is sufficiently small and negative, there is a chance that adding the attention mask for these tokens before images will result in -inf for each value, causing NaNs after softmax.

I have been trying to reproduce you NaNs issue, but can't so far. There is a [colab notebook](https://colab.research.google.com/drive/1RltyDpv7Fbu_My03RyZ7ftavEQyoxbek#scrollTo=prXRsUiXCII9) for doing QLoRA PEFT finetuning. I used a similar setup, using almost the same libraries as you (except for cu17 which doesn't work in my env, so I used cu18) and didn't get NaNs even when placing text before the image.

Did you perform the QLoRA fine tuning with the same setup as described in the colab? 

Also side note: the image_attention_mask I described in the comment above is the one fed to the model, but it gets modified before reaching the cross-attention block. The idea stays the same though.

I have been finetuning IDEFICS on a separate task with unreleased data and have also not been using the Trainer module for finetuning, so there is a good chance I am introducing some error of my own for the NaNs. I also recently had the same NaN problem with padding tokens in regular self-attention (where the pad tokens also have an attention mask with all entries set to the smallest value), so the NaN problem I have is not about cross-attention. I'll see if I can replicate the problem with publicly available data and share my code in a separate repository. Thanks for looking into this!

As an aside, the problem itself is essentially discriminative image captioning, with the model being fed 10 images and being asked to produce a caption for a target image. To give some more information, the model input in training is structured in this manner:
```
prompt = [
"Image 0", img_0, "Image 1", img_1, ..., "Image 9", img_9,
"Instruction: You will provide a discriminative caption for the target image.",
f"The target image is Image {target_idx}. Caption: {caption}"
]
```


PR looks good and should be merged soon 🤗 thanks for you patience 

I believe the problem @momergul mentioned is related to the `pytorch 2.1` release as mentioned in https://github.com/pytorch/pytorch/issues/110213. In which, the output of `nn.functional.scaled_dot_product_attention` will be `NaN` if the whole row is masked. It is the situation I met when I am finetuning the Idefics model. Would you mind provide any hints to fix this?

Thanks for bringing this back to my attention!
[PR #26839](https://github.com/huggingface/transformers/pull/26839) answers this issue:
> Expected behavior
> Hello! I believe there is a bug in how cross attention is performed within IdeficsGatedCrossAttentionLayer in models/idefics/modeling_idefics.py for text tokens appearing before any images are given to the model. As IDEFICS is autoregressive, the hidden state for a text token appearing before any image is observed should not be changed after cross attention. During the forward pass in the code snippet I provided, I expect the following behavior immediately after Line 911 of models/idefics/modeling_idefics.py:
> 
> Expected behavior:
> torch.all(residual[0, 0:4] == hidden_states[0, 0:4]) evaluates to True
> 
> Observed behavior:
> torch.all(residual[0, 0:4] == hidden_states[0, 0:4]) evaluates to False
> 
> I believe this is due to how the attention mask is applied. For the first 4 tokens which appear before any image, all values of image_attention_mask are set to the smallest possible value. This results in the attention weights during the call to nn.functional.scaled_dot_product_attention in Line 692 to each be equal to each other. This in turn means that these four text tokens appearing before any image each attend to the image embeddings.
> 
> Is my understanding correct here? I would greatly appreciate it if you could look into this.

You are right though that it does not answer the NaNs problem. I looked into the issue you showed back then but 2 things made me think it was okay: the torch version specified by @momergul was: 2.0.1+cu117 , and when I double checked for torch 2.1, I did the inference in bf16. 
It turns out, as I realized when diving back into this after seeing you message, this bug only happens in fp32. 
The best quick fix right now would be to use bf16, or fp16 if you don't have access to bf16-friendly gpus.

Otherwise, since it affects multiple models, I believe there will be a fix soon in transformers. 
@ydshieh for visibility

Hi @momergul @king159 the issue in IDEFICS related to PyTorch SDPA being used in fp32 on CUDA device with the memory-efficient attention backend yielding NaN (https://github.com/pytorch/pytorch/issues/110213) will be fixed with https://github.com/huggingface/transformers/pull/26572.