Restructure `multi_head_attention_forward`

[Enealor]

🚀 Feature

Restructure the function multi_head_attention_forward in nn.functional into several functions to improve the ability to experiment. In particular, decompose the function so that the following are available:

• The input embedding functions.
• The computation of attention weights.
• The output embedding function.

This will allow users to try different embeddings or attention mechanisms without having to recode the rest.

Motivation

Addresses the issue of decomposing the function as mentioned in #32590. It also moves forward on including more support for attention mechanisms.

Pitch

Currently, the mutli_head_attention_forward function encapsulates the projection of the query, key, and value, computing attention for these projections, and computing the output projection after applying attention. Furthermore, the input embedding utilizes several code paths that are different embeddings. By decomposing the function into several parts, we can make it more readable and open to experimentation.

The following plan is based on the above:

• Functions for computing the input embeddings q, k, and v. There are currently four code paths used for doing this, and three unique embeddings are used. Each embedding should be an individual function so that it's clearer what method is being used. The embeddings used are labeled 'self-attention' (where query = key = value), 'encoder-decoder attention' (where key = value) and one that is unlabeled but is probably just called attention. The last embedding has two code paths depending on whether in_proj_weight is used or separate weights are used for query, key and value. (See L3669-L3748.)
• A function for applying attention to get a new query. Some models rely on computing attention in different ways, and separating this out would allow us to use those more freely. This should optionally return the attention weights. Specifically, this is the Scaled Dot-Product Attention. (See L3750-L3824.)
• A function for computing the output projection of the query. There is currently only one function needed for doing this. (See L3826-L3836.)

Alternatives

Some of the restructurings I have suggested could be skipped in favor of introducing fewer functions. In particular, only one of the input embeddings needs to be provided. The rest could be left to the end-user.

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki

[ngimel]

cc @zhangguanheng66 @cpuhrsch.
@Enealor do you want to work on this? If so, can you please coordinate with @zhangguanheng66 and @cpuhrsch .
Thanks!

[Enealor]

@ngimel I am happy to coordinate with them. The issue that spawned this asked for helpers to first open an issue and wait for feedback. I am willing to work on it further given the go-ahead.

[cpuhrsch]

@Enealor - We're actively working on this for several of the reasons you stated. We'll need review / input of our designs and eventually also work to implement. I'm happy to see we have community support for this and are happy to loop you in.

[zhangguanheng66]

Definitely. As the first step, please split the MHA function into several building blocks in nn.Module and nn.functional. Something like this should be a good start. Feel free to ping us once you have a PR.

Please keep in mind that those building blocks should be JIT-ability, quantization, and potentially porting to C++ in the future.

[zhangguanheng66]

@Enealor I starts a PR here. Feel free to continue the work if you have time. After we collect some feedback, we can move the code to torch.nn.functional

[Enealor]

@zhangguanheng66 Thank you for the link. I have started restructuring the code in functional. Based on my work so far, I have changed my pitch and added an alternative.

I could use feedback on how to handle the input embeddings. As I explain above, there are three different ways the input embedding is calculated. Should I add a function for each, or provide one function that can be used by the end-user to build up the rest? I ask because I had started the former, while I noticed in the PR that you are doing the latter.

[Enealor]

Follow-up: I have created a PR with the approach that would introduce less specialized functions. I believe it is ready for feedback.

[zhangguanheng66]

Follow-up: I have created a PR with the approach that would introduce less specialized functions. I believe it is ready for feedback.

Thanks. @Enealor Made some comments there. Just wondering if you want to merge your work to this one. Then, you can take charge of this PR from now on.

[netw0rkf10w]

Hi guys!
Could you please tell me what is the current state of this?
I would like to implement a different method for computing the attention weights, and was about to split the current multi_head_attention_forward into smaller functions, but then fortunately I found this thread so I thought maybe I shouldn't start from scratch.
Thanks a lot!

[zhangguanheng66]

The new MHA has been landed in torchtext.
https://github.com/pytorch/text/blob/60907bf3394a97eb45056a237ca0d647a6e03216/torchtext/modules/multiheadattention.py#L5