google / google/flax

Add `BatchRenorm` layer to `linen.normalization`

Open
#3,822 0 comments 0 reactions 0 assignees View on GitHub
Priority: P2 - no schedule
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

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.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.