tf.keras.ops.rms_normalization

Performs Root Mean Square (RMS) normalization on x.

View aliases

Main aliases

tf.keras.ops.nn.rms_normalization

tf.keras.ops.rms_normalization(
    x, scale=None, axis=-1, epsilon=None
)

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

The operation is different from LayerNormalization with RMS scaling.

It is defined as rms_normalization(x) = x * rsqrt(mean(square(x))) * scale

Args
x	Input tensor.
scale	Optional scaling factor for the normalization.
axis	The axis or axes along which to perform normalization. Defaults to -1.
epsilon	A lower bound value for the norm. Defaults to backend.epsilon().

Returns
The normalized array.

Example:

x = keras.random.normal((1, 10))
keras.ops.rms_normalization(x)
array([[0.69384296, 0.94444374, 0.16551171, 0.05749961, 1.11008865,
        0.52475186, 1.57686807, 1.69893307, 1.27292764, 0.30819128]])