Why does tfa.layers.GroupNormalization(groups=1) produce different output than LayerNormalization?

Viewed 304

From the group normalization documentation in tensorflow addons, it states that the group norm layer should become layer normalization if the number of groups is set to one.

However when I try this by calling the layers one a test tensor the results differ. It appears that the group norm layer computes the mean and variance across the time as well as the channel axis, whereas the layer norm computes it for each channel's vector independently.

Is this a bug or am I missing something? The current behavior of layer norm is actually desirable for what I am doing.

Here is the documentation for GroupNormalization:

In [5]: x = tf.constant([[[1, 2], [3, 40]], [[1 , -1], [2, 200]]], dtype = tf.float32)
In [6]: tf.keras.layers.LayerNormalization()(x)                                                                                                               
Out[6]: 
<tf.Tensor: shape=(2, 2, 2), dtype=float32, numpy=
array([[[-0.99800587,  0.99800587],
        [-0.99999857,  0.99999857]],

       [[ 0.9995002 , -0.9995002 ],
        [-1.        ,  1.        ]]], dtype=float32)>

In [7]: tfa.layers.GroupNormalization(groups = 1)(x)                                                                                                          
Out[7]: 
<tf.Tensor: shape=(2, 2, 2), dtype=float32, numpy=
array([[[-0.6375344 , -0.57681686],
        [-0.5160993 ,  1.7304504 ]],

       [[-0.5734435 , -0.5966129 ],
        [-0.5618587 ,  1.7319152 ]]], dtype=float32)>
1 Answers

According to the doc of tf.keras.layers.LayerNormalization in TF 2.4.1, source:

Note that other implementations of layer normalization may choose to define gamma and beta over a separate set of axes from the axes being normalized across. For example, Group Normalization (Wu et al. 2018) with a group size of 1 corresponds to a Layer Normalization that normalizes across height, width, and channel and has gamma and beta span only the channel dimension. So, this Layer Normalization implementation will not match a Group Normalization layer with group size set to 1.

Related