Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

RMSNorm › Gradient Flow Through RMSNorm

Michael BrenndoerferUpdated June 7, 20257 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 105RMSNorm › Gradient Flow Through RMSNorm

Reading section 8 of 15 · Gradient Flow Through RMSNorm · 7 minSection 8 of 15 · 7 min

Section 8 of 15 · 7 min read

RMSNorm › Gradient Flow Through RMSNormOne section · Sections in this chapter8 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Understanding the backward pass helps us see why RMSNorm is computationally cheaper and how gradients propagate through the network. The gradient computation for RMSNorm is simpler than LayerNorm because it doesn't need to backpropagate through the mean computation. Deriving this backward pass from scratch is also a useful exercise in applying the chain rule to shared computations.

The core challenge in computing gradients for normalization layers is that every output element depends on every input element through the shared normalization constant. When you change input element $x_j$ , you change output element $y_j$ directly and also change the RMS value, which affects all output elements. This coupling through the shared denominator means the gradient of the loss with respect to $x_j$ has two terms: a direct term and a distributed term.

Consider the forward pass where each output element $y_i$ is computed as:

y_i = \gamma_i \cdot \frac{x_i}{\text{RMS}(\mathbf{x})}

where:

$y_i$ : the $i$ -th element of the output
$\gamma_i$ : the $i$ -th learnable scale parameter
$x_i$ : the $i$ -th element of the input
$\text{RMS}(\mathbf{x})$ : the root mean square computed over all input elements, shared by all output elements

To compute gradients, we need to determine how each input element $x_j$ affects each output element $y_i$ . This requires the partial derivative:

\frac{\partial y_i}{\partial x_j} = \gamma_i \cdot \frac{\partial}{\partial x_j}\left(\frac{x_i}{\text{RMS}(\mathbf{x})}\right)

For notational convenience, let $r = \text{RMS}(\mathbf{x}) = \sqrt{\frac{1}{d}\sum_k x_k^2}$ . We first need the derivative of $r$ with respect to $x_j$ . Using the chain rule on the square root and sum:

\frac{\partial r}{\partial x_j} = \frac{\partial}{\partial x_j}\left(\frac{1}{d}\sum_k x_k^2\right)^{1/2} = \frac{1}{2}\left(\frac{1}{d}\sum_k x_k^2\right)^{-1/2} \cdot \frac{2x_j}{d} = \frac{x_j}{d \cdot r}

where:

$\frac{\partial r}{\partial x_j}$ : how much the RMS changes when $x_j$ changes
$x_j$ : the $j$ -th input element (the one we're differentiating with respect to)
$d$ : the dimension of the input vector
$r$ : the RMS value (appears in the denominator because of the square root derivative)

Now we apply the quotient rule to $\frac{x_i}{r}$ . The quotient rule states that $\frac{d}{dx}\left(\frac{u}{v}\right) = \frac{u'v - uv'}{v^2}$ :

\frac{\partial}{\partial x_j}\left(\frac{x_i}{r}\right) = \frac{\delta_{ij} \cdot r - x_i \cdot \frac{x_j}{d \cdot r}}{r^2} = \frac{\delta_{ij}}{r} - \frac{x_i x_j}{d \cdot r^3}

where:

$\delta_{ij}$ : the Kronecker delta, equal to 1 if $i = j$ and 0 otherwise
The first term $\frac{\delta_{ij}}{r}$ : the direct effect when $i = j$ (changing $x_j$ directly changes the numerator $x_i$ )
The second term $\frac{x_i x_j}{d \cdot r^3}$ : the indirect effect through the RMS in the denominator (changing any $x_j$ affects the RMS, which affects all outputs)

To compute the gradient of a loss $L$ with respect to input $x_j$ , we apply the chain rule, summing over all output elements:

\frac{\partial L}{\partial x_j} = \sum_i \frac{\partial L}{\partial y_i} \cdot \gamma_i \cdot \left(\frac{\delta_{ij}}{r} - \frac{x_i x_j}{d \cdot r^3}\right)

The Kronecker delta $\delta_{ij}$ selects only the $i = j$ term from the sum for the first part, giving us:

\frac{\partial L}{\partial x_j} = \frac{\gamma_j}{r} \cdot \frac{\partial L}{\partial y_j} - \frac{x_j}{d \cdot r^3} \sum_i \gamma_i \cdot x_i \cdot \frac{\partial L}{\partial y_i}

where:

$\frac{\partial L}{\partial y_j}$ : the upstream gradient for output element $j$
The first term: the direct gradient path from $x_j$ through $y_j$
The second term: the indirect gradient path where $x_j$ affects all outputs via the shared RMS denominator
The sum $\sum_i \gamma_i \cdot x_i \cdot \frac{\partial L}{\partial y_i}$ : aggregates the indirect effects across all output dimensions, and can be computed once and reused for all $j$

The key efficiency advantage is that the sum $\sum_i \gamma_i \cdot x_i \cdot \frac{\partial L}{\partial y_i}$ can be computed once and reused for all $j$ . This makes the backward pass $O(d)$ in complexity, the same as the forward pass. For LayerNorm, the backward pass is more complex because it must also propagate through the mean subtraction, leading to additional terms in the gradient expression.

In[18]:

Code

def rmsnorm_backward(dout, x, gamma, eps=1e-6):
    """
    Backward pass for RMSNorm.

    Args:
        dout: Upstream gradient, shape (..., d)
        x: Original input, shape (..., d)
        gamma: Scale parameter, shape (d,)
        eps: Numerical stability constant

    Returns:
        dx: Gradient with respect to input
        dgamma: Gradient with respect to scale
    """
    # Compute RMS from forward pass
    mean_sq = np.mean(x**2, axis=-1, keepdims=True)
    rms = np.sqrt(mean_sq + eps)

    # Normalized input
    x_norm = x / rms

    # Gradient for gamma
    dgamma = np.sum(dout * x_norm, axis=tuple(range(dout.ndim - 1)))

    # Gradient for x
    d = x.shape[-1]

    # Term 1: direct gradient through division
    dx_norm = dout * gamma / rms

    # Term 2: gradient through RMS (chain rule)
    # d(1/rms)/dx_j = -x_j / (d * rms^3)
    sum_term = np.sum(dout * gamma * x, axis=-1, keepdims=True)
    dx_rms = -sum_term * x / (d * rms**3)

    dx = dx_norm + dx_rms

    return dx, dgamma

In[19]:

Code

# Verify gradients numerically
def numerical_gradient(f, x, h=1e-5):
    """Compute numerical gradient using central difference."""
    grad = np.zeros_like(x)
    it = np.nditer(x, flags=["multi_index"], op_flags=["readwrite"])
    while not it.finished:
        idx = it.multi_index
        old_val = x[idx]

        x[idx] = old_val + h
        fxph = f(x.copy())

        x[idx] = old_val - h
        fxmh = f(x.copy())

        grad[idx] = (fxph - fxmh) / (2 * h)
        x[idx] = old_val
        it.iternext()

    return grad


# Test setup
d_test = 8
x_test = np.random.randn(2, d_test)
gamma_test = np.random.randn(d_test)
dout_test = np.random.randn(2, d_test)


def forward_scalar(x):
    """Forward pass returning scalar loss."""
    rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + 1e-6)
    y = gamma_test * (x / rms)
    return np.sum(y * dout_test)


# Analytical gradient
dx_analytical, _ = rmsnorm_backward(dout_test, x_test.copy(), gamma_test)

# Numerical gradient
dx_numerical = numerical_gradient(forward_scalar, x_test.copy())

Out[20]:

Console

Gradient verification:
  Max relative error: 1.33e-10
  Gradient check passed

The gradient check confirms our analytical backward pass implementation is correct. A relative error below $10^{-4}$ indicates that the analytical and numerical gradients agree to high precision, validating both the mathematical derivation and the code implementation.

Out[21]:

Visualization

Heatmap of the Jacobian matrix for RMSNorm showing diagonal dominance with subtle off-diagonal coupling through the shared RMS normalization. — Jacobian matrix of RMSNorm showing the partial derivative of each output element with respect to each input element. The diagonal shows the direct gradient path where input j directly influences output j. The off-diagonal elements capture the indirect coupling through the shared RMS denominator: changing any input element changes the RMS, which affects every output element. This coupling is what makes RMSNorm gradients flow to all inputs, not just the corresponding one.

The Jacobian heatmap makes the gradient structure concrete. The strong diagonal shows that each input element has a direct gradient path to its corresponding output element. The subtle off-diagonal values reveal the coupling effect: every input element has a small indirect influence on every output element via the shared RMS denominator. This coupling ensures that even elements whose direct contribution to the loss is small still receive gradients when the overall RMS scale needs to change.