SelfAttention misses Linear after attention, wrong for Conformer, Transformer

[albertz]

There is a linear projection after the attention.

ESPNet MultiHeadedAttention has it.
PyTorch torch.nn.MultiheadAttention does not have it.
Keras tf.keras.layers.MultiHeadAttention has it.
torchaudio.models.wav2vec2.components.SelfAttention has it.
Fairseq MultiheadAttention has it.

Our nn.GenericSelfAttention (and thus nn.SelfAttention) does not have it.
The RETURNN SelfAttentionLayer also does not have it.

But then we also don't have it in ConformerEncoderLayer, so it's clearly missing.
Also we don't have it in our Transformer, so it is missing there as well.

So, should we change nn.GenericSelfAttention? Always include it? Or optionally include it? Make it a required argument that there is no confusion about it, like out_dim: Optional[nn.Dim] (without default). In case the user sets None, no linear transformation at the end, otherwise there is.

If we don't change nn.GenericSelfAttention, we must fix the Transformer and Conformer.

[JackTemaki]

It should not be part of nn.SelfAttention, but definitely be part of the ConformerEncoderLayer and other related layers.

[albertz]

It should not be part of nn.SelfAttention

Why? This might confuse everyone else (everyone who has not used maybe RETURNN SelfAttentionLayer before). This also now lead to these quite serious bugs, so it even confused ourselves.

[JackTemaki]

I think now that I looked at the code the whole naming is flawed right now. In our case "GenericSelfAttention" is not generic at all, but is fixed to a dot attention as far as I can see. our current GenericSelfAttention is rather a MultiHeadedDotSelfAttention, and then yes, it should include the Linear layer to be correct. (Or lets say compatible to ESPNet and what people expect from the layer)

[albertz]

GenericSelfAttention is called "generic" because it is the base class for SelfAttention and CausalSelfAttention.

[JackTemaki]

Yes, but this is highly misleading. Lets say I now implement an MLP based self-attention instead of a dot based, how would I call that then? And I would expect this should inherit from GenericSelfAttention as template, but this would be incorrect.

[albertz]

I think self-attention always uses dot product. I have never seen sth else. So this is pretty standard. You can of course always implement sth custom, but then this would be a separate implementation. Why would you expect to inherit from sth else if you would not share any code with the standard self-attention? I don't think this is a problem.

Look at any other framework. It's always just called SelfAttention for self-attention or MultiHeadAttention (or alike) for dot attention.

[albertz]

Change nn.GenericSelfAttention? Make it a required argument that there is no confusion about it, like out_dim: Optional[nn.Dim] (without default). In case the user sets None, no linear transformation at the end, otherwise there is.

I tend to prefer this variant, i.e. having it explicit. The majority of other code seems to have it included, just like the original paper describes it, but there are some exceptions.

[JackTemaki]

I think self-attention always uses dot product.

I just wanted to clarify that this is a hidden implication.

Why would you expect to inherit from sth else if you would not share any code with the standard self-attention?

If dot_attention would not be called explicitly but be e.g. an overridable function, the GenericSelfAttention would be more "Generic", all I am saying it that it makes more fixed assumptions than the name hints at.