Repository: pytorch
Issue: More general MultiHeadAttention and Transformer modules
Labels: triaged,enhancement


## 🚀 Feature

<!-- A clear and concise description of the feature proposal --> 

Modularize `nn.MultiHeadAttention` / `nn.functional.multi_head_attention_forward()` and generalize the `nn.Transformer` modules so that:
1. Novel attention functions can be used with minimal effort (i.e., without having to rewrite the entire Transformer architecture), and
2. The attention functionality can be used in modules other than `nn.MultiHeadAttention`.
3. There is more flexibility for the `nn.Transformer` sublayers/components

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

Recently, there has been a huge amount of interest in making Attention -- the core component of Transformer modes -- more efficient and performant by reducing the L^2 scaling of attention (where L is the length of a sequence) by making various approximations. E.g.:
- [Adaptive Attention](https://arxiv.org/abs/1905.07799)
- [Locality Sensitive Hashing (Reformer)](https://arxiv.org/abs/2001.04451)
- [Routing Attention](https://arxiv.org/abs/2003.05997)
- [Random Fourier Features (Performer)](https://arxiv.org/abs/2009.14794)
- and lots more

Currently, to use any new Attention modules in PyTorch, the entire `nn.MultiHeadAttention` and `nn.Transformer` modules needs to be rewritten making experimentation of new Attention modules slower than it needs to be. Attention has also been used in other novel contexts such [Product Key Memory](https://arxiv.org/abs/1907.05242) and it would be beneficial to have different attention approximations available for these contexts too.

Furthermore, novel sublayers (such as the Product Key Memory) have been used to replace the feedforward network which is not possible with the current implementation of `nn.Transformer`.

## Pitch

<!-- A clear and concise description of what you want to happen. -->

### General Idea: MultiHeadAttention
Separate the `nn.MuliHeadAttention` / `nn.multi_head_attention_forward()` functionality into three sections:
1. Create the query-key-value
2. Split / merge heads
3. Attention

It would look like:
```python
def query_key_value(...):
    ...

def split_heads(...):
    ...

def merge_heads(...):
    ...

def attention(...):
    ...
```

and have corresponding Modules for each one:
```python
class QueryKeyValue(nn.Module):
    ...

    def forward(self, ...):
        return query_key_value(...)

class SplitHeads(nn.Module):
    ...

    def forward(self, ...):
        return split_heads(...)

class MergeHeads(nn.Module):
    ...

    def forward(self, ...):
        return merge_heads(...)

class Attention(nn.Module):
    ...

    def forward(self, ...):
        return attention(...)
```

then have a `nn.MultiHeadAttention` module that looked like:
```python
class MultiHeadAttention(nn.Module):
    def __init__(self, ...):
        ...
        self.query_key_value = QueryKeyValue(...)
        self.split_heads = SplitHeads(...)
        self.merge_heads = MergeHeads(...)
        self.attention = Attention(...)

    def forward(self, ...):
        q, k, v = self.query_key_value(...)
        q = self.split_heads(q)
        k = self.split_heads(k)
        v = self.split_heads(v)
        x, a = self.attention(...)
        x = self.merge_heads(x)
        return x, a
```
Users would then rewrite the MultiHeadAttention module using their own custom `Attention` module, reusing the other modules and using the above implementation as a template.

### General Idea: Transformer
One glaring problem with this is that all the Transformer modules are hard-coded for the current `nn.MultiHeadAttention` implementation. The only way I can currently see around that is to actually pass the Transformer modules a `nn.MultiHeadAttention` module (or whatever your custom version is) in a similar way `nn.Module`s are passed to `nn.Transforer{En,De}coder` or even `nn.Sequential`. A huge benefit to this is that you could pass a sublayer1 and sublayer2 to the `nn.Transformer` modules and mix-and-match sublayers:
```python
class Transformer(nn.Module):
    def __init__(
        self,
        encoder_sublayer1: nn.Module,
        encoder_sublayer2: nn.Module,
        decoder_sublayer1: nn.Module,
        decoder_sublayer2: nn.Module,
        ...
     ):
        # init stuff here
```
and the same for `nn.TransformerEncoderLayer` etc. That would mean that you could also do things like:
1. Easily swap the order of the sublayers,
2. Easily use alternatives to the 2-layer relu feedforward network (like a deeper feedforward network, a feedforward network with a novel activation function or Product Key Memory)
3. Use different architectures/sizes for the encoder and decoder if desired

### Specific implementations
To test and demonstrate the above ideas, I also propose adding (at least) the following specific implementations:
* [Product Key Memory](https://arxiv.org/abs/1907.05242)
* [Random Fourier Features](https://arxiv.org/abs/2009.14794)

I currently have an implementation of the former which is being used at my company (although it would probably need to be reworked a little bit to fit this paradigm), and am working on an implementation of the latter.

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

### MultiHeadAttention
Rewrite the `nn.MultiHeadAttention` module something like:
```python
class MultiHeadAttention(nn.Module):
    def __init__(self, ...):
        ...

    def query_key_value(self, ...):
        ...

    def split_heads(self, ...):
        ...

    def merge_heads(self, ...):
        ...

    def attention(self, ...):
        ...

    def forward(self, ...):
        q, k, v = self.query_key_value(...)
        q = self.split_heads(q)
        k = self.split_heads(k)
        v = self.split_heads(v)
        x, a = self.attention(...)
        x = self.merge_heads(x)
        return x, a
```
and then overload the `attention` method with different kinds of attention. However, this is definitely less flexible than the proposed solution.


## Additional context

<!-- Add any other context or screenshots about the feature request here. -->

The above proposal is based on the NLP applications I am familiar with. I've noticed recently that attention/transformers are being applied in computer vision, however I am unfamiliar with if and how they differ from NLP. Input from someone more experienced in computer vision would be greatly appreciated!

cc @zhangguanheng66