From 346e4720cf375068ebe8b85972ba0e1e38ad97a0 Mon Sep 17 00:00:00 2001
From: andrewdipper <andrewdipper11235@gmail.com>
Date: Wed, 6 Mar 2024 16:00:08 -0800
Subject: [PATCH 1/6] Update BatchNorm to add mode for using batch statisics /
 increase stability when using running statistics.

---
 equinox/nn/_batch_norm.py | 62 +++++++++++++++++++++++++++++++--------
 tests/test_nn.py          | 39 +++++++++++++++++++-----
 tests/test_stateful.py    |  4 +--
 3 files changed, 83 insertions(+), 22 deletions(-)

diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index 45802d10..c778dfdd 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -1,10 +1,11 @@
+import warnings
 from collections.abc import Hashable, Sequence
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 import jax
 import jax.lax as lax
 import jax.numpy as jnp
-from jaxtyping import Array, Bool, Float, PRNGKeyArray
+from jaxtyping import Array, Float, Int, PRNGKeyArray
 
 from .._misc import default_floating_dtype
 from .._module import field
@@ -44,24 +45,29 @@ class BatchNorm(StatefulLayer, strict=True):
 
     weight: Optional[Float[Array, "input_size"]]
     bias: Optional[Float[Array, "input_size"]]
-    first_time_index: StateIndex[Bool[Array, ""]]
+    count_index: StateIndex[Int[Array, ""]]
     state_index: StateIndex[
         tuple[Float[Array, "input_size"], Float[Array, "input_size"]]
     ]
+    zero_frac_index: StateIndex[Float[Array, ""]]
     axis_name: Union[Hashable, Sequence[Hashable]]
     inference: bool
     input_size: int = field(static=True)
+    approach: Union[None, str] = field(static=True)
     eps: float = field(static=True)
     channelwise_affine: bool = field(static=True)
     momentum: float = field(static=True)
+    warmup_period: int = field(static=True)
 
     def __init__(
         self,
         input_size: int,
         axis_name: Union[Hashable, Sequence[Hashable]],
+        approach: Optional[Literal["batch", "ema"]] = None,
         eps: float = 1e-5,
         channelwise_affine: bool = True,
         momentum: float = 0.99,
+        warmup_period: int = 1000,
         inference: bool = False,
         dtype=None,
     ):
@@ -71,11 +77,17 @@ def __init__(
         - `axis_name`: The name of the batch axis to compute statistics over, as passed
             to `axis_name` in `jax.vmap` or `jax.pmap`. Can also be a sequence (e.g. a
             tuple or a list) of names, to compute statistics over multiple named axes.
+        - `approach`: The approach to use for the running statistics. If `approach=None`
+            a warning will be raised and approach will default to `"batch"`. During
+            training `"batch"` only uses batch statisics while`"ema"` uses the running
+            statistics.
         - `eps`: Value added to the denominator for numerical stability.
         - `channelwise_affine`: Whether the module has learnable channel-wise affine
             parameters.
         - `momentum`: The rate at which to update the running statistics. Should be a
             value between 0 and 1 exclusive.
+        - `warmup_period`: The period to warm up the running statistics. Only used when
+            `approach=\"ema\"`.
         - `inference`: If `False` then the batch means and variances will be calculated
             and used to update the running statistics. If `True` then the running
             statistics are directly used for normalisation. This may be toggled with
@@ -86,26 +98,37 @@ def __init__(
             64-bit mode.
         """
 
+        if approach is None:
+            warnings.warn('BatchNorm approach is None, defaults to approach="batch"')
+            approach = "batch"
+
+        valid_approaches = {"batch", "ema"}
+        if approach not in valid_approaches:
+            raise ValueError(f"approach must be one of {valid_approaches}")
+        self.approach = approach
+
         if channelwise_affine:
             self.weight = jnp.ones((input_size,))
             self.bias = jnp.zeros((input_size,))
         else:
             self.weight = None
             self.bias = None
-        self.first_time_index = StateIndex(jnp.array(True))
+        self.count_index = StateIndex(jnp.array(0, dtype=jnp.int32))
         if dtype is None:
             dtype = default_floating_dtype()
         init_buffers = (
-            jnp.empty((input_size,), dtype=dtype),
-            jnp.empty((input_size,), dtype=dtype),
+            jnp.zeros((input_size,), dtype=dtype),
+            jnp.zeros((input_size,), dtype=dtype),
         )
         self.state_index = StateIndex(init_buffers)
+        self.zero_frac_index = StateIndex(jnp.array(1.0, dtype=dtype))
         self.inference = inference
         self.axis_name = axis_name
         self.input_size = input_size
         self.eps = eps
         self.channelwise_affine = channelwise_affine
         self.momentum = momentum
+        self.warmup_period = max(1, warmup_period)
 
     @jax.named_scope("eqx.nn.BatchNorm")
     def __call__(
@@ -143,7 +166,10 @@ def __call__(
         if inference is None:
             inference = self.inference
         if inference:
+            zero_frac = state.get(self.zero_frac_index)
             running_mean, running_var = state.get(self.state_index)
+            norm_mean = running_mean / jnp.maximum(1.0 - zero_frac, self.eps)
+            norm_var = running_var / jnp.maximum(1.0 - zero_frac, self.eps)
         else:
 
             def _stats(y):
@@ -154,23 +180,35 @@ def _stats(y):
                 var = jnp.maximum(0.0, var)
                 return mean, var
 
-            first_time = state.get(self.first_time_index)
-            state = state.set(self.first_time_index, jnp.array(False))
+            momentum = self.momentum
+            zero_frac = state.get(self.zero_frac_index)
+            zero_frac *= momentum
+            state = state.set(self.zero_frac_index, zero_frac)
 
             batch_mean, batch_var = jax.vmap(_stats)(x)
             running_mean, running_var = state.get(self.state_index)
-            momentum = self.momentum
             running_mean = (1 - momentum) * batch_mean + momentum * running_mean
             running_var = (1 - momentum) * batch_var + momentum * running_var
-            running_mean = lax.select(first_time, batch_mean, running_mean)
-            running_var = lax.select(first_time, batch_var, running_var)
             state = state.set(self.state_index, (running_mean, running_var))
 
+            if self.approach == "ema":
+                warmup_count = state.get(self.count_index)
+                warmup_count = jnp.minimum(warmup_count + 1, self.warmup_period)
+                state = state.set(self.count_index, warmup_count)
+
+                warmup_frac = warmup_count / self.warmup_period
+                norm_mean = zero_frac * batch_mean + running_mean
+                norm_mean = (1.0 - warmup_frac) * batch_mean + warmup_frac * norm_mean
+                norm_var = zero_frac * batch_var + running_var
+                norm_var = (1.0 - warmup_frac) * batch_var + warmup_frac * norm_var
+            else:
+                norm_mean, norm_var = batch_mean, batch_var
+
         def _norm(y, m, v, w, b):
             out = (y - m) / jnp.sqrt(v + self.eps)
             if self.channelwise_affine:
                 out = out * w + b
             return out
 
-        out = jax.vmap(_norm)(x, running_mean, running_var, self.weight, self.bias)
+        out = jax.vmap(_norm)(x, norm_mean, norm_var, self.weight, self.bias)
         return out, state
diff --git a/tests/test_nn.py b/tests/test_nn.py
index babdd07d..6d754461 100644
--- a/tests/test_nn.py
+++ b/tests/test_nn.py
@@ -124,7 +124,7 @@ def test_sequential(getkey):
         [
             eqx.nn.Linear(2, 4, key=getkey()),
             eqx.nn.Linear(4, 1, key=getkey()),
-            eqx.nn.BatchNorm(1, axis_name="batch"),
+            eqx.nn.BatchNorm(1, axis_name="batch", approach="batch"),
             eqx.nn.Linear(1, 3, key=getkey()),
         ]
     )
@@ -158,7 +158,7 @@ def make():
         inner_seq = eqx.nn.Sequential(
             [
                 eqx.nn.Linear(2, 4, key=getkey()),
-                eqx.nn.BatchNorm(4, axis_name="batch")
+                eqx.nn.BatchNorm(4, axis_name="batch", approach="batch")
                 if inner_stateful
                 else eqx.nn.Identity(),
                 eqx.nn.Linear(4, 3, key=getkey()),
@@ -168,7 +168,7 @@ def make():
             [
                 eqx.nn.Linear(5, 2, key=getkey()),
                 inner_seq,
-                eqx.nn.BatchNorm(3, axis_name="batch")
+                eqx.nn.BatchNorm(3, axis_name="batch", approach="batch")
                 if outer_stateful
                 else eqx.nn.Identity(),
                 eqx.nn.Linear(3, 6, key=getkey()),
@@ -825,18 +825,35 @@ def test_batch_norm(getkey):
     x2 = jrandom.uniform(getkey(), (10, 5, 6))
     x3 = jrandom.uniform(getkey(), (10, 5, 7, 8))
 
-    # Test that it works with a single vmap'd axis_name
+    # Test that it warns with no approach - defaulting to batch
+    with pytest.warns(UserWarning):
+        bn = eqx.nn.BatchNorm(5, "batch")
+        assert bn.approach == "batch"
 
-    bn = eqx.nn.BatchNorm(5, "batch")
+    # Test initialization
+    bn_momentum = 0.99
+    bn = eqx.nn.BatchNorm(5, "batch", approach="ema", momentum=bn_momentum)
     state = eqx.nn.State(bn)
     vbn = jax.vmap(bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
+    running_mean, running_var = state.get(bn.state_index)
+    zero_frac = state.get(bn.zero_frac_index)
+    warmup_count = state.get(bn.count_index)
+    assert jnp.array_equal(running_mean, jnp.zeros(running_mean.shape))
+    assert jnp.array_equal(running_var, jnp.zeros(running_var.shape))
+    assert jnp.array_equal(zero_frac, jnp.array(1.0))
+    assert jnp.array_equal(warmup_count, jnp.array(0))
 
-    for x in (x1, x2, x3):
+    # Test that it works with a single vmap'd axis_name
+    for i, x in enumerate([x1, x2, x3]):
         out, state = vbn(x, state)
         assert out.shape == x.shape
         running_mean, running_var = state.get(bn.state_index)
+        zero_frac = state.get(bn.zero_frac_index)
+        warmup_count = state.get(bn.count_index)
         assert running_mean.shape == (5,)
         assert running_var.shape == (5,)
+        assert jnp.array_equal(warmup_count, jnp.array(i + 1))
+        assert jnp.allclose(zero_frac, jnp.array(bn_momentum ** (i + 1)))
 
     # Test that it fails without any vmap'd axis_name
 
@@ -861,7 +878,7 @@ def test_batch_norm(getkey):
 
     # Test that it handles multiple axis_names
 
-    vvbn = eqx.nn.BatchNorm(6, ("batch1", "batch2"))
+    vvbn = eqx.nn.BatchNorm(6, ("batch1", "batch2"), approach="ema")
     vvstate = eqx.nn.State(vvbn)
     for axis_name in ("batch1", "batch2"):
         vvbn = jax.vmap(
@@ -876,7 +893,7 @@ def test_batch_norm(getkey):
     # Test that it normalises
 
     x1alt = jrandom.normal(jrandom.PRNGKey(5678), (10, 5))  # avoid flakey test
-    bn = eqx.nn.BatchNorm(5, "batch", channelwise_affine=False)
+    bn = eqx.nn.BatchNorm(5, "batch", channelwise_affine=False, approach="ema")
     state = eqx.nn.State(bn)
     vbn = jax.vmap(bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
     out, state = vbn(x1alt, state)
@@ -890,6 +907,8 @@ def test_batch_norm(getkey):
     running_mean, running_var = state.get(bn.state_index)
     out, state = vbn(3 * x1 + 10, state)
     running_mean2, running_var2 = state.get(bn.state_index)
+    zero_frac2 = state.get(bn.zero_frac_index)
+    warmup_count2 = state.get(bn.count_index)
     assert not jnp.allclose(running_mean, running_mean2)
     assert not jnp.allclose(running_var, running_var2)
 
@@ -899,8 +918,12 @@ def test_batch_norm(getkey):
     vibn = jax.vmap(ibn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
     out, state = vibn(4 * x1 + 20, state)
     running_mean3, running_var3 = state.get(bn.state_index)
+    zero_frac3 = state.get(bn.zero_frac_index)
+    warmup_count3 = state.get(bn.count_index)
     assert jnp.array_equal(running_mean2, running_mean3)
     assert jnp.array_equal(running_var2, running_var3)
+    assert jnp.array_equal(zero_frac2, zero_frac3)
+    assert jnp.array_equal(warmup_count2, warmup_count3)
 
     # Test that we can differentiate through it
 
diff --git a/tests/test_stateful.py b/tests/test_stateful.py
index d7bc632a..de75319b 100644
--- a/tests/test_stateful.py
+++ b/tests/test_stateful.py
@@ -7,7 +7,7 @@
 
 
 def test_delete_init_state():
-    model = eqx.nn.BatchNorm(3, "batch")
+    model = eqx.nn.BatchNorm(3, "batch", approach="batch")
     eqx.nn.State(model)
     model2 = eqx.nn.delete_init_state(model)
 
@@ -17,7 +17,7 @@ def test_delete_init_state():
 
     leaves = [x for x in jtu.tree_leaves(model) if eqx.is_array(x)]
     leaves2 = [x for x in jtu.tree_leaves(model2) if eqx.is_array(x)]
-    assert len(leaves) == len(leaves2) + 3
+    assert len(leaves) == len(leaves2) + 4
 
 
 def test_double_state():

From cc021f63379ba14ead7d9bb1ff691703448a2dff Mon Sep 17 00:00:00 2001
From: andrewdipper <andrewdipper11235@gmail.com>
Date: Wed, 6 Mar 2024 18:16:16 -0800
Subject: [PATCH 2/6] fix batchnorm type definiton

---
 equinox/nn/_batch_norm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index c778dfdd..8c9adabf 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -53,7 +53,7 @@ class BatchNorm(StatefulLayer, strict=True):
     axis_name: Union[Hashable, Sequence[Hashable]]
     inference: bool
     input_size: int = field(static=True)
-    approach: Union[None, str] = field(static=True)
+    approach: Literal["batch", "ema"] = field(static=True)
     eps: float = field(static=True)
     channelwise_affine: bool = field(static=True)
     momentum: float = field(static=True)

From 6bfa77fe8e7389f18e15cc03b0f76673caa0ca87 Mon Sep 17 00:00:00 2001
From: andrewdipper <andrewdipper11235@gmail.com>
Date: Thu, 14 Mar 2024 12:44:03 -0700
Subject: [PATCH 3/6] use unbiased variance + minor fixes

---
 equinox/nn/_batch_norm.py | 78 +++++++++++++++++++++++++++++++++++----
 tests/test_nn.py          | 32 ++++++++++++++--
 2 files changed, 99 insertions(+), 11 deletions(-)

diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index 8c9adabf..67795ffb 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -41,6 +41,61 @@ class BatchNorm(StatefulLayer, strict=True):
     statistics updated. During inference then just the running statistics are used.
     Whether the model is in training or inference mode should be toggled using
     [`equinox.nn.inference_mode`][].
+
+    With `approach = "batch"` during training the batch mean and variance are used
+    for normalization.  For inference the exponential running mean and ubiased
+    variance are used for normalization in accordance with the cited paper below:
+
+    $\text{TrainStats}_t = \text{BatchStats}_t$
+
+    $\text{InferenceStats}_t = \frac{\left(1.0 - m\right)\sum_{i=0}^{t}m^{t-i}
+    \text{BatchStats}_i}{\text{max} \left(1.0 - m^{t+1}, \varepsilon \right)}$
+
+    With `approach = "ema"` exponential running means and variances are kept.  During
+    training the batch statistics are used to fill in the running statistics until
+    they are populated.  In addition a linear iterpolation is used between the batch
+    and running statistics over the `warmup_period`.  During inference the running
+    statistics are used for normalization:
+
+
+
+    $\text{WarmupFrac}_t = \text{min} \left(1.0, \frac{t}{\text{WarmupPeriod}} \right)$
+
+    $\text{TrainStats}_t = (1.0 - \text{WarmupFrac}_t) * BatchStats_t +
+    \text{WarmupFrac}_t * \left(1.0 - m\right)\sum_{i=0}^{t}m^{t-i}\text{BatchStats}_i$
+
+    $\text{InferenceStats}_t = \frac{\left(1.0 - m\right)\sum_{i=0}^{t}m^{t-i}
+    \text{BatchStats}_i}{\text{Max} \left(1.0 - m^{t+1}, \varepsilon \right)}$
+
+
+    $\text{Note: } \frac{(1.0 - m)\sum_{i=0}^{t}m^{t-i}}{1.0 - m^{t+1}} =
+    \frac{(1.0 - m)\sum_{i=0}^{t}m^{i}}{1.0 - m^{t+1}}$
+    $= \frac{(1.0 - m)\frac{1.0 - m^{t+1}}{1.0 - m}}{1.0 - m^{t+1}} = 1$
+
+
+    ??? cite
+
+        [Batch Normalization: Accelerating Deep Network Training by Reducing
+         Internal Covariate Shift](https://arxiv.org/abs/1502.03167)
+
+        ```bibtex
+        @article{DBLP:journals/corr/IoffeS15,
+        author       = {Sergey Ioffe and
+                        Christian Szegedy},
+        title        = {Batch Normalization: Accelerating Deep Network Training
+                        by Reducing Internal Covariate Shift},
+        journal      = {CoRR},
+        volume       = {abs/1502.03167},
+        year         = {2015},
+        url          = {http://arxiv.org/abs/1502.03167},
+        eprinttype    = {arXiv},
+        eprint       = {1502.03167},
+        timestamp    = {Mon, 13 Aug 2018 16:47:06 +0200},
+        biburl       = {https://dblp.org/rec/journals/corr/IoffeS15.bib},
+        bibsource    = {dblp computer science bibliography, https://dblp.org}
+        }
+        ```
+
     """  # noqa: E501
 
     weight: Optional[Float[Array, "input_size"]]
@@ -67,7 +122,7 @@ def __init__(
         eps: float = 1e-5,
         channelwise_affine: bool = True,
         momentum: float = 0.99,
-        warmup_period: int = 1000,
+        warmup_period: int = 1,
         inference: bool = False,
         dtype=None,
     ):
@@ -86,8 +141,8 @@ def __init__(
             parameters.
         - `momentum`: The rate at which to update the running statistics. Should be a
             value between 0 and 1 exclusive.
-        - `warmup_period`: The period to warm up the running statistics. Only used when
-            `approach=\"ema\"`.
+        - `warmup_period`: The interpolation period between batch and running
+            statistics. Only used when `approach=\"ema\"`.
         - `inference`: If `False` then the batch means and variances will be calculated
             and used to update the running statistics. If `True` then the running
             statistics are directly used for normalisation. This may be toggled with
@@ -107,6 +162,9 @@ def __init__(
             raise ValueError(f"approach must be one of {valid_approaches}")
         self.approach = approach
 
+        if warmup_period < 1:
+            raise ValueError("warmup_period must be >= 1")
+
         if channelwise_affine:
             self.weight = jnp.ones((input_size,))
             self.bias = jnp.zeros((input_size,))
@@ -128,7 +186,7 @@ def __init__(
         self.eps = eps
         self.channelwise_affine = channelwise_affine
         self.momentum = momentum
-        self.warmup_period = max(1, warmup_period)
+        self.warmup_period = warmup_period
 
     @jax.named_scope("eqx.nn.BatchNorm")
     def __call__(
@@ -182,16 +240,15 @@ def _stats(y):
 
             momentum = self.momentum
             zero_frac = state.get(self.zero_frac_index)
-            zero_frac *= momentum
+            zero_frac = zero_frac * momentum
             state = state.set(self.zero_frac_index, zero_frac)
 
             batch_mean, batch_var = jax.vmap(_stats)(x)
             running_mean, running_var = state.get(self.state_index)
             running_mean = (1 - momentum) * batch_mean + momentum * running_mean
-            running_var = (1 - momentum) * batch_var + momentum * running_var
-            state = state.set(self.state_index, (running_mean, running_var))
 
             if self.approach == "ema":
+                running_var = (1 - momentum) * batch_var + momentum * running_var
                 warmup_count = state.get(self.count_index)
                 warmup_count = jnp.minimum(warmup_count + 1, self.warmup_period)
                 state = state.set(self.count_index, warmup_count)
@@ -202,8 +259,15 @@ def _stats(y):
                 norm_var = zero_frac * batch_var + running_var
                 norm_var = (1.0 - warmup_frac) * batch_var + warmup_frac * norm_var
             else:
+                axis_size = jax.lax.psum(jnp.array(1.0), self.axis_name)
+                debias_coef = (axis_size) / jnp.maximum(axis_size - 1, self.eps)
+                running_var = (
+                    1 - momentum
+                ) * debias_coef * batch_var + momentum * running_var
                 norm_mean, norm_var = batch_mean, batch_var
 
+            state = state.set(self.state_index, (running_mean, running_var))
+
         def _norm(y, m, v, w, b):
             out = (y - m) / jnp.sqrt(v + self.eps)
             if self.channelwise_affine:
diff --git a/tests/test_nn.py b/tests/test_nn.py
index 6d754461..4685f3d1 100644
--- a/tests/test_nn.py
+++ b/tests/test_nn.py
@@ -830,9 +830,14 @@ def test_batch_norm(getkey):
         bn = eqx.nn.BatchNorm(5, "batch")
         assert bn.approach == "batch"
 
+    with pytest.raises(ValueError):
+        bn = eqx.nn.BatchNorm(5, "batch", approach="ema", warmup_period=0)
+
     # Test initialization
     bn_momentum = 0.99
-    bn = eqx.nn.BatchNorm(5, "batch", approach="ema", momentum=bn_momentum)
+    bn = eqx.nn.BatchNorm(
+        5, "batch", approach="ema", warmup_period=10, momentum=bn_momentum
+    )
     state = eqx.nn.State(bn)
     vbn = jax.vmap(bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
     running_mean, running_var = state.get(bn.state_index)
@@ -890,8 +895,7 @@ def test_batch_norm(getkey):
     assert running_mean.shape == (6,)
     assert running_var.shape == (6,)
 
-    # Test that it normalises
-
+    # Test that approach=ema normalises
     x1alt = jrandom.normal(jrandom.PRNGKey(5678), (10, 5))  # avoid flakey test
     bn = eqx.nn.BatchNorm(5, "batch", channelwise_affine=False, approach="ema")
     state = eqx.nn.State(bn)
@@ -902,6 +906,27 @@ def test_batch_norm(getkey):
     )
     assert jnp.allclose(out, true_out)
 
+    # Test that approach=batch normalises in training mode
+    bn = eqx.nn.BatchNorm(
+        5, "batch", channelwise_affine=False, approach="batch", momentum=0.9
+    )
+    state = eqx.nn.State(bn)
+    vbn = jax.vmap(bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
+    out, state = vbn(x1alt, state)
+    true_out = (x1alt - jnp.mean(x1alt, axis=0, keepdims=True)) / jnp.sqrt(
+        jnp.var(x1alt, axis=0, keepdims=True) + 1e-5
+    )
+    assert jnp.allclose(out, true_out)
+    # Test that approach=batch normaises in inference mode
+    bn_inf = eqx.nn.inference_mode(bn, value=True)
+    vbn_inf = jax.vmap(bn_inf, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
+    out, state = vbn_inf(x1alt, state)
+    debias_coef = x1alt.shape[0] / (x1alt.shape[0] - 1)
+    true_out = (x1alt - jnp.mean(x1alt, axis=0, keepdims=True)) / jnp.sqrt(
+        debias_coef * jnp.var(x1alt, axis=0, keepdims=True) + 1e-5
+    )
+    assert jnp.allclose(out, true_out)
+
     # Test that the statistics update during training
     out, state = vbn(x1, state)
     running_mean, running_var = state.get(bn.state_index)
@@ -913,7 +938,6 @@ def test_batch_norm(getkey):
     assert not jnp.allclose(running_var, running_var2)
 
     # Test that the statistics don't update at inference
-
     ibn = eqx.nn.inference_mode(bn, value=True)
     vibn = jax.vmap(ibn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
     out, state = vibn(4 * x1 + 20, state)

From e7705b69593ad24f4286ae41b04b49b809d80fcf Mon Sep 17 00:00:00 2001
From: andrewdipper <andrewdipper11235@gmail.com>
Date: Thu, 14 Mar 2024 13:06:07 -0700
Subject: [PATCH 4/6] reformat + comments

---
 equinox/nn/_batch_norm.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index 67795ffb..f24ca918 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -224,6 +224,7 @@ def __call__(
         if inference is None:
             inference = self.inference
         if inference:
+            # renormalize running stats to account for the zeroed part
             zero_frac = state.get(self.zero_frac_index)
             running_mean, running_var = state.get(self.state_index)
             norm_mean = running_mean / jnp.maximum(1.0 - zero_frac, self.eps)
@@ -253,17 +254,23 @@ def _stats(y):
                 warmup_count = jnp.minimum(warmup_count + 1, self.warmup_period)
                 state = state.set(self.count_index, warmup_count)
 
+                # fill in unpopulated part of running stats with batch stats
                 warmup_frac = warmup_count / self.warmup_period
                 norm_mean = zero_frac * batch_mean + running_mean
-                norm_mean = (1.0 - warmup_frac) * batch_mean + warmup_frac * norm_mean
                 norm_var = zero_frac * batch_var + running_var
+
+                # apply warmup interpolation between batch and running statistics
+                norm_mean = (1.0 - warmup_frac) * batch_mean + warmup_frac * norm_mean
                 norm_var = (1.0 - warmup_frac) * batch_var + warmup_frac * norm_var
             else:
+                # calculate unbiased variance for saving
                 axis_size = jax.lax.psum(jnp.array(1.0), self.axis_name)
                 debias_coef = (axis_size) / jnp.maximum(axis_size - 1, self.eps)
                 running_var = (
                     1 - momentum
                 ) * debias_coef * batch_var + momentum * running_var
+
+                # just use batch statistics when not in inference mode
                 norm_mean, norm_var = batch_mean, batch_var
 
             state = state.set(self.state_index, (running_mean, running_var))

From 7a99c7f044e6c9c7e5ef3e1d84c23081d9cdaf91 Mon Sep 17 00:00:00 2001
From: andrewdipper <andrewdipper11235@gmail.com>
Date: Sun, 31 Mar 2024 22:05:49 -0700
Subject: [PATCH 5/6] add BatchNorm backward compatibility mode

---
 equinox/nn/_batch_norm.py | 46 ++++++++++++++++++++++++++++-----------
 tests/test_nn.py          |  2 +-
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index f24ca918..ab565ffc 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -44,7 +44,8 @@ class BatchNorm(StatefulLayer, strict=True):
 
     With `approach = "batch"` during training the batch mean and variance are used
     for normalization.  For inference the exponential running mean and ubiased
-    variance are used for normalization in accordance with the cited paper below:
+    variance are used for normalization in accordance with the cited paper below.
+    Let `m` be momentum:
 
     $\text{TrainStats}_t = \text{BatchStats}_t$
 
@@ -65,13 +66,16 @@ class BatchNorm(StatefulLayer, strict=True):
     \text{WarmupFrac}_t * \left(1.0 - m\right)\sum_{i=0}^{t}m^{t-i}\text{BatchStats}_i$
 
     $\text{InferenceStats}_t = \frac{\left(1.0 - m\right)\sum_{i=0}^{t}m^{t-i}
-    \text{BatchStats}_i}{\text{Max} \left(1.0 - m^{t+1}, \varepsilon \right)}$
+    \text{BatchStats}_i}{\text{max} \left(1.0 - m^{t+1}, \varepsilon \right)}$
 
 
     $\text{Note: } \frac{(1.0 - m)\sum_{i=0}^{t}m^{t-i}}{1.0 - m^{t+1}} =
     \frac{(1.0 - m)\sum_{i=0}^{t}m^{i}}{1.0 - m^{t+1}}$
     $= \frac{(1.0 - m)\frac{1.0 - m^{t+1}}{1.0 - m}}{1.0 - m^{t+1}} = 1$
 
+    `approach = "ema_compatibility"` reproduces the original equinox BatchNorm
+    behavior.  It often results in training instabilities and `approach = "batch"`
+    or `"ema"` is recommended.
 
     ??? cite
 
@@ -108,7 +112,7 @@ class BatchNorm(StatefulLayer, strict=True):
     axis_name: Union[Hashable, Sequence[Hashable]]
     inference: bool
     input_size: int = field(static=True)
-    approach: Literal["batch", "ema"] = field(static=True)
+    approach: Literal["batch", "ema", "ema_compatibility"] = field(static=True)
     eps: float = field(static=True)
     channelwise_affine: bool = field(static=True)
     momentum: float = field(static=True)
@@ -118,7 +122,7 @@ def __init__(
         self,
         input_size: int,
         axis_name: Union[Hashable, Sequence[Hashable]],
-        approach: Optional[Literal["batch", "ema"]] = None,
+        approach: Optional[Literal["batch", "ema", "ema_compatibility"]] = None,
         eps: float = 1e-5,
         channelwise_affine: bool = True,
         momentum: float = 0.99,
@@ -154,10 +158,16 @@ def __init__(
         """
 
         if approach is None:
-            warnings.warn('BatchNorm approach is None, defaults to approach="batch"')
-            approach = "batch"
-
-        valid_approaches = {"batch", "ema"}
+            warnings.warn(
+                "BatchNorm approach is None, defaults to "
+                'approach="ema_compatibility".  This is not recommended as '
+                'it can lead to training instability.  Use "batch" or '
+                'alternatively "ema" with appropriately selected warmup '
+                "instead."
+            )
+            approach = "ema_compatibility"
+
+        valid_approaches = {"batch", "ema", "ema_compatibility"}
         if approach not in valid_approaches:
             raise ValueError(f"approach must be one of {valid_approaches}")
         self.approach = approach
@@ -240,15 +250,13 @@ def _stats(y):
                 return mean, var
 
             momentum = self.momentum
-            zero_frac = state.get(self.zero_frac_index)
-            zero_frac = zero_frac * momentum
-            state = state.set(self.zero_frac_index, zero_frac)
-
             batch_mean, batch_var = jax.vmap(_stats)(x)
+            zero_frac = state.get(self.zero_frac_index)
             running_mean, running_var = state.get(self.state_index)
-            running_mean = (1 - momentum) * batch_mean + momentum * running_mean
 
             if self.approach == "ema":
+                zero_frac = zero_frac * momentum
+                running_mean = (1 - momentum) * batch_mean + momentum * running_mean
                 running_var = (1 - momentum) * batch_var + momentum * running_var
                 warmup_count = state.get(self.count_index)
                 warmup_count = jnp.minimum(warmup_count + 1, self.warmup_period)
@@ -262,7 +270,18 @@ def _stats(y):
                 # apply warmup interpolation between batch and running statistics
                 norm_mean = (1.0 - warmup_frac) * batch_mean + warmup_frac * norm_mean
                 norm_var = (1.0 - warmup_frac) * batch_var + warmup_frac * norm_var
+
+            elif self.approach == "ema_compatibility":
+                running_mean = (1 - momentum) * batch_mean + momentum * running_mean
+                running_var = (1 - momentum) * batch_var + momentum * running_var
+                running_mean = lax.select(zero_frac == 1.0, batch_mean, running_mean)
+                running_var = lax.select(zero_frac == 1.0, batch_var, running_var)
+                norm_mean, norm_var = running_mean, running_var
+                zero_frac = 0.0 * zero_frac
+
             else:
+                zero_frac = zero_frac * momentum
+                running_mean = (1 - momentum) * batch_mean + momentum * running_mean
                 # calculate unbiased variance for saving
                 axis_size = jax.lax.psum(jnp.array(1.0), self.axis_name)
                 debias_coef = (axis_size) / jnp.maximum(axis_size - 1, self.eps)
@@ -273,6 +292,7 @@ def _stats(y):
                 # just use batch statistics when not in inference mode
                 norm_mean, norm_var = batch_mean, batch_var
 
+            state = state.set(self.zero_frac_index, zero_frac)
             state = state.set(self.state_index, (running_mean, running_var))
 
         def _norm(y, m, v, w, b):
diff --git a/tests/test_nn.py b/tests/test_nn.py
index 4685f3d1..fb3704b0 100644
--- a/tests/test_nn.py
+++ b/tests/test_nn.py
@@ -828,7 +828,7 @@ def test_batch_norm(getkey):
     # Test that it warns with no approach - defaulting to batch
     with pytest.warns(UserWarning):
         bn = eqx.nn.BatchNorm(5, "batch")
-        assert bn.approach == "batch"
+        assert bn.approach == "ema_compatibility"
 
     with pytest.raises(ValueError):
         bn = eqx.nn.BatchNorm(5, "batch", approach="ema", warmup_period=0)

From 96f4a9652aaab8e014ea2b008824a07813560550 Mon Sep 17 00:00:00 2001
From: andrewdipper <andrewdipper11235@gmail.com>
Date: Sun, 31 Mar 2024 22:32:04 -0700
Subject: [PATCH 6/6] fix description

---
 equinox/nn/_batch_norm.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index ab565ffc..0ac300aa 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -137,9 +137,9 @@ def __init__(
             to `axis_name` in `jax.vmap` or `jax.pmap`. Can also be a sequence (e.g. a
             tuple or a list) of names, to compute statistics over multiple named axes.
         - `approach`: The approach to use for the running statistics. If `approach=None`
-            a warning will be raised and approach will default to `"batch"`. During
-            training `"batch"` only uses batch statisics while`"ema"` uses the running
-            statistics.
+            a warning will be raised and approach will default to `"ema_compatibility"`.
+            During training `"batch"` only uses batch statisics while`"ema"` and
+            `"ema_compatibility"` uses the running statistics.
         - `eps`: Value added to the denominator for numerical stability.
         - `channelwise_affine`: Whether the module has learnable channel-wise affine
             parameters.