Add BatchRenorm
layer to linen.normalization
#3822
Labels
Priority: P2 - no schedule
Best effort response and resolution. We have no plan to work on this at the moment.
I propose adding a batch renormalization (BatchRenorm) layer to flax.
I would be happy to make a PR.
BatchRenorm (https://arxiv.org/pdf/1702.03275.pdf) is an improved version of the vanilla BatchNorm layer. The difference to BatchNorm is that after a warm-up period, the running statistics are used to normalize the batch, both in train and eval mode. This helps to deal with BatchNorm's stability issues during long training runs. In contrast, BatchNorm uses the min batch statistics during train mode.
Alternatively, the BatchNorm layer could be refactored to support renormalization. However, I believe that it would be cleaner to put this into a separate BatchRenorm class.
Just recently, BatchRenorm has been shown to yield new state-of-the-art results in deep reinforcement learning (https://openreview.net/pdf?id=PczQtTsTIX), and I believe this might also lead to wider adoption in this community.
The text was updated successfully, but these errors were encountered: