3-dimensional attention_mask in LongformerSelfAttention

# ❓ Questions & Help

<!-- The GitHub issue tracker is primarly intended for bugs, feature requests,
     new models, benchmarks, and migration questions. For all other questions,
     we direct you to the Hugging Face forum: https://discuss.huggingface.co/ .
     -->

## Details

<!-- Description of your issue -->

<!-- You should first ask your question on the forum, and only if
     you didn't get an answer after a few days ask it here on GitHub. -->
Hi everyone, I am learning NLP and I am doing a project that use ```distilroberta``` as an encoder model. I have read about [K-BERT](https://arxiv.org/pdf/1909.07606), and I want to implement that model with ```distilroberta```. To say it briefly, **K-BERT** has a parallel Knowledge Graph with it. Each word in an input sentence will be queried from that Knowledge Graph and the information is injected to the input sentence beside that word. The information from the Knowledge Graph is just related to their token onwer. So, there is a visible matrix that decides which tokens are related to a token, it acts as an ```attention_mask``` an ```attention_mask``` is usually ```[batch, seq_len]``` but this kind of attention mask is ```[batch, seq_len, seq_len]```, and ```position index``` to indicate which token is related to the others. Therefore, when using distilroberta, I need to pass ```position``` as ```position_ids``` and ```visible matrix``` as ```attention_mask``` argument. Thankfully, in ```modeling_roberta``` it accept the ```attention_mask``` with 3 dimension and create an extended attention mask as ```[:, None, :, :]```, my visible matrix works. However, When I try to convert ```distilroberta``` to ```longformer```, there is a problem in ```LongformerSelfAttention```: 
```python 
# values to pad for attention probs
remove_from_windowed_attention_mask = (attention_mask != 0)[:, :, None, None]
```
If ```attention_mask``` is ```[batch, seq_len]```. After transforming like that, ```remove_from_windowed_attention_mask``` will be ```[batch, seq_len, 1, 1]```. However, with my ```visible matrix```, it results in ```[batch, seq_len, 1, 1, seq_len]```, then that has an error at: 
```python 
def _sliding_chunks_query_key_matmul(self, query: torch.Tensor, key: torch.Tensor, window_overlap: int):
        """
        Matrix multiplication of query and key tensors using with a sliding window attention pattern. This
        implementation splits the input into overlapping chunks of size 2w (e.g. 512 for pretrained Longformer) with an
        overlap of size window_overlap
        """
        batch_size, seq_len, num_heads, head_dim = query.size()
        assert (
            seq_len % (window_overlap * 2) == 0
        ), f"Sequence length should be multiple of {window_overlap * 2}. Given {seq_len}"
        assert query.size() == key.size()
 
 ValueError: too many values to unpack (expected 4)
```
I actually do not know how to solve this problem. I am afraid whether I change that attention mask transformation to be suitable for my ```visible matrix```, it will affect something unexpectedly or not. I would be really appreciated because of your help. 

**A link to original question on the forum**: https://discuss.huggingface.co/t/3-dimensional-attention-mask-in-longformerselfattention/16496?u=khangnguyen2907

<!-- Your issue will be closed if you don't fill this part. -->

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.