Error in masked BatchNormalization with > 3 dimensions

[drasmuss]

import keras

x = keras.ops.ones((1, 2, 3, 4))
mask = keras.ops.ones((1, 2), dtype="bool")
y = keras.layers.BatchNormalization()(x, mask=mask, training=True)

gives

tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling BatchNormalization.call().

{{function_node __wrapped__Mul_device_/job:localhost/replica:0/task:0/device:CPU:0}} Incompatible shapes: [1,2,1] vs. [1,2,3,4] [Op:Mul] name: 

Arguments received by BatchNormalization.call():
  • inputs=tf.Tensor(shape=(1, 2, 3, 4), dtype=float32)
  • training=True
  • mask=tf.Tensor(shape=(1, 2), dtype=bool)

[Grvzard]

According to the example above, I think what we should do in Keras is automatically expand the mask from shape (1, 2) to (1, 2, 3) right?

[fchollet]

You should generally make sure that the mask only has one less dimension than the input, e.g. in this case mask = keras.ops.ones((1, 2, 1), dtype="bool")

[pandyah5]

This raises an exception and explains the mismatch in dimensions. Is this a bug to be fixed on the Keras side? I believe it's more of an expected exception right?

[Grvzard]

Is this a bug to be fixed on the Keras side?

As mentioned above. What we can do is make the exception message clearer.

[pandyah5]

Gotcha! I can try modifying it tonight

[drasmuss]

I'm not sure if this is under the control of the Keras team, but the docs here https://www.tensorflow.org/guide/keras/understanding_masking_and_padding#mask-generating_layers_embedding_and_masking, explicitly state that masks will be shape (batch_size, sequence_length). It would be nice to update that to clarify that they should actually be input.ndims - 1.

[drasmuss]

While implementing @fchollet suggested fix here #19818 (comment) I ran into a different issue. The mean/variance calculations are incorrect when that mask is specified:

import keras

x = keras.ops.ones((1, 2, 3, 4))
layer = keras.layers.BatchNormalization()
layer.build(x.shape)
print(layer._moments(x, None))
print(layer._moments(x, keras.ops.ones((1, 2, 1), dtype="bool")))
print(layer._moments(x, keras.ops.ones((1, 2, 3), dtype="bool")))

gives output

(<tf.Tensor: shape=(4,), dtype=float32, numpy=array([1., 1., 1., 1.], dtype=float32)>, <tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>)
(<tf.Tensor: shape=(4,), dtype=float32, numpy=array([3., 3., 3., 3.], dtype=float32)>, <tf.Tensor: shape=(4,), dtype=float32, numpy=array([12., 12., 12., 12.], dtype=float32)>)
(<tf.Tensor: shape=(4,), dtype=float32, numpy=array([1., 1., 1., 1.], dtype=float32)>, <tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>)

We can see that a mask with shape (1, 2, 1) gives incorrect mean/variance, but (1, 2, 3) gives correct output. It wasn't a problem for me to add trailing size-1 dimensions to the mask, but if I have to add a fully-tiled dimension to all the non-sequence axes of my sequence mask that will dramatically increase the memory usage of the model, as those dimensions are quite large.

Is the intended behaviour for the sequence masks that they need to be fully specified for all dimensions except the last (like (1, 2, 3))? Or should (1, 2, 1) be working, and this is just a bug in the mean/variance calculations in BatchNormalization?

[pandyah5]

@drasmuss That's an interesting observation, given its a tensor of ones, we should not expect a mean greater than 1. I believe the error lies here:

keras/keras/src/layers/normalization/batch_normalization.py

Line 299 in efaaf85

mask_weights_broadcasted = ops.expand_dims(

When I tried to debug it, I narrowed down the issue to this. I am just trying to make sense of how the computations are being carried out here:

[pandyah5]

The reason we are getting a mean of three is as follows:

• If the tensors in Keras behave the same way as NumPy then the mask is being implicitly broadcasted from (1, 2, 1, 1) to (1, 2, 3, 4) in the line below. What is even more confusing is that the mask_weights_broadcasted variable is of shape (1, 1, 2, 1). It is not exactly broadcasted but rather extended to meet the number of dimensions. The implicit broadcasting happens here:

  keras/keras/src/layers/normalization/batch_normalization.py

  Line 303 in efaaf85

  weighted_inputs = mask_weights_broadcasted * inputs

• As seen from the screenshot above this does not cause any masking and the input remains the same and we get [6, 6, 6, 6].

• However, the sum of weights is 2, because in the tensor of shape (1, 1, 2, 1) we just have two ones hence the 2.

This gives us a mean of 3. @fchollet I'll need some advice over here, but the issue seems like the mask is not doing its job to trim down the input. Let me know if I am on the correct path :)

[drasmuss]

#19829 doesn't fix the bug in the mean/variance calculations. Would you rather I open a different issue for that bug?

[pandyah5]

HI @drasmuss I would recommend opening a new issue for this to make it more tractable. If possible you could also include the triage information I provided before to give a lead to the contributors.