Stabilize reduced-precision InstanceNorm (#4230)

Author: ternaus

python/mlx/nn/layers/normalization.py

@@ -64,12 +64,19 @@ def __call__(self, x: mx.array) -> mx.array:
                 f"InstanceNorm expects inputs with at least 3 dimensions"
                 f" (N, ..., C) but the input has {x.ndim} dimensions."
             )
-        reduction_axes = tuple(range(1, x.ndim - 1))
-        # Compute stats
-        mean = mx.mean(x, axis=reduction_axes, keepdims=True)
-        var = mx.var(x, axis=reduction_axes, keepdims=True)
-        # Normalize
-        x = (x - mean) * mx.rsqrt(var + self.eps)
+        batch_size, features = x.shape[0], x.shape[-1]
+        spatial_shape = x.shape[1:-1]
+        channels_first = mx.transpose(x, (0, x.ndim - 1, *range(1, x.ndim - 1)))
+        x = mx.fast.layer_norm(
+            channels_first.reshape(batch_size, features, -1),
+            None,
+            None,
+            self.eps,
+        )
+        x = mx.transpose(
+            x.reshape(batch_size, features, *spatial_shape),
+            (0, *range(2, len(spatial_shape) + 2), 1),
+        )
         # Scale and shift if necessary
         return (self.weight * x + self.bias) if "weight" in self else x
 

python/tests/test_nn.py

@@ -693,6 +693,21 @@ def test_instance_norm(self):
         ]
         self.assertTrue(x.shape == y.shape)
         self.assertTrue(np.allclose(y, expected_y, atol=1e-5))
+        # Reduced-precision statistics must not overflow for finite feature maps.
+        checkerboard = np.indices((4, 4, 4)).sum(axis=0) % 2
+        x = mx.array(
+            np.stack(
+                [
+                    np.where(checkerboard, -512, 512),
+                    np.where(checkerboard, -256, 256),
+                ],
+                axis=-1,
+            ).astype(np.float16)
+        )[None]
+        y = nn.InstanceNorm(dims=2)(x)
+        self.assertEqual(y.dtype, mx.float16)
+        self.assertTrue(mx.allclose(y.min(), mx.array(-1.0, dtype=mx.float16)))
+        self.assertTrue(mx.allclose(y.max(), mx.array(1.0, dtype=mx.float16)))
         # Test repr
         self.assertTrue(str(inorm) == "InstanceNorm(3, eps=1e-05, affine=False)")
         # Raise for inputs without spatial dimensions