Function merge_padding_and_attention_mask does not return an output with the desired shape when both padding and attention masks are given

[abodinier]

Describe the bug
Function merge_padding_and_attention_mask does not return an output with the desired shape when both padding and attention masks are given.

In this example, the output should be of size (10, 20, 20), instead, it is (10, 10, 20, 20).

This is because a tf.newaxis is applied two times lines 70 https://github.com/keras-team/keras-nlp/blob/1eeec3b6d5de3796d1dbc28c4930c8189bac80c3/keras_nlp/layers/transformer_layer_utils.py#L70 and 77 https://github.com/keras-team/keras-nlp/blob/1eeec3b6d5de3796d1dbc28c4930c8189bac80c3/keras_nlp/layers/transformer_layer_utils.py#L77.

A fix to me would be to return return tf.minimum(mask, attention_mask) instead of tf.minimum(mask[:, tf.newaxis, :],attention_mask)

To Reproduce
Cf screenshot

Expected behavior
The function should return a tensor of shape (batch_size, S, T)

Additional context

Would you like to help us fix it?
Yes I would!

[blazingbhavneek]

Hey! I am available for this. I am new to the code base though, please refer me to some files/documentation so that I can understand the issue better

[Samarendra109]

I can take up this issue? Please let me know if this issue can be assigned to me.

[mattdangerw]

@abodinier thanks for finding this and filing!

The fix you are proposing seems reasonable, though it might be good to first look at a more end-to-end usage of this utility. E.g. if you provide both a padding_mask and attention_mask to a TransformerEncoder layer, will things break today?

Feel free to put a PR if our overall flow with a manual attention_mask is broken and the fix seems clear. One thing that might help the utility and overall usage is to have some nice error messages around rank. E.g. if padding_mask.rank != 2: raise ValueError(...) and if padding_mask.rank != 3: raise ValueError(...).

[abodinier]

@mattdangerw

Yes, today if we provide both padding_mask and attention_mask it breaks at line: https://github.com/keras-team/keras/blob/a63bef2ac504d651568c8800a9c1fcdcbd1cb41c/keras/layers/attention/multi_head_attention.py#L494 because because of the duplicated tf.newaxis, attention_scores and attention_mask have not the same shape (typical error message: InvalidArgumentError: Incompatible shapes: [2,8,3,3] vs. [2,2,3,3] [Op:AddV2])

The only case where it works is if by chance the batch size equals the num_heads, in this case, the second dimension will be the same.

Exs :
Only padding_mask: OK

Only attention_mask: OK

Both of them: KO

And the complete trace in the latter case:
`InvalidArgumentError Traceback (most recent call last)
/Users/alexandre/Desktop/issue keras_nlp/keras-nlp/keras_nlp/layers/test.ipynb Cell 4 in <cell line: 6>()
1 layer = TransformerEncoder(
2 intermediate_dim=10,
3 num_heads=4,
4 )
----> 6 layer.call(
7 inputs=inputs,
8 padding_mask=padding_mask,
9 attention_mask=attention_mask
10 )

File ~/Desktop/issue keras_nlp/keras-nlp/keras_nlp/layers/transformer_encoder.py:193, in TransformerEncoder.call(self, inputs, padding_mask, attention_mask)
191 if self.normalize_first:
192 x = self._self_attention_layernorm(x)
--> 193 x = self._self_attention_layer(
194 query=x,
195 value=x,
196 attention_mask=self_attention_mask,
197 )
198 x = self._self_attention_dropout(x)
199 x = x + residual

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/keras/engine/base_layer.py:1037, in Layer.call(self, *args, **kwargs)
1033 inputs = self._maybe_cast_inputs(inputs, input_list)
1035 with autocast_variable.enable_auto_cast_variables(
1036 self._compute_dtype_object):
-> 1037 outputs = call_fn(inputs, *args, **kwargs)
1039 if self._activity_regularizer:
1040 self._handle_activity_regularization(inputs, outputs)

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/keras/layers/multi_head_attention.py:503, in MultiHeadAttention.call(self, query, value, key, attention_mask, return_attention_scores, training)
500 # value = [B, S, N, H]
501 value = self._value_dense(value)
--> 503 attention_output, attention_scores = self._compute_attention(
504 query, key, value, attention_mask, training)
505 attention_output = self._output_dense(attention_output)
507 if return_attention_scores:

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/keras/layers/multi_head_attention.py:468, in MultiHeadAttention._compute_attention(self, query, key, value, attention_mask, training)
463 # Take the dot product between "query" and "key" to get the raw
464 # attention scores.
465 attention_scores = tf.einsum(self._dot_product_equation, key,
466 query)
--> 468 attention_scores = self._masked_softmax(attention_scores, attention_mask)
470 # This is actually dropping out entire tokens to attend to, which might
471 # seem a bit unusual, but is taken from the original Transformer paper.
472 attention_scores_dropout = self._dropout_layer(
473 attention_scores, training=training)

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/keras/layers/multi_head_attention.py:431, in MultiHeadAttention._masked_softmax(self, attention_scores, attention_mask)
428 for _ in range(len(attention_scores.shape) - len(attention_mask.shape)):
429 attention_mask = tf.expand_dims(
430 attention_mask, axis=mask_expansion_axis)
--> 431 return self._softmax(attention_scores, attention_mask)

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/keras/engine/base_layer.py:1037, in Layer.call(self, *args, **kwargs)
1033 inputs = self._maybe_cast_inputs(inputs, input_list)
1035 with autocast_variable.enable_auto_cast_variables(
1036 self._compute_dtype_object):
-> 1037 outputs = call_fn(inputs, *args, **kwargs)
1039 if self._activity_regularizer:
1040 self._handle_activity_regularization(inputs, outputs)

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/keras/layers/advanced_activations.py:341, in Softmax.call(self, inputs, mask)
336 adder = (1.0 - tf.cast(mask, inputs.dtype)) * (
337 _large_compatible_negative(inputs.dtype))
339 # Since we are adding it to the raw scores before the softmax, this is
340 # effectively the same as removing these entirely.
--> 341 inputs += adder
342 if isinstance(self.axis, (tuple, list)):
343 if len(self.axis) > 1:

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback..error_handler(*args, **kwargs)
151 except Exception as e:
152 filtered_tb = _process_traceback_frames(e.traceback)
--> 153 raise e.with_traceback(filtered_tb) from None
154 finally:
155 del filtered_tb

File /opt/homebrew/Caskroom/miniforge/base/lib/python3.9/site-packages/tensorflow/python/framework/ops.py:7107, in raise_from_not_ok_status(e, name)
7105 def raise_from_not_ok_status(e, name):
7106 e.message += (" name: " + name if name is not None else "")
-> 7107 raise core._status_to_exception(e) from None

InvalidArgumentError: Incompatible shapes: [2,4,3,3] vs. [2,2,3,3] [Op:AddV2]`

[abodinier]

The PR #790 (comment) has been created.

[dhantule]

@abodinier, Thanks for reporting this.
Can you confirm if the issue is resolved? If it's resolved could you please close the issue. Thanks !

[abodinier]

Hi @dhantule, I confirm it has been fixed and merged long ago, I close the issue.