RMSNormalization¶

RMSNormalization - 23¶

Version¶

name: RMSNormalization (GitHub)
domain: main
since_version: 23
function: True
support_level: SupportType.COMMON
shape inference: True

This version of the operator has been available since version 23.

Summary¶

This is RMS normalization defined in ONNX as function as described in the paper https://arxiv.org/pdf/1910.07467. The overall computation can be split into two stages. The root mean squared norm is taken over the last D dimensions, where D is the dimension of normalized_shape. For example, if normalized_shape is (3, 5) (a 2-dimensional shape), the rms norm is computed over the last 2 dimensions of the input. The computation required by standardization can be described by the following equations.

XSquared = Mul(X, X)
XSquaredMean = ReduceMean&lt;axes=normalized_axes&gt;(XSquared)
MeanSquareEpsilon = Add(XSquaredMean, epsilon)
RMS = Sqrt(MeanSquareEpsilon)
Normalized = Div(X, RMS)

where normalized_axes is [axis, ..., rank of X - 1]. The variables RMS stand for root mean square, Depending on stash_type attribute, the actual computation must happen in different floating-point precision. For example, if stash_type is 1, this operator casts all input variables to 32-bit float, perform the computation, and finally cast Normalized back to the original type of X. The second stage then scales the outcome of the first stage using:

Y= Mul(Normalized, Scale)

Let d[i] indicate the i-th dimension of X. If X’s shape is [d[0], ..., d[axis-1], d[axis], ..., d[rank-1]], the shape of RMS is [d[0], ..., d[axis-1], 1, ..., 1]. Y and X have the same shape. This operator supports unidirectional broadcasting (Scale should be unidirectional broadcastable to tensor X); for more details please check Broadcasting in ONNX.

Attributes¶

axis - INT (default is -1):

The first normalization dimension. If rank(X) is r, axis’ allowed range is [-r, r). Negative value means counting dimensions from the back.
epsilon - FLOAT (default is 1e-05):

The epsilon value to use to avoid division by zero.
stash_type - INT (default is 1):

The floating-point precision used in stage one of the computation.

Inputs¶

X (heterogeneous) - T:

The input tensor to be normalized. In general, the shape is (D1, D2, … , Dn) for n-dimensional data, where the root mean squared norm is taken over the last D dimensions, D is determined by the axis attribute.
scale (heterogeneous) - V:

Scale tensor. Scale tensor shape should be broadcastable to the normalized shape.

Outputs¶

Y (heterogeneous) - V:

Output data tensor. Same shape as X

Type Constraints¶

T in ( tensor(bfloat16), tensor(double), tensor(float), tensor(float16) ):

Constrain input X type to float tensors.
V in ( tensor(bfloat16), tensor(double), tensor(float), tensor(float16) ):

Constrain output Y and scale type to float tensors.

Examples¶

default¶

import numpy as np
import onnx

X = np.random.randn(2, 3, 4, 5).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    Y = _rms_normalization(X, W, axis=axis)

    node = onnx.helper.make_node(
        "RMSNormalization",
        inputs=["X", "W"],
        outputs=["Y"],
        axis=axis,
    )

    if axis < 0:
        name = f"test_rms_normalization_4d_axis_negative_{-axis}"
    else:
        name = f"test_rms_normalization_4d_axis{axis}"

    expect(node, inputs=[X, W], outputs=[Y], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

_default_axis¶

import numpy as np
import onnx

X = np.random.randn(2, 3, 4, 5).astype(np.float32)

# Default axis in RMSNormalization is -1.
normalized_shape = calculate_normalized_shape(X.shape, -1)
W = np.random.randn(*normalized_shape).astype(np.float32)
# Axis is default to -1 in the reference implementation.
Y = _rms_normalization(X, W)

# Not specifying axis attribute means -1.
node = onnx.helper.make_node(
    "RMSNormalization",
    inputs=["X", "W"],
    outputs=["Y"],
)

expect(
    node,
    inputs=[X, W],
    outputs=[Y],
    name="test_rms_normalization_default_axis",
)

2d¶

import numpy as np
import onnx

X = np.random.randn(3, 4).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    Y = _rms_normalization(X, W, axis=axis)

    node = onnx.helper.make_node(
        "RMSNormalization",
        inputs=["X", "W"],
        outputs=["Y"],
        axis=axis,
    )

    if axis < 0:
        name = f"test_rms_normalization_2d_axis_negative_{-axis}"
    else:
        name = f"test_rms_normalization_2d_axis{axis}"

    expect(node, inputs=[X, W], outputs=[Y], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))

3d_epsilon¶

import numpy as np
import onnx

epsilon = 1e-1
X = np.random.randn(2, 3, 5).astype(np.float32)

def case(axis: int) -> None:
    normalized_shape = calculate_normalized_shape(X.shape, axis)
    W = np.random.randn(*normalized_shape).astype(np.float32)
    Y = _rms_normalization(X, W, axis=axis, epsilon=epsilon)
    node = onnx.helper.make_node(
        "RMSNormalization",
        inputs=["X", "W"],
        outputs=["Y"],
        axis=axis,
        epsilon=epsilon,
    )

    if axis < 0:
        name = f"test_rms_normalization_3d_axis_negative_{-axis}_epsilon"
    else:
        name = f"test_rms_normalization_3d_axis{axis}_epsilon"

    expect(node, inputs=[X, W], outputs=[Y], name=name)

for i in range(len(X.shape)):
    case(i)
    case(i - len(X.shape))