BatchNorm training instability fix by andrewdipper · Pull Request #675 · patrick-kidger/equinox

andrewdipper · 2024-03-07T02:23:59Z

This is in reference to issue 659.

I modified BatchNorm to have two approaches "batch" and "ema". "batch" just uses the batch statistics during training time. If approach is not specified it defaults to "batch" with a warning. It's robust and seems to be the standard choice - it's far less likely to kill a model just by adding it.

"ema" is based of the smooth start method in the above issue. So keep a running mean and variance but instead of renormalizing Adam style the parts of the running averages that are zeroed are filled with the batch statistics. The problem is it's still not robust - the momentum parameter is simultaneously specifying a warmup period (when we're expecting the input distribution to change significantly) and how long we want the running average to be. So I added a linear warmup period.

Now for any choice of momentum there seems to be a warmup_period choice that will give good results. And validation performance was at least as good as with batch mode for my tests. However, I don't see a good default for warmup_period.

Some considerations:

having approach="batch" and the common axis_name="batch" is a little awkward
There's an example using BatchNorm - that will start raising a warning and should probably get changed
The current BatchNorm behavior can't be exactly replicated (ema / momentum=0.99 / warmup_period=1) is close but different at the start
There's one more piece of state hence the test_stateful.py change. Though this could be conditionally removed for approach="batch" if desired

Let me know what you think or if any changes or tests need to be added

…bility when using running statistics.

patrick-kidger · 2024-03-10T08:23:14Z

        self.eps = eps
        self.channelwise_affine = channelwise_affine
        self.momentum = momentum
+        self.warmup_period = max(1, warmup_period)


Why the max? Perhaps it would be better to just error out on values that are too small?

warmup_period=0 seemed natural for off - Changed to just check and error out

patrick-kidger · 2024-03-10T08:24:01Z

+        self.warmup_period = max(1, warmup_period)

    @jax.named_scope("eqx.nn.BatchNorm")
    def __call__(


It's not completely obvious to me that the ema implementation, with default arguments, reproduces the previous behaviour. (For example, we have warmup_period=1000 by default?)

Can you add some comments explaining what each approach corresponds to?

ema with warmup_period=1 approximately reproduces previous behavior. As I noted the start is different because of how the running statistics are initially populated. With warmup_period=1 there's no interpolation between the batch and running stats - the running stats are always used as with the previous behavior. I can give an exact replication with an extra approach if necessary.

Added some to the documentation

I think an exact replication is probably important for the default behaviour, just because I'd like to be sure that we're bit-for-bit backward compatible.

Makes sense, it was different enough that I added it as "ema_compatibility". I changed the warning to rather strongly recommend against using "ema_compatibility". I haven't found a use case where I wouldn't expect to see the instability (at least with a larger learning rate) but that could very much be due to a lack of imagination on my part.. That part can definitely change if needed

patrick-kidger · 2024-03-10T08:26:39Z

-            state = state.set(self.first_time_index, jnp.array(False))
+            momentum = self.momentum
+            zero_frac = state.get(self.zero_frac_index)
+            zero_frac *= momentum


Stylistic nit: I tend not to use the inplace operations in JAX code. This (a) fits with the functional style a bit better, and (b) emphasises that we're definitely falling back to the zero_frac = zero_frac * momentum interpretation of the syntax. (Gosh, why does Python has two different meanings for the same syntax?)

makes sense, done

patrick-kidger · 2024-03-10T08:27:19Z

            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


These don't appear to be used on the batch branch. I think the lines here can be reorganised to keep each approach only using the things it needs.

these are used by the batch branch when we're in inference mode so they still need to be computed and stored

patrick-kidger · 2024-03-10T08:29:36Z

+                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


I'm definitely going to have to sit down and grok what's going on here more carefully! As above it would be good to have some comments / docstrings / references / etc. describing what each approach is meant to do.

(C.f. something like the MultiheadAttention docstring for an example on how to use LaTeX if it'd be helpful.)

Added some commentary and tried making it a bit cleaner.

But overall batch mode should follow the cited paper. Ema follows the prior behavior but changes the initialization of the running stats and adds interpolation so it can be stable while training.

andrewdipper · 2024-03-14T20:21:39Z

+                debias_coef = (axis_size) / jnp.maximum(axis_size - 1, self.eps)
+                running_var = (
+                    1 - momentum
+                ) * debias_coef * batch_var + momentum * running_var


I neglected to use unbiased variance so corrected that here

andrewdipper · 2024-03-14T20:22:36Z

                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)


I'm using this to get the length of the "batch" axis - but not sure it's the best / correct way

I think this is the correct way! IIRC psum(1) is actually special-cased for this purpose.

patrick-kidger · 2024-03-23T22:03:02Z

+        - `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.


So continuing from my previous comment -- probably the default should be ema if approach=None.

patrick-kidger

Okay! Sorry for taking so long to getting back around to reviewing this.

Lmk once you're happy that the previous behaviour is replicated by default, and I'll sit down with a pen and paper and satisfy myself that the calculations all look reasonable!

andrewdipper · 2024-04-01T05:26:19Z

All good - I got caught up in other things myself!

From my tests the replication is exact now. It added another approach that is very similar to "ema" but it seemed like the most reasonable way to organize it. Let me know if anything isn't clear

lockwo · 2025-02-06T17:08:57Z

What's the status of this PR? I'm encountering some equinox batch norm related trouble as I switch over some code from haiku, and I think it's related to some of the zero debiased issues also tackled in this PR

patrick-kidger · 2025-02-06T20:06:13Z

Ah, the refactor here got large enough that I never found time to check that it was (a) still doing the same thing for backward compatibility; (b) was doing things that I thought were correct going forward.

If you're interested -- and using such functionality in practice (I pretty much never use BatchNorm, which is part of why it's hard for me to really confirm this is doing the right thing) -- then I'd welcome your input here!

lockwo · 2025-02-06T21:44:47Z

Makes sense, I'll see about taking a stab at a small-ish change to allow for an option to align it with other packages. Tbh I haven't really followed the lore of batch norm ever since the whole internal covariate shift illusion debacle. I did follow the whole drama in 2022 around the post-batch norm models and ViT vs CNN with stuff like https://arxiv.org/abs/2102.06171 and https://arxiv.org/abs/2201.03545 (although idk what is SotA these days), but for this specific project I'm just working on some AlphaZero based stuff now and it has batch norm.

Mostly I just want the operations to be the same between packages (or at least, there is an option to be the same) rather than the leaning too heavily on the theory (since even if flax and haiku are doing things wrong in a different way most probably just use it, and it basically becomes the expected way).

andrewdipper added 2 commits March 6, 2024 16:00

Update BatchNorm to add mode for using batch statisics / increase sta…

346e472

…bility when using running statistics.

fix batchnorm type definiton

cc021f6

patrick-kidger reviewed Mar 10, 2024

View reviewed changes

andrewdipper added 2 commits March 14, 2024 12:44

use unbiased variance + minor fixes

6bfa77f

reformat + comments

e7705b6

andrewdipper commented Mar 14, 2024

View reviewed changes

andrewdipper requested a review from patrick-kidger March 15, 2024 06:45

patrick-kidger reviewed Mar 23, 2024

View reviewed changes

add BatchNorm backward compatibility mode

7a99c7f

fix description

96f4a96

andrewdipper requested a review from patrick-kidger April 1, 2024 15:07

patrick-kidger mentioned this pull request Apr 18, 2024

Possible bug in AdaptivePool #641

Open

patrick-kidger force-pushed the main branch 2 times, most recently from 2905390 to a386f27 Compare April 20, 2024 09:26

lockwo mentioned this pull request Feb 10, 2025

Batch BatchNorm #948

Merged

lockwo mentioned this pull request Feb 24, 2025

BatchNorm initialization #955

Closed

Uh oh!

Conversation

andrewdipper commented Mar 7, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

patrick-kidger left a comment

Choose a reason for hiding this comment

Uh oh!

andrewdipper commented Apr 1, 2024

Uh oh!

lockwo commented Feb 6, 2025

Uh oh!

patrick-kidger commented Feb 6, 2025

Uh oh!

lockwo commented Feb 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lockwo commented Feb 6, 2025 •

edited

Loading