From fb5340967eba6f385c90fa30c8449b5a894e6bb3 Mon Sep 17 00:00:00 2001 From: ishtihoss Date: Tue, 7 Jul 2026 19:45:09 -0700 Subject: [PATCH 1/3] Fix BatchNorm running variance estimator --- python/mlx/nn/layers/normalization.py | 8 ++- python/tests/test_nn.py | 78 ++++++++++++++++++++++----- 2 files changed, 71 insertions(+), 15 deletions(-) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index f7d2d78c3b..67dd5d5c1a 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -371,8 +371,14 @@ def __call__(self, x: mx.array) -> mx.array: mean, var = self._calc_stats(x) if self.training and self.track_running_stats: mu = self.momentum + stats_size = 1 + for size in x.shape[:-1]: + stats_size *= size + running_var = ( + var * (stats_size / (stats_size - 1)) if stats_size > 1 else var + ) self.running_mean = (1 - mu) * self.running_mean + mu * mean - self.running_var = (1 - mu) * self.running_var + mu * var + self.running_var = (1 - mu) * self.running_var + mu * running_var elif self.track_running_stats: mean = self.running_mean var = self.running_var diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 8ad72a323c..7d977f9a09 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -672,7 +672,9 @@ def test_batch_norm(self): ], ) expected_mean = mx.array([0.008929, 0.005680, -0.016092, 0.027778]) - expected_var = mx.array([0.928435, 1.00455, 1.04117, 0.94258]) + expected_var = 0.9 * mx.ones_like(bn.running_var) + 0.1 * mx.var( + x, axis=0, ddof=1 + ) self.assertTrue(x.shape == y.shape) self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) self.assertTrue(mx.allclose(bn.running_mean, expected_mean, atol=1e-5)) @@ -681,14 +683,8 @@ def test_batch_norm(self): # test eval mode bn.eval() y = bn(x) - expected_y = mx.array( - [ - [-0.15984, 1.73159, -1.25456, 1.57891], - [-0.872193, -1.4281, -0.414439, -0.228678], - [0.602743, -0.30566, -0.554687, 0.139639], - [0.252199, 0.29066, -0.599572, -0.0512532], - [0.594096, -0.0334829, 2.11359, -0.151081], - ] + expected_y = ( + bn.weight * (x - expected_mean) * mx.rsqrt(expected_var + bn.eps) + bn.bias ) self.assertTrue(x.shape == y.shape) @@ -742,7 +738,9 @@ def test_batch_norm(self): expected_mean = mx.array( [[[0.00207845, -5.3259e-05, 0.04755, -0.0697296, 0.0236228]]] ) - expected_var = mx.array([[[0.968415, 1.05322, 0.96913, 0.932305, 0.967224]]]) + expected_var = 0.9 * mx.ones_like(bn.running_var) + 0.1 * mx.var( + x, axis=(0, 1), ddof=1 + ) self.assertTrue(mx.allclose(bn.running_mean, expected_mean, atol=1e-5)) self.assertTrue(mx.allclose(bn.running_var, expected_var, atol=1e-5)) @@ -797,9 +795,42 @@ def test_batch_norm_stats(self): normalized_data = batch_norm(data) means = mx.mean(data, axis=0) - variances = mx.var(data, axis=0) + batch_variances = mx.var(data, axis=0) + running_variances = mx.var(data, axis=0, ddof=1) + running_mean = (1 - momentum) * running_mean + momentum * means + running_var = (1 - momentum) * running_var + momentum * running_variances + expected_normalized_var = batch_variances / (batch_variances + batch_norm.eps) + self.assertTrue( + mx.allclose( + mx.var(normalized_data, axis=0), + expected_normalized_var, + atol=1e-5, + ) + ) + self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) + self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) + + batch_norm = nn.BatchNorm(num_features) + + batch_norm.train() + running_mean = batch_norm.running_mean + running_var = batch_norm.running_var + data = mx.random.normal((batch_size, h, num_features)) + + normalized_data = batch_norm(data) + means = mx.mean(data, axis=(0, 1)) + batch_variances = mx.var(data, axis=(0, 1)) + running_variances = mx.var(data, axis=(0, 1), ddof=1) running_mean = (1 - momentum) * running_mean + momentum * means - running_var = (1 - momentum) * running_var + momentum * variances + running_var = (1 - momentum) * running_var + momentum * running_variances + expected_normalized_var = batch_variances / (batch_variances + batch_norm.eps) + self.assertTrue( + mx.allclose( + mx.var(normalized_data, axis=(0, 1)), + expected_normalized_var, + atol=1e-5, + ) + ) self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) @@ -812,12 +843,31 @@ def test_batch_norm_stats(self): normalized_data = batch_norm(data) means = mx.mean(data, axis=(0, 1, 2)) - variances = mx.var(data, axis=(0, 1, 2)) + batch_variances = mx.var(data, axis=(0, 1, 2)) + running_variances = mx.var(data, axis=(0, 1, 2), ddof=1) running_mean = (1 - momentum) * running_mean + momentum * means - running_var = (1 - momentum) * running_var + momentum * variances + running_var = (1 - momentum) * running_var + momentum * running_variances + expected_normalized_var = batch_variances / (batch_variances + batch_norm.eps) + self.assertTrue( + mx.allclose( + mx.var(normalized_data, axis=(0, 1, 2)), + expected_normalized_var, + atol=1e-5, + ) + ) self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) + batch_norm = nn.BatchNorm(num_features) + + batch_norm.train() + running_var = batch_norm.running_var + data = mx.random.normal((1, num_features)) + + batch_norm(data) + running_var = (1 - momentum) * running_var + momentum * mx.var(data, axis=0) + self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) + self.assertEqual(batch_norm.running_mean.shape, running_mean.shape) self.assertEqual(batch_norm.running_var.shape, running_var.shape) From 9a4c1be1694c0822fd3cb09d7f18202ed3885f20 Mon Sep 17 00:00:00 2001 From: ishtihoss Date: Thu, 9 Jul 2026 09:02:53 -0700 Subject: [PATCH 2/3] Address BatchNorm review feedback --- python/mlx/nn/layers/normalization.py | 6 +- python/tests/test_nn.py | 162 ++++++++++++++------------ 2 files changed, 94 insertions(+), 74 deletions(-) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index 67dd5d5c1a..a079eddcda 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -374,8 +374,12 @@ def __call__(self, x: mx.array) -> mx.array: stats_size = 1 for size in x.shape[:-1]: stats_size *= size + # Normalize with biased batch variance, but store the unbiased + # estimate for inference. Keep stats_size == 1 finite. running_var = ( - var * (stats_size / (stats_size - 1)) if stats_size > 1 else var + mx.var(x, axis=tuple(range(0, x.ndim - 1)), ddof=1) + if stats_size > 1 + else var ) self.running_mean = (1 - mu) * self.running_mean + mu * mean self.running_var = (1 - mu) * self.running_var + mu * running_var diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 7d977f9a09..89680968cf 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -10,6 +10,13 @@ import numpy as np from mlx.utils import tree_flatten, tree_map, tree_reduce +try: + import torch + + has_torch = True +except ImportError: + has_torch = False + class TestBase(mlx_tests.MLXTestCase): def test_module_utilities(self): @@ -672,9 +679,7 @@ def test_batch_norm(self): ], ) expected_mean = mx.array([0.008929, 0.005680, -0.016092, 0.027778]) - expected_var = 0.9 * mx.ones_like(bn.running_var) + 0.1 * mx.var( - x, axis=0, ddof=1 - ) + expected_var = mx.array([0.935544, 1.030691, 1.076463, 0.953224]) self.assertTrue(x.shape == y.shape) self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) self.assertTrue(mx.allclose(bn.running_mean, expected_mean, atol=1e-5)) @@ -683,8 +688,14 @@ def test_batch_norm(self): # test eval mode bn.eval() y = bn(x) - expected_y = ( - bn.weight * (x - expected_mean) * mx.rsqrt(expected_var + bn.eps) + bn.bias + expected_y = mx.array( + [ + [-0.159232, 1.70949, -1.23382, 1.57007], + [-0.868873, -1.40987, -0.407588, -0.227397], + [0.600449, -0.301759, -0.545518, 0.138857], + [0.251239, 0.286951, -0.589661, -0.0509662], + [0.591834, -0.0330556, 2.07865, -0.150235], + ] ) self.assertTrue(x.shape == y.shape) @@ -736,11 +747,9 @@ def test_batch_norm(self): ) self.assertTrue(mx.allclose(y, expected_y, atol=1e-5)) expected_mean = mx.array( - [[[0.00207845, -5.3259e-05, 0.04755, -0.0697296, 0.0236228]]] - ) - expected_var = 0.9 * mx.ones_like(bn.running_var) + 0.1 * mx.var( - x, axis=(0, 1), ddof=1 + [0.00207845, -5.3259e-05, 0.04755, -0.0697296, 0.0236228] ) + expected_var = mx.array([0.978188, 1.07511, 0.979006, 0.93692, 0.976827]) self.assertTrue(mx.allclose(bn.running_mean, expected_mean, atol=1e-5)) self.assertTrue(mx.allclose(bn.running_var, expected_var, atol=1e-5)) @@ -778,6 +787,70 @@ def test_batch_norm(self): self.assertTrue(mx.allclose(y.mean(axis=(0, 1, 2)), mx.zeros((6,)), atol=1e-5)) self.assertTrue(mx.allclose(y.var(axis=(0, 1, 2)), mx.ones((6,)), atol=1e-2)) + @unittest.skipIf(not has_torch, "requires Torch") + def test_batch_norm_matches_torch(self): + rng = np.random.default_rng(0) + momentum = 0.1 + eps = 1e-5 + + def check_batch_norm(shape, torch_module, to_torch=None, from_torch=None): + features = shape[-1] + x_np = rng.normal(size=shape).astype(np.float32) + weight_np = rng.normal(size=(features,)).astype(np.float32) + bias_np = rng.normal(size=(features,)).astype(np.float32) + + mlx_bn = nn.BatchNorm(features, eps=eps, momentum=momentum) + mlx_bn.weight = mx.array(weight_np) + mlx_bn.bias = mx.array(bias_np) + mlx_y = mlx_bn(mx.array(x_np)) + mx.eval(mlx_y, mlx_bn.running_mean, mlx_bn.running_var) + + torch_bn = torch_module(features, eps=eps, momentum=momentum) + with torch.no_grad(): + torch_bn.weight.copy_(torch.from_numpy(weight_np)) + torch_bn.bias.copy_(torch.from_numpy(bias_np)) + x_torch_np = x_np.transpose(to_torch) if to_torch else x_np + torch_y = torch_bn(torch.from_numpy(x_torch_np)).detach().numpy() + if from_torch: + torch_y = torch_y.transpose(from_torch) + + np.testing.assert_allclose(np.array(mlx_y), torch_y, rtol=1e-4, atol=1e-4) + np.testing.assert_allclose( + np.array(mlx_bn.running_mean), + torch_bn.running_mean.detach().numpy(), + rtol=1e-5, + atol=1e-5, + ) + np.testing.assert_allclose( + np.array(mlx_bn.running_var), + torch_bn.running_var.detach().numpy(), + rtol=1e-4, + atol=1e-4, + ) + + mlx_bn.eval() + torch_bn.eval() + mlx_y = mlx_bn(mx.array(x_np)) + mx.eval(mlx_y) + torch_y = torch_bn(torch.from_numpy(x_torch_np)).detach().numpy() + if from_torch: + torch_y = torch_y.transpose(from_torch) + np.testing.assert_allclose(np.array(mlx_y), torch_y, rtol=1e-4, atol=1e-4) + + check_batch_norm((5, 4), torch.nn.BatchNorm1d) + check_batch_norm( + (2, 4, 5), + torch.nn.BatchNorm1d, + to_torch=(0, 2, 1), + from_torch=(0, 2, 1), + ) + check_batch_norm( + (2, 3, 3, 6), + torch.nn.BatchNorm2d, + to_torch=(0, 3, 1, 2), + from_torch=(0, 2, 3, 1), + ) + def test_batch_norm_stats(self): batch_size = 2 num_features = 4 @@ -786,77 +859,22 @@ def test_batch_norm_stats(self): momentum = 0.1 batch_norm = nn.BatchNorm(num_features) - batch_norm.train() - running_mean = batch_norm.running_mean - running_var = batch_norm.running_var - - data = mx.random.normal((batch_size, num_features)) - - normalized_data = batch_norm(data) - means = mx.mean(data, axis=0) - batch_variances = mx.var(data, axis=0) - running_variances = mx.var(data, axis=0, ddof=1) - running_mean = (1 - momentum) * running_mean + momentum * means - running_var = (1 - momentum) * running_var + momentum * running_variances - expected_normalized_var = batch_variances / (batch_variances + batch_norm.eps) - self.assertTrue( - mx.allclose( - mx.var(normalized_data, axis=0), - expected_normalized_var, - atol=1e-5, - ) - ) - self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) - self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) - - batch_norm = nn.BatchNorm(num_features) - - batch_norm.train() - running_mean = batch_norm.running_mean - running_var = batch_norm.running_var - data = mx.random.normal((batch_size, h, num_features)) + data = mx.random.normal((batch_size, h, w, num_features)) normalized_data = batch_norm(data) - means = mx.mean(data, axis=(0, 1)) - batch_variances = mx.var(data, axis=(0, 1)) - running_variances = mx.var(data, axis=(0, 1), ddof=1) - running_mean = (1 - momentum) * running_mean + momentum * means - running_var = (1 - momentum) * running_var + momentum * running_variances - expected_normalized_var = batch_variances / (batch_variances + batch_norm.eps) self.assertTrue( mx.allclose( - mx.var(normalized_data, axis=(0, 1)), - expected_normalized_var, - atol=1e-5, + mx.mean(normalized_data, axis=(0, 1, 2)), mx.zeros((4,)), atol=1e-5 ) ) - self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) - self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) - - batch_norm = nn.BatchNorm(num_features) - - batch_norm.train() - running_mean = batch_norm.running_mean - running_var = batch_norm.running_var - data = mx.random.normal((batch_size, h, w, num_features)) - - normalized_data = batch_norm(data) - means = mx.mean(data, axis=(0, 1, 2)) - batch_variances = mx.var(data, axis=(0, 1, 2)) - running_variances = mx.var(data, axis=(0, 1, 2), ddof=1) - running_mean = (1 - momentum) * running_mean + momentum * means - running_var = (1 - momentum) * running_var + momentum * running_variances - expected_normalized_var = batch_variances / (batch_variances + batch_norm.eps) self.assertTrue( mx.allclose( - mx.var(normalized_data, axis=(0, 1, 2)), - expected_normalized_var, - atol=1e-5, + mx.var(normalized_data, axis=(0, 1, 2)), mx.ones((4,)), atol=1e-2 ) ) - self.assertTrue(mx.allclose(batch_norm.running_mean, running_mean, atol=1e-5)) - self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) + self.assertEqual(batch_norm.running_mean.shape, (num_features,)) + self.assertEqual(batch_norm.running_var.shape, (num_features,)) batch_norm = nn.BatchNorm(num_features) @@ -864,13 +882,11 @@ def test_batch_norm_stats(self): running_var = batch_norm.running_var data = mx.random.normal((1, num_features)) - batch_norm(data) + normalized_data = batch_norm(data) + self.assertTrue(np.all(np.isfinite(normalized_data))) running_var = (1 - momentum) * running_var + momentum * mx.var(data, axis=0) self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) - self.assertEqual(batch_norm.running_mean.shape, running_mean.shape) - self.assertEqual(batch_norm.running_var.shape, running_var.shape) - def test_conv1d(self): N = 5 L = 12 From 41d888abc0b348ee2a6d78fd3c0dc65aa93f0bd8 Mon Sep 17 00:00:00 2001 From: ishtihoss Date: Thu, 9 Jul 2026 18:01:00 -0700 Subject: [PATCH 3/3] Address BatchNorm follow-up review --- python/mlx/nn/layers/normalization.py | 25 +++++++++--------- python/tests/test_nn.py | 37 +++++++++++++-------------- 2 files changed, 31 insertions(+), 31 deletions(-) diff --git a/python/mlx/nn/layers/normalization.py b/python/mlx/nn/layers/normalization.py index a079eddcda..339dfaf1e0 100644 --- a/python/mlx/nn/layers/normalization.py +++ b/python/mlx/nn/layers/normalization.py @@ -333,13 +333,14 @@ def _extra_repr(self): f"track_running_stats={self.track_running_stats}" ) - def _calc_stats(self, x: mx.array) -> Tuple[mx.array, mx.array]: + def _calc_stats(self, x: mx.array, ddof: int = 0) -> Tuple[mx.array, mx.array]: """ Calculate the mean and variance of the input tensor across the batch and spatial dimensions. Args: x (array): Input tensor. + ddof (int): Delta degrees of freedom for variance. Returns: tuple: Tuple containing mean and variance. @@ -347,7 +348,7 @@ def _calc_stats(self, x: mx.array) -> Tuple[mx.array, mx.array]: reduction_axes = tuple(range(0, x.ndim - 1)) mean = mx.mean(x, axis=reduction_axes) - var = mx.var(x, axis=reduction_axes) + var = mx.var(x, axis=reduction_axes, ddof=ddof) return mean, var @@ -366,21 +367,21 @@ def __call__(self, x: mx.array) -> mx.array: f"Expected input tensor to have 2, 3 or 4 dimensions, but got {x.ndim}" ) + if self.training: + stats_size = 1 + for size in x.shape[:-1]: + stats_size *= size + if stats_size == 1: + raise ValueError( + "BatchNorm training requires more than one value per channel." + ) + # Calculate the mean and variance used to normalize the input x. If we # are in training mode update the running stats if needed. mean, var = self._calc_stats(x) if self.training and self.track_running_stats: mu = self.momentum - stats_size = 1 - for size in x.shape[:-1]: - stats_size *= size - # Normalize with biased batch variance, but store the unbiased - # estimate for inference. Keep stats_size == 1 finite. - running_var = ( - mx.var(x, axis=tuple(range(0, x.ndim - 1)), ddof=1) - if stats_size > 1 - else var - ) + _, running_var = self._calc_stats(x, ddof=1) self.running_mean = (1 - mu) * self.running_mean + mu * mean self.running_var = (1 - mu) * self.running_var + mu * running_var elif self.track_running_stats: diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 89680968cf..0df667a64d 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -814,18 +814,22 @@ def check_batch_norm(shape, torch_module, to_torch=None, from_torch=None): if from_torch: torch_y = torch_y.transpose(from_torch) - np.testing.assert_allclose(np.array(mlx_y), torch_y, rtol=1e-4, atol=1e-4) - np.testing.assert_allclose( - np.array(mlx_bn.running_mean), - torch_bn.running_mean.detach().numpy(), - rtol=1e-5, - atol=1e-5, + self.assertTrue(mx.allclose(mlx_y, mx.array(torch_y), rtol=1e-4, atol=1e-4)) + self.assertTrue( + mx.allclose( + mlx_bn.running_mean, + mx.array(torch_bn.running_mean.detach().numpy()), + rtol=1e-5, + atol=1e-5, + ) ) - np.testing.assert_allclose( - np.array(mlx_bn.running_var), - torch_bn.running_var.detach().numpy(), - rtol=1e-4, - atol=1e-4, + self.assertTrue( + mx.allclose( + mlx_bn.running_var, + mx.array(torch_bn.running_var.detach().numpy()), + rtol=1e-4, + atol=1e-4, + ) ) mlx_bn.eval() @@ -835,7 +839,7 @@ def check_batch_norm(shape, torch_module, to_torch=None, from_torch=None): torch_y = torch_bn(torch.from_numpy(x_torch_np)).detach().numpy() if from_torch: torch_y = torch_y.transpose(from_torch) - np.testing.assert_allclose(np.array(mlx_y), torch_y, rtol=1e-4, atol=1e-4) + self.assertTrue(mx.allclose(mlx_y, mx.array(torch_y), rtol=1e-4, atol=1e-4)) check_batch_norm((5, 4), torch.nn.BatchNorm1d) check_batch_norm( @@ -856,7 +860,6 @@ def test_batch_norm_stats(self): num_features = 4 h = 3 w = 3 - momentum = 0.1 batch_norm = nn.BatchNorm(num_features) batch_norm.train() @@ -877,15 +880,11 @@ def test_batch_norm_stats(self): self.assertEqual(batch_norm.running_var.shape, (num_features,)) batch_norm = nn.BatchNorm(num_features) - batch_norm.train() - running_var = batch_norm.running_var data = mx.random.normal((1, num_features)) - normalized_data = batch_norm(data) - self.assertTrue(np.all(np.isfinite(normalized_data))) - running_var = (1 - momentum) * running_var + momentum * mx.var(data, axis=0) - self.assertTrue(mx.allclose(batch_norm.running_var, running_var, atol=1e-5)) + with self.assertRaises(ValueError): + batch_norm(data) def test_conv1d(self): N = 5