How to do Batch Normalization in JAX/Flax?

[marcvanzee]

[marcvanzee]

The BatchNorm module is in normalization.py. The canonical example using it is Imagenet

In a multi-device setting, every device updates its normalizing parameters (ra_mean and ra_var in the code) based on its own batch statistics (unless we specify axis_name, which makes it global). However, these parameters are only stored in in Module variablebatch_stats and not in params, so they are never synced across devices unless specifically done so.

If they aren't synced, they can theoretically diverge, but if your data is fairly uniform across shards they're likely to trend towards similar values. Syncing before eval is definitely a good idea though, since otherwise your eval results will depend on which devices process which examples.

axis_index_groups can be used to set for vanilla batchnorm (aggregate from a single device) or synchronized batchnorm (aggregate from all devices). It controls how large a device group you want to sync over. The current example in the comment is a bit confusing. If you just want to aggregate over all devices in the axis you don't need to pass axis_index_groups though.