Return updated attention mask from Wav2Vec 2.0
Author: gau-nernst
Pull request: no
Labels: name: WIP

### Feature request

In Wav2Vec 2.0, the first few convolution layers affect the attention mask. Thus, if I want to use all Wav2Vec 2.0 outputs (last_hidden_state), I need access to the updated attention mask. Currently the workaround is to call the private method `._get_feature_vector_attention_mask()`. To return this attention mask directly will make the experience much better.

Concretely, it can be an additional field in Wav2Vec2BaseModelOutput.

Additionally, data2vec-audio and HuBERT also have the same problem, since they are also based on Wav2Vec 2.0.

### Motivation

Fairseq implementation returns the updated attention mask for downstream uses.

https://github.com/facebookresearch/fairseq/blob/100cd91db19bb27277a06a25eb4154c805b10189/fairseq/models/wav2vec/wav2vec2.py#L696-L700

Then their Wav2Vec 2.0 for classification implementation is more elegant.

https://github.com/facebookresearch/fairseq/blob/100cd91db19bb27277a06a25eb4154c805b10189/fairseq/models/wav2vec/wav2vec2_classification.py#L84-L93

In HF, Wav2Vec 2.0 for classification requires calling the private method.

https://github.com/huggingface/transformers/blob/641adca55832ed9c5648f54dcd8926d67d3511db/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L2127

### Your contribution

I can help to submit a PR for this. It will be an additional field in Wav2Vec2BaseModelOutput. I'm not quite sure about other types of Wav2Vec2 output objects.

Comments:
cc @sanchit-gandhi 

Hey @gau-nernst - thanks for the great issue report. Having gone through the fariseq code, I could see this being a valuable contribution, since the original implementation does indeed return the attention mask for downstream use cases. However, I'm not entirely sure what such use cases would be? Perhaps you could provide a few examples based on your use cast?

Since Wav2Vec2 is quite a used model, we'd want to be quite certain on the improvements that we'd get by changing the signature of the model output before jumping into a PR!


Virtually any processing to be done on `last_hidden_state` requires the updated attention mask to avoid including padded data. Some use cases
- (my current use case) Extract embeddings from audio samples using various wav2vec 2.0 models. After getting `last_hidden_state`, I need to do pooling (either mean or max). Padded data should be excluded, thus the updated attention mask is required.
- Add a seq2seq module on top of wav2vec 2.0 e.g. transformer, ECAPA-TDNN. Again, the updated attention mask is required.

In HF, almost (if not all) derived Wav2Vec 2.0 models (e.g. Wav2Vec2ForCTC, Wav2Vec2ForSequenceClassification, Wav2Vec2ForXVector) require calling the private method, because "any processing to be done on `last_hidden_state` requires the updated attention mask". Although HF provides some custom Wav2Vec2 for downstream tasks, power users like myself prefer writing our own extra modules, either for more flexibility or because we want to use new/SOTA modules that are not available in HF.

Regarding changing the model signature, I have always thought that HF wraps model output in a dataclass so that you can always extend it later without breaking old code. There is no change to current fields. The breaking change would be when a tuple is returned (`return_dict=False`).

Thanks for the descriptive explanation @gau-nernst - I agree that this definitely does help power users like yourself, and facilitates the community building on-top of HF Wav2Vec2 models.

Let's go for it and add the `attention_mask` field in a PR! Would you like to open a PR to do this?

@sanchit-gandhi yes, sure, I will make a PR for this and request you to review.