Layer Normalization › The Role of Learnable Parameters

Michael BrenndoerferUpdated June 11, 20254 min read
Language AI HandbookAll chapters
Part 15 of 67Transformer Blocks
Chapter 104Layer Normalization › The Role of Learnable Parameters
Reading section 5 of 16 · The Role of Learnable Parameters · 4 minSection 5 of 16 · 4 min

Book contents

Section 5 of 16 · 4 min read

Layer Normalization › The Role of Learnable ParametersOne section · Sections in this chapter5 / 16 ↓
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

After normalizing activations to zero mean and unit variance, we've effectively forced all features into a standardized distribution. But what if the network needs some features to have a larger spread, or to be centered around a non-zero value? The learnable parameters γ\gamma (scale) and β\beta (shift) solve this problem.

For each feature dimension ii, the final output is:

yi=γi⋅x^i+βiy_i = \gamma_i \cdot \hat{x}_i + \beta_i

where:

  • yiy_i: the final output for the ii-th feature
  • x^i\hat{x}_i: the normalized value (zero mean, unit variance)
  • γi\gamma_i: the learned scale for feature ii, which controls the spread of values
  • βi\beta_i: the learned shift for feature ii, which controls the center of the distribution

If the network learns γi=σoriginal\gamma_i = \sigma_{\text{original}} and βi=μoriginal\beta_i = \mu_{\text{original}}, it can completely undo the normalization and recover the original distribution. This means layer normalization can never hurt the network's representational capacity; in the worst case, it learns to bypass itself entirely. In practice, the network finds an intermediate setting that benefits from stable optimization while still representing the patterns it needs.

This "can always undo itself" property is worth examining more carefully. Suppose a particular layer in the network has learned to encode certain semantic properties through the absolute scale of its activations. Without layer normalization, those large activations pass through unchanged. With layer normalization, they get normalized down to unit variance. But then γ\gamma can scale them back up. The network does not lose information; it just processes it differently. The training process can discover, for instance, that one feature dimension should carry a signal with large variance (indicating strong activation) while another should operate at a smaller scale. The γ\gamma vector encodes exactly this per-feature preference.

The β\beta vector serves a complementary role. Downstream computations like softmax, activation functions, and biased projections all shift values around. If a particular layer normalization is followed by a ReLU, having outputs that start from a non-zero center may help more of the features remain active. The network can learn this through β\beta, shifting the normalized distribution to a mean that works well with the downstream computation.

In[12]:
Code
# Demonstrate how gamma and beta affect the output
gamma_custom = torch.tensor([2.0, 0.5, 1.5, 0.8, 1.0, 3.0, 0.3, 2.5])
beta_custom = torch.tensor([1.0, -1.0, 0.0, 2.0, -0.5, 0.5, 0.0, -2.0])

output_custom, _ = layer_norm_forward(x, gamma_custom, beta_custom)
Out[13]:
Console
Feature-wise comparison (first token):
Feature  gamma    beta     Output mean 
----------------------------------------
0        2.0      1.0      1.663       
1        0.5      -1.0     -0.945      
2        1.5      0.0      0.672       
3        0.8      2.0      2.192       
4        1.0      -0.5     -0.485      
5        3.0      0.5      -0.862      
6        0.3      0.0      -0.143      
7        2.5      -2.0     -2.539

The output distribution for each feature is controlled by its corresponding γ\gamma and β\beta values. Features with larger γ\gamma values have wider distributions, while β\beta shifts the center. This per-feature control lets different dimensions of the representation operate at different scales, which matters in transformers because attention heads and feature dimensions may need different dynamic ranges.

Out[14]:
Visualization
Bar chart comparing feature output means with identity parameters versus custom gamma and beta values.
Effect of learned gamma and beta parameters on output distributions. Each bar shows the mean output value for one feature. With identity parameters (gamma=1, beta=0), all features have mean near zero. With custom parameters, each feature has its own center (controlled by beta) and spread (controlled by gamma). This shows the per-dimension flexibility that layer normalization provides.

With identity parameters (gamma=1, beta=0), all feature means are near zero as expected. With custom parameters, each feature shifts to its corresponding beta value. This shows how the learnable parameters give each dimension independent control over its output distribution.