Layer Normalization › Gradient Flow Through Layer Normalization

Understanding how gradients flow through layer normalization is essential for debugging training issues and understanding why normalization stabilizes training.

The backward pass through layer normalization is more complex than a simple element-wise operation because each output depends on all inputs through the mean and variance computation. When we change a single input $x_i$, it affects its own normalized output, the mean $\mu$, and the variance $\sigma^2$, which in turn affects every output element. This coupling makes the gradient computation more involved.

Given the loss $L$ and the upstream gradient $\frac{\partial L}{\partial y}$ (the gradient flowing back from later layers), we need to compute three gradients: $\frac{\partial L}{\partial x}$ for backpropagation, and $\frac{\partial L}{\partial \gamma}$ and $\frac{\partial L}{\partial \beta}$ for updating the learnable parameters.

The output of layer normalization is $y_i = \gamma_i \cdot \hat{x}_i + \beta_i$, where $\hat{x}_i$ is the normalized input. Since this is a simple affine transformation, the gradients follow directly from the chain rule.

For the scale parameter $\gamma_i$:

$$\frac{\partial L}{\partial \gamma_i} = \sum_{n} \frac{\partial L}{\partial y_{n,i}} \cdot \hat{x}_{n,i}$$

where:

• $\frac{\partial L}{\partial \gamma_i}$: the gradient of the loss with respect to the $i$-th scale parameter
• $n$: an index over all tokens (across batch and sequence dimensions)
• $\frac{\partial L}{\partial y_{n,i}}$: the upstream gradient for the $i$-th feature of the $n$-th token
• $\hat{x}_{n,i}$: the normalized value of the $i$-th feature for the $n$-th token

Intuitively, this sums up how much each token's normalized value contributed to the loss through this scale parameter.

For the shift parameter $\beta_i$:

$$\frac{\partial L}{\partial \beta_i} = \sum_{n} \frac{\partial L}{\partial y_{n,i}}$$

where:

• $\frac{\partial L}{\partial \beta_i}$: the gradient of the loss with respect to the $i$-th shift parameter

This is simply the sum of upstream gradients, since $\beta_i$ adds directly to the output.

The gradient with respect to input is more involved because $\hat{x}_i$ depends on $x_i$ in three ways: directly, through the mean $\mu$, and through the variance $\sigma^2$. Applying the chain rule carefully yields:

$$\frac{\partial L}{\partial x_i} = \frac{\gamma_i}{\sigma} \left( \frac{\partial L}{\partial y_i} - \frac{1}{d}\sum_{j=1}^{d}\frac{\partial L}{\partial y_j} - \frac{\hat{x}_i}{d}\sum_{j=1}^{d}\frac{\partial L}{\partial y_j}\hat{x}_j \right)$$

where:

• $\frac{\partial L}{\partial x_i}$: the gradient of the loss with respect to the $i$-th input element
• $\gamma_i$: the learned scale parameter for the $i$-th feature
• $\sigma = \sqrt{\sigma^2 + \epsilon}$: the standard deviation (with epsilon for stability)
• $d$: the feature dimension (number of elements in the input vector)
• $\hat{x}_i$: the normalized input, equal to $(x_i - \mu) / \sigma$

Let's break down the three terms inside the parentheses:

1. Direct contribution $\frac{\partial L}{\partial y_i}$: The gradient that would flow if normalization were a simple scaling operation.

2. Mean correction $-\frac{1}{d}\sum_{j=1}^{d}\frac{\partial L}{\partial y_j}$: Accounts for how changing $x_i$ affects $\mu$, which affects all outputs. This term subtracts the average gradient, centering the gradient distribution.

3. Variance correction $-\frac{\hat{x}_i}{d}\sum_{j=1}^{d}\frac{\partial L}{\partial y_j}\hat{x}_j$: Accounts for how changing $x_i$ affects $\sigma^2$, which scales all outputs. This term is proportional to the normalized value $\hat{x}_i$, meaning inputs far from the mean get larger corrections.

This formula reveals something important: the gradient for each input element depends on the gradients of all other elements through the mean and variance terms. This coupling helps distribute gradient information across features, which can improve training stability.

The mean correction term subtracts a constant from every gradient. The effect is similar to centering the gradients, preventing them from all shifting in the same direction. If the loss pushes all features to increase, the mean correction tempers this. This ensures the relative ordering of features matters more than their absolute magnitude. This is analogous to how the forward pass removes the mean from the input; the backward pass removes the mean from the gradient.

The variance correction term is more subtle. It is proportional to $\hat{x}_i$, the normalized value at position $i$. Features that are far from the normalized mean (large $|\hat{x}_i|$) receive a larger variance correction. This makes sense: those features have a larger influence on the variance computation, so their gradient needs a larger adjustment to account for that influence. In practice, this term prevents individual features from receiving gradients that are systematically larger simply because they happen to have larger normalized values.