Add support for Keras mask & causal mask to MultiHeadAttention

[ageron]

This is a first go at adding Keras masking support and causal support to the MultiHeadAttention layer.

See the discussion in #16248.

[fchollet]

Thank you for the PR! For consistency, I would favor the same API as the one used in KerasNLP here: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/layers/transformer_decoder.py#L191

As in, a use_causal_mask argument in call(). It would not change the PR much otherwise.

[fchollet]

FYI @mattdangerw @chenmoneygithub

[ageron]

Hi @fchollet , thanks for reviewing this PR.

Note that the keras.layers.Attention layer uses an init parameter called causal, as noted by @haifeng-jin in #16248, so it seems to me that consistency within Keras is more important than consistency with Keras-NLP?

That said, I just made the change you requested, and I'm happy to push it if you prefer, I don't have a strong opinion about this: perhaps it's the Attention layer that should be updated to use use_causal_mask in the call() method.

[fchollet]

This is a good question -- we should make a choice one way or the other and standardize everything (Keras, KerasNLP) on that choice. Let me chat with the team and we'll figure out what to do.

[chenmoneygithub]

@fchollet To me causal is not the best argument name as it is too abstract. I did a codesearch and there are only a few usages of the causal argument of tf.keras.layers.Attention in google3, so arg name changing won't be too hard.

[ageron]

Sounds good. I just replaced causal in the constructor with use_causal_mask in the call() method.

[fchollet]

I also think it makes sense to have the argument in call() because that's where we pass the other masking-related arguments. So let's keep use_causal_mask in call().

@chenmoneygithub could you please file an issue to change the Attention layer API to adopt this convention?

[chenmoneygithub]

@fchollet Sure!

[fchollet]

LGTM

[mattdangerw]

@ageron did some more testing a came across a few things.

1. I think we need to be more defensive about assuming mask types. We have real world usage where the mask (explicit or implicit) is floating point, which leads to errors when doing the bitwise &. Maybe we can just make sure we cast all masks to bool type before working with them? We should probably add a test for this too.

2. In an encoder where QKV are all the same, with an implicit mask, this code looks like it will make an attention mask with zeros along the bottom and along the right. I am used to seeing an attention mask with zeros only along one of those dimensions. This colab the difference. I am reading your code correct? Do you know what the right approach here is?

[ageron]

Hi @mattdangerw ,
Sorry for the late response, got Covid in the family.
Thanks for the careful review, you make two good points.
I'll take care of the mask types.
However, regarding the padding difference, after some thought I'm not sure the difference matters, since it only affects masked out tokens, so normally these outputs will be ignored anyway.

For example, let's take a three word sentence padded to 5 tokens like this:

"I am happy <pad> <pad>"

With causal self-attention, we want "I" to attend only to "I", "am" should attend only to "I am", and "happy" should attend only to "I am happy". So far so good. Next, it doesn't really matter what both padding tokens attend to, since their output representation will be ignored anyway (downstream).

I tried to think of ways in which it could matter, but I couldn't find any, other than possibly a performance difference. Wdyt?

[mattdangerw]

I tried to think of ways in which it could matter, but I couldn't find any, other than possibly a performance difference. Wdyt?

Yeah that makes sense to me. Thanks for the explanation. Let's go with what we have up.

Could you add the defensive casting to bool types for the masks, and a test?

[ageron]

Hi @mattdangerw ,
Thanks for your feedback, I'll add the defensive casting + test by ~Wednesday, sounds good?

[ageron]

Hi @mattdangerw, I just added defensive casting to bool (with tests).

[mattdangerw]

@ageron thank you so much! I will kick off another round of testing on this.

[mattdangerw]

(approving just to trigger our import flow for testing)

[mattdangerw]

Still looking at this. We have one production failure from an oddly shaped Keras implicit mask we are trying to figure out.

Hopefully no more action needed here, but will ping if anything comes up!