Added RMS normalization layer #2881
Conversation
Codecov Report
@@ Coverage Diff @@
## main #2881 +/- ##
==========================================
+ Coverage 81.45% 81.47% +0.02%
==========================================
Files 55 55
Lines 5779 5798 +19
==========================================
+ Hits 4707 4724 +17
- Misses 1072 1074 +2
Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here. |
| @@ -335,6 +343,70 @@ def __call__(self, x): | |||
| self.bias_init, self.scale_init) | |||
|
|
|||
|
|
|||
| class RMSNorm(Module): | |||
There was a problem hiding this comment.
Could you please add an example of how to use the layer? I think we should starting doing this for every layer, like @cgarciae does in his RNN PR: https://github.com/google/flax/pull/2604/files#r1107264719.
chiamp
force-pushed
the
rmslayernorm
branch
2 times, most recently
from
4a05ad4
to
42ba933
Compare
|
@marcvanzee @levskaya I added a docstring, let me know if this works! |
|
@chiamp - I added a exception for the deprecation warning, your tests all seem to pass now! |
flax/linen/normalization.py
Outdated
| epsilon: float = 1e-6 | ||
| dtype: Optional[Dtype] = None | ||
| param_dtype: Dtype = jnp.float32 | ||
| use_bias: bool = True |
There was a problem hiding this comment.
sorry I just noticed this - we probably don't want use_bias and bias_init here since we're never adjusting the offset?
Resolves #2849.
Added an optional argument
use_meanin the_compute_statsfunction inflax/linen/normalization.py, which will compute the mean and variance if set toTrue, and will set the mean to 0 and compute the variance without subtracting the mean if set toFalse. The latter mode is useful as square rooting this "variance" value (which is done in the_normalizefunction) will give you the RMS.