class

RMSNormBackward

RMSNormBackward

source edit

Python wrappers

pylucid.nn.RMSNormclass

Autograd node for RMS Normalization.

Computes the per-slice reciprocal RMS $\text{rstd} = 1 / \sqrt{\tfrac{1}{N}\sum_i x_i^2 + \epsilon}$ over the trailing $D$ dimensions of x, scales the input by it, and applies the per-element gain $\gamma$ . Only saved_rstd_ is cached for backward — no mean is needed because RMSNorm does not centre the input.

Inherits FuncOp<RMSNormBackward, 2> with two saved-input slots — x and gamma (RMSNorm has no bias). outer_ * N_ always equals numel(x).

Math

For each leading-axis slice of x (flattened length $N$ ):

\text{rms}(x) = \sqrt{\frac{1}{N}\sum_{i=1}^{N} x_i^2 + \epsilon}, \qquad y_i = \frac{x_i}{\text{rms}(x)} \cdot \gamma_i

Equivalently, with the saved $\text{rstd} = 1/\text{rms}(x)$ :

y_i = x_i \cdot \text{rstd} \cdot \gamma_i

References

Zhang & Sennrich, "Root Mean Square Layer Normalization" (NeurIPS 2019, arXiv:1910.07467).

Attributes

schema_v1OpSchema

Registered schema (name="rms_norm", version=1, amp_policy=ForceFP32, produces_grad=true).

saved_rstd_Storage

Per-slice reciprocal RMS, shape (outer_,).

outer_std::size_t

Product of leading (non-normalised) dimensions of x.

N_std::size_t

Product of normalised (trailing) dimensions of x (== numel(gamma)).

Notes

Thread safety: instances are created during forward and consumed during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

int forward(const int & x, const int & gamma, double eps)

source edit

Run the forward pass.

Reduces x over its last $D$ dimensions (where $D$ is the rank of gamma), divides by the RMS, applies the gain, and (when grad mode is on) registers this as the grad_fn of the returned output. Caches saved_rstd_ for backward.

Parameters

xTensorImplPtr

Input tensor of any shape (*, normalized_shape). Rank must be >= rank(gamma).

gammaTensorImplPtr

Per-element scale; its shape must match the trailing dims of x.

epsdouble

Small constant added inside the square root for numerical stability.

Returns

TensorImplPtr

Normalised output, same shape as x.

Raises

ShapeMismatch

If gamma.ndim > x.ndim or if the trailing dims of x do not match gamma.

Methods

apply

→std::vector<Storage>

int apply(int grad_out)

source edit

Compute gradients for the two saved inputs.

Returns {dx, d_gamma} in the saved-input order. Closed-form (per leading slice of length $N$ , using the cached $\text{rstd} = 1/\text{rms}(x)$ ):

d\gamma_i = \sum_{\text{slice}} \text{grad\_out}_i\,x_i\,\text{rstd}

and the input gradient is the standard RMSNorm backward formula (the centred grad is replaced by the un-centred grad since RMSNorm has no mean subtraction).

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

Two-element vector {dx, d_gamma}.

class

RMSNormBackward

RMSNormBackward

source edit

Python wrappers

pylucid.nn.RMSNormclass

Autograd node for RMS Normalization.

Inherits FuncOp<RMSNormBackward, 2> with two saved-input slots — x and gamma (RMSNorm has no bias). outer_ * N_ always equals numel(x).

Math

For each leading-axis slice of x (flattened length $N$ ):

\text{rms}(x) = \sqrt{\frac{1}{N}\sum_{i=1}^{N} x_i^2 + \epsilon}, \qquad y_i = \frac{x_i}{\text{rms}(x)} \cdot \gamma_i

Equivalently, with the saved $\text{rstd} = 1/\text{rms}(x)$ :

y_i = x_i \cdot \text{rstd} \cdot \gamma_i

References

Zhang & Sennrich, "Root Mean Square Layer Normalization" (NeurIPS 2019, arXiv:1910.07467).

Attributes

schema_v1OpSchema

Registered schema (name="rms_norm", version=1, amp_policy=ForceFP32, produces_grad=true).

saved_rstd_Storage

Per-slice reciprocal RMS, shape (outer_,).

outer_std::size_t

Product of leading (non-normalised) dimensions of x.

N_std::size_t

Product of normalised (trailing) dimensions of x (== numel(gamma)).

Notes

Thread safety: instances are created during forward and consumed during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

int forward(const int & x, const int & gamma, double eps)

source edit

Run the forward pass.

Parameters

xTensorImplPtr

Input tensor of any shape (*, normalized_shape). Rank must be >= rank(gamma).

gammaTensorImplPtr

Per-element scale; its shape must match the trailing dims of x.

epsdouble

Small constant added inside the square root for numerical stability.

Returns

TensorImplPtr

Normalised output, same shape as x.

Raises

ShapeMismatch

If gamma.ndim > x.ndim or if the trailing dims of x do not match gamma.

Methods

apply

→std::vector<Storage>

int apply(int grad_out)

source edit

Compute gradients for the two saved inputs.

Returns {dx, d_gamma} in the saved-input order. Closed-form (per leading slice of length $N$ , using the cached $\text{rstd} = 1/\text{rms}(x)$ ):

d\gamma_i = \sum_{\text{slice}} \text{grad\_out}_i\,x_i\,\text{rstd}

and the input gradient is the standard RMSNorm backward formula (the centred grad is replaced by the un-centred grad since RMSNorm has no mean subtraction).

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

Two-element vector {dx, d_gamma}.