Why tf.contrib.layers.instance_norm layer contain StopGradient operation?

user_1179925 0
Viewed 289

Why tf.contrib.layers.instance_norm layer contain StopGradient operation? i.e. why it's needed?

enter image description here

Seems there is StopGradient even in simpler layer tf.nn.moments (that can be building block of tf.contrib.layers.instance_norm).

x_m, x_v = tf.nn.moments(x, [1, 2], keep_dims=True)

enter image description here

Also I find a note on StopGradient in tf.nn.moments source code:

# The dynamic range of fp16 is too limited to support the collection of
# sufficient statistics. As a workaround we simply perform the operations
# on 32-bit floats before converting the mean and variance back to fp16
y = math_ops.cast(x, dtypes.float32) if x.dtype == dtypes.float16 else x
# Compute true mean while keeping the dims for proper broadcasting.
mean = math_ops.reduce_mean(y, axes, keepdims=True, name="mean")
# sample variance, not unbiased variance
# Note: stop_gradient does not change the gradient that gets
#       backpropagated to the mean from the variance calculation,
#       because that gradient is zero
variance = math_ops.reduce_mean(
    math_ops.squared_difference(y, array_ops.stop_gradient(mean)),
    axes,
    keepdims=True,
    name="variance")

So it's sort of optimisation because gradient is always zero?

batch-normalization tensorflow deep-learning
1 Answers

Attempt of an answer.

This design tells us that minimizing second moment we would not want to propagate gradients through the first moment. Does it make sense? If we try to minimize E[x^2]-E[x]^2 we would minimize E[x^2] while simultaneously maximizing E[x]^2. First term would decrease absolute values of each element (drag them to the center). Second term would increase all values by gradient which would do nothing to minimize variance but might negatively affect other gradient paths.

So, we don't propagate gradient of second moment through the first moment because this gradient would not effect second moment whatsoever, at least when using plain SGD.

user_8678890
Related