Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add BatchRenorm layer to linen.normalization #3822

Open
danielpalen opened this issue Apr 3, 2024 · 0 comments
Open

Add BatchRenorm layer to linen.normalization #3822

danielpalen opened this issue Apr 3, 2024 · 0 comments
Labels
Priority: P2 - no schedule Best effort response and resolution. We have no plan to work on this at the moment.

Comments

@danielpalen
Copy link

danielpalen commented Apr 3, 2024

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.

@danielpalen danielpalen changed the title Add BatchRenorm layer to linen.normalization Add BatchRenorm layer to linen.normalization Apr 3, 2024
@chiamp chiamp added the Priority: P2 - no schedule Best effort response and resolution. We have no plan to work on this at the moment. label Apr 10, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Priority: P2 - no schedule Best effort response and resolution. We have no plan to work on this at the moment.
Projects
None yet
Development

No branches or pull requests

2 participants