Use random_attention_mask for TF tests

[ydshieh]

What does this PR do?

• Change TF's random_attention_mask to match its PT/Flax equivalence.

• Use random_attention_mask defined in test_modeling_tf_common.py to generate attention mask in TF tests.

  • so TF code has the same logic as in PT/Flax tests (regarding this attention mask part in tests)
  • avoid large difference between PT/TF outputs. (In particular, TFGPT2EncoderDecoderModelTest in here)
    • In the case of TFBERTEncoderDecoderModelTest or TFGPT2EncoderDecoderModelTest, it is caused by some sequence in a batch which gets all 0s as attention mask (generated by ids_tensor) - may happens on both encoder and decoder (especially after combining with the causal mask).

More context

Currently, most of TF tests still uses

input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

while in PT/Flax tests, they call

input_mask = random_attention_mask([self.batch_size, self.seq_length])

(defined in the comment test file).

In particular, random_attention_mask has

    # make sure that at least one token is attended to for each batch
    attn_mask[:, -1] = 1

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[ydshieh]

This is added to make TF CLIP pass.
(as TF's random_attention_mask is changed too in this PR)

[ydshieh]

This is changed to match PT/Flax's random_attention_mask.

[sgugger]

Thanks for making this more consistent with the rest of the library!

[gante]

Interesting. It moves the column of 1s for the start to the end and now becomes like a left-padded input. It could help with GPT-2, indeed

[ydshieh]

Interesting. It moves the column of 1s for the start to the end and now becomes like a left-padded input. It could help with GPT-2, indeed

(If you are interested to know a bit more the detail, @gante )

Actually, moving 1 to the end will cause problem (when a model uses causal mask.). This is why I needed to update the code in TFCLIPModelTest.

In general, current library has a bit issue when the final attention mask (after combining the causal mask if any) received by the attention layer has a sequence (in the batch) having all 0s as mask. One thing (but maybe not only) involved is the different values (-1e4, -1e9, -1e30, -inf) used.

Put 1 at the start will avoid this situation (when combining the causal mask).
(But I don't want to change the PT/Flax logic in this PR. This should be addressed in a separate PR after discussion.)

Regarding the tests like TFGPT2EncoderDecoderModelTest, this PR only helps partially (the encoder part). The decoder part needs extra logic for now (to address the above situation regarding the causal mask)

[Rocketknight1]

Writing as I go to make sure I follow:

• TF tests used to call ids_tensor with a vocabulary of 2 to generate a random attention mask
• They now call random_attention_mask, which also generates a tensor containing 0 and 1 only, but guarantees that at least one token will have a value of 1.
• This matches the behaviour in the rest of the library, guarantees we will never get a fully-masked input, and slightly increases the expected number of unmasked tokens in each input.

Seems like a great change for both test reliability and consistency with the rest of the library!

[ydshieh]

Hi, @Rocketknight1,

Yes, all the points are right -- except

• I am not sure about this statement and slightly increases the expected number of unmasked tokens in each input.: I would say not this case, but I might misunderstand the sentence.

• but guarantees that at least one token will have a value of 1:

  • Yes, but not guarantee the same thing for the final attention mask used by attention layers to compute the softmax - because the final mask might be the one after combining the causal mask (for decoder models).
  • Some more future PRs to improve these kinds of things.

[Rocketknight1]

That makes sense! And my comment about "increases the expected number of unmasked tokens" was just an irrelevant observation - the average number of unmasked tokens is very slightly larger since we guarantee that one of them will have value 1. Ignore me!