Expose Attention Weights from MultiHeadDotProductAttention

[Xiaoming-Zhao]

Hi, I am wondering whether it is possible to support returning attention weights computed by MultiHeadDotProductAttention, similar to what is proposed #2869 (but instead of sow, I am hoping that we can directly support returning it more easily).

I understand that there is a function dot_product_attention_weights to compute attention weights. However, it is not trivial to get the exact weights MultiHeadDotProductAttention used during the forward pass as there are projections etc.

The ways I am thinking about to obtain the attention weights is to let dot_product_attention directly return the weights it has already computed and change __call__ of MultiHeadDotProductAttention accordingly.

I am happy to create a PR if the above sounds reasonable to you.

Besides the available option in tensorflow as discussed in #2869, torch also supports returning attention weights by setting need_weights=True, which are quite helpful.

Arguably, the attention weights are one of the most important features that are worth investigating. They are quite useful in my personal experiences across various projects.

Any suggestions would be much appreciated. Thanks a lot in advance.

[chiamp]

This feature has been added in #3529. Pass return_weights=True to MultiHeadDotProductAttention's __call__ method to sow the attention weights into the intermediates collection.