add key_value_mask argument for cross attention

[jglaser]

This change set enables the use of two different masks for queries and keys/values. This is necessary for using the module as a cross attention.

[lucidrains]

@jglaser hello there! looks good :) have you had much luck with this efficient attention variant?

[lucidrains]

hmm actually, how would cross attention work if the tokens and the cross-attended contexts have different dimensions? i thought this scheme requires a square matrix

[jglaser]

@jglaser hello there! looks good :) have you had much luck with this efficient attention variant?

I am still testing if it results in better model performance for long sentences, but so far it looks good in terms of memory efficiency

[jglaser]

hmm actually, how would cross attention work if the tokens and the cross-attended contexts have different dimensions? i thought this scheme requires a square matrix

You're right, I have only tested this for a square matrix. Both in sentence length, and in hidden dimension.

I achieve the former by extending the shorter of the two input sentences with the pad token, and for the latter I implement a linear layer that translates between keys and values of different hidden dimensions (per token).

Both can be done either in the forward function outside the self attention module, or by replacing the to_qkv and to_out member variables with custom layers, without any further changes to this module.

That said, we may eventually want to generalize this self-attention to accept inputs for the queries and keys/values of different dimensions.

[jglaser]

Closing this: there is no need for a separate mask for queries and keys/values. Queries can always be masked by truncating the output of the self attention layer appropriately, only a single mask (for keys and values) is needed.