tf.keras.layers.RMSNormalization

Root Mean Square (RMS) Normalization layer.

Inherits From: Layer, Operation

tf.keras.layers.RMSNormalization(
    axis=-1, epsilon=1e-06, **kwargs
)

This layer normalizes the input tensor based on its RMS value.

The Keras layer performs the operation as described in Root Mean Square Layer Normalization by Biao Zhang et al.

If scale is enabled, the layer will scale the normalized outputs via a learnable scaling factor.

So, with scaling enabled, the normalization equations are as follows:

Let the intermediate activations for a mini-batch to be the inputs.

rms_normalization(x) = x * rsqrt(mean(square(x))) * scale

For example:

layer = keras.layers.RMSNormalization()
layer.build([5, 20, 30, 10])
print(layer.scale.shape)
(10,)
layer(np.random.rand(1, 10)).numpy()
array([[0.35098287, 1.0495652 , 1.4645109 , 1.2944688 , 0.31124955,
        1.2768592 , 1.184331  , 0.17474432, 0.49955517, 1.2428929 ]],
    dtype=float32)

axis int. The axis on which to perform the normalization.
epsilon float. A small number to add to avoid division by zero.

input Retrieves the input tensor(s) of a symbolic operation.

Only returns the tensor(s) corresponding to the first time the operation was called.

output Retrieves the output tensor(s) of a layer.

Only returns the tensor(s) corresponding to the first time the operation was called.

Methods

from_config

View source

@classmethod
from_config(
    config
)

Creates an operation from its config.

This method is the reverse of get_config, capable of instantiating the same operation from the config dictionary.

if "dtype" in config and isinstance(config["dtype"], dict):
    policy = dtype_policies.deserialize(config["dtype"])

Args
config A Python dictionary, typically the output of get_config.

Returns
An operation instance.