First of all, thank you for your great library!
The problem
- Initialize a model with multiple batch norms (e.g. ResNet)
- Set model to inference mode
- Observe numerical instability (values > 10^9 after a few ResNet-Blocks)
If I read your code correctly variance in BatchNorm is initialized to zero, meaning that if you use BatchNorm in inference without any training (think: testing random model performance on dataset or the first iteration of an RL algorithm), the BatchNorm scales all values by $\frac{1}{\sqrt{eps}} = 10^4$. Both PyTorch and Flax use $\sigma^2 = 1$ in their code. I think this is a much saner default, considering that $\sigma^2 = 0$ doesn't really make sense.
Tl;dr:
- BatchNorm causes instability due to the initial value of the variance
- Both Flax and PyTorch use variance = 1 as their initial value
First of all, thank you for your great library!
The problem
If I read your code correctly variance in BatchNorm is initialized to zero, meaning that if you use BatchNorm in inference without any training (think: testing random model performance on dataset or the first iteration of an RL algorithm), the BatchNorm scales all values by$\frac{1}{\sqrt{eps}} = 10^4$ . Both PyTorch and Flax use $\sigma^2 = 1$ in their code. I think this is a much saner default, considering that $\sigma^2 = 0$ doesn't really make sense.
Tl;dr: