Support for 4D attention mask for T5

[Aethor]

Feature request

Currently, T5 cannot take 4D attention masks (batch_size, num_heads, seq_len, seq_len) as inputs. Passing a 4D attention_mask and a 4D decoder_attention_mask like so leads to a shape-related exception :

import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration

tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small")
model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small")

input_ids = tokenizer("Where is", return_tensors="pt").input_ids
decoder_input_ids = tokenizer("<pad>", return_tensors="pt").input_ids

batch_size, seq_len = input_ids.shape
tgt_len = decoder_input_ids.shape[1]
num_heads = model.config.num_heads

attention_mask = torch.ones(batch_size, num_heads, seq_len, seq_len)
decoder_attention_mask = torch.ones(batch_size, num_heads, tgt_len, tgt_len).tril(0)

model(
    input_ids,
    decoder_input_ids=decoder_input_ids,
    attention_mask=attention_mask,
    decoder_attention_mask=decoder_attention_mask,
)

One of the problems in the current code is in the handling of the cross-attention mask. Currently, it is created using the 1D encoder attention mask when supplied. However, in the case of a 4D mask, it seems unclear how to correctly use the encoder mask: therefore, the best solution might be to introduce a new 4D mask argument cross_attention_mask of shape (batch_size, num_heads, tgt_len, seq_len)`. This lets the user controls all attention masks if necessary.

Motivation

4D masks are useful for many purposes, as outlined by #27539 and this blog post, but not all models support them.

Your contribution

I propose to fix the code to handle 4D attention masks, and to add a new cross_attention_mask argument to add the possibility to control the cross attention mask manually. I wrote a version of that code in this fork.

I'm happy to create a PR with my code, but:

1. This is my first transformers contribution, I need help with some things such as handling the "Copy" code duplication mechanism of transformers. Should other similar models with copied functions from T5 be changed as well?
2. Although I wrote a first test with trivial masks, I am not entirely sure how to test this
3. I want to be sure that adding the new cross_attention mask parameter is the right way to do this and will be approved

[Rocketknight1]

cc @ArthurZucker for text models

[vasqu]

Hey @Aethor

I think it would be better to work on integrating/creating the new attn mask interface instead of reworking on old mask functions, which will be deprecated at some point. I.e. encoder mask and encoder-decoder mask functions similar to

transformers/src/transformers/masking_utils.py

Lines 744 to 780 in def7558

	def create_causal_mask(
	config: PretrainedConfig,
	input_embeds: torch.Tensor,
	attention_mask: Optional[torch.Tensor],
	cache_position: torch.Tensor,
	past_key_values: Optional[Cache],
	position_ids: Optional[torch.Tensor] = None,
	or_mask_function: Optional[Callable] = None,
	and_mask_function: Optional[Callable] = None,
	) -> Optional[Union[torch.Tensor, BlockMask]]:
	"""
	Create a standard causal mask based on the attention implementation used (stored in the config). If `past_key_values`
	has an hybrid cache structure, this function will return the mask corresponding to one of the "full_attention" layers (to align
	to what is needed in the `modeling_xxx.py` files).
	Args:
	config (`PretrainedConfig`):
	The model config.
	input_embeds (`torch.Tensor`):
	The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the
	batch size, query length and dtype.
	attention_mask (`torch.Tensor`, optional):
	The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).
	It can also be an already prepared 4D mask, in which case it is returned as-is.
	cache_position (`torch.Tensor`):
	A tensor of shape (query_length,) indicating the current indices of the input sequence elements.
	past_key_values (`Cache`, optional):
	The past key values, if we use a cache.
	position_ids (`torch.Tensor`, optional)
	A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.
	or_mask_function (`Callable`, optional):
	An optional mask function to combine with the causal mask function (by doing the union of both). This is
	useful to easily overlay another mask on top of the causal one, for example for image tokens handling.
	and_mask_function (`Callable`, optional):
	An optional mask function to combine with the causal mask function (by doing the intersection of both). This is
	useful to easily overlay another mask on top of the causal one, for example for image tokens handling.
	"""

The old mask functions are less flexible while the new mask function will also be extendable to new custom attentions + it already allows for passing custom 4D masks.

To answer your questions:

• T5 is also quite old, so changing the mask logic there shouldn't affect many models. If there is copied code from T5 (on the parts you modify), you can apply these changes via make fix-copies. Some docs are also at https://huggingface.co/docs/transformers/contributing (not completely up to date but the utilities are the same).
• A test like this would fall under the base T5 model tests, e.g. something like https://github.com/Aethor/transformers/blob/22dc62edbdbc3f2afeb90a31c75047711c1afc5c/tests/models/t5/test_modeling_t5.py#L1039-L1041. The test should also check for padding to work.
• The cross attention mask == attention mask, no need to add another argument here. See

  transformers/src/transformers/models/t5/modeling_t5.py

  Line 1513 in def7558

  encoder_attention_mask=attention_mask,

  We should only need to check if that mask was already 4D and ignore any preparation (assumes correctly prepared mask).

[liygzting]

mark

[Aethor]

Hi @vasqu, thank you for your answer. I agree that if there is a new interface for attention masks, it should be prioritized. I have no problem working on that, if this is something that is welcomed for T5 (or other similar old models)? Regarding the specific points you answered:

• thanks, I did try make fix-copies but I didn't know if the scope of the changes were acceptable.
• thank you for the pointer.
• mmh, it seems to me that an other argument should be added here actually. If I understand correctly, the shape of a 4D attention_mask is (batch_size, num_heads, seq_len, seq_len), but the shape of the cross attention mask should be (batch_size, num_heads, target_len, seq_len) since decoder outputs are attending encoder inputs. This is why I introduced a new argument: if we simply use the attention mask which is what I did at the beginning, we have another shape issue.

To summarize, for now I'm happy using my little hack for my specific needs, and I understand that it's not ideal to change the existing old implementations. If converting T5/other similar older models to use the new attention interface is a welcome addition, I'm open to working on it.

[vasqu]

Going through bit by bit @Aethor

I have no problem working on that, if this is something that is welcomed for T5 (or other similar old models)?

Very much so! The masks would transfer directly to most encoder and cross-attention models. (There are only a few specific outliers)

thanks, I did try make fix-copies but I didn't know if the scope of the changes were acceptable.

That's not a problem (as long as it doesn't break other models - we have tests for that to check).

This is why I introduced a new argument: if we simply use the attention mask which is what I did at the beginning, we have another shape issue.

Ah ic, so we have the encoder mask for the encoder and the cross attn - in 2D it doesn't pose problems but in 4D we get the mismatch. It sounds like a reasonable solution. However, I'm a bit unsure about the usefulness here and if it's worth the maintenance as we want the reduce the amount of super specific features. I'm down for it if more of the community reacts to it.

If converting T5/other similar older models to use the new attention interface is a welcome addition, I'm open to working on it.

Definitely would be!