Weight Initialization › BERT and GPT Initialization Schemes

Michael BrenndoerferApril 25, 20255 min read
Language AI HandbookAll chapters
Part 10 of 67Neural Network Foundations
Chapter 70Weight Initialization › BERT and GPT Initialization Schemes
Reading section 10 of 19 · BERT and GPT Initialization Schemes · 5 minSection 10 of 19 · 5 min

Book contents

Language AI Handbook67 parts · 513 chaptersView the whole book →
Active chapter 70 of 513

Section 10 of 19 · 5 min read

Weight Initialization › BERT and GPT Initialization SchemesOne section · Sections in this chapter10 / 19 ↓
Sections in this chapter19
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

Modern large language models take a more pragmatic approach to initialization than the theoretically derived schemes above. Both BERT and GPT use a simple, uniform strategy: normal distribution with a fixed small standard deviation, regardless of layer size. This works because layer normalization, applied throughout these architectures, normalizes activations during training and reduces the network.s sensitivity to the exact initialization scale.

The choice of a fixed standard deviation rather than a layer-size-dependent formula is a deliberate engineering decision. In transformer architectures, every attention layer and feed-forward layer is followed by layer normalization, which rescales activations to have unit variance at every step. This means the network continuously self-corrects its activation distribution, making the initialization less critical for preventing vanishing or exploding activations. The initialization instead needs to be "reasonable enough" that the network can start making productive gradient updates from the first iteration.

BERT Initialization

BERT uses N(0,0.02)\mathcal{N}(0, 0.02) for all weight matrices and zeros for all bias vectors. The 0.02 standard deviation was chosen empirically by the original BERT authors based on what produced stable training for their specific architecture. The layer normalization parameters (scale γ\gamma and bias β\beta) are initialized to 1 and 0, respectively, establishing an initial identity transformation.

The choice of 0.02 is smaller than what Xavier or He would prescribe for typical transformer layer dimensions. For a layer with nin=768n_{\text{in}} = 768 (the hidden dimension of BERT-base), He initialization gives σ=2/768≈0.051\sigma = \sqrt{2/768} \approx 0.051, and Xavier gives σ=2/(768+768)≈0.036\sigma = \sqrt{2/(768+768)} \approx 0.036. The BERT value of 0.02 is smaller. This conservatism is intentional: smaller initialization means smaller initial activations, and with layer normalization continuously rescaling, the network quickly adapts to useful activation scales regardless of the initialization.

GPT-2 Initialization

GPT-2 uses the same N(0,0.02)\mathcal{N}(0, 0.02) base distribution but introduces a residual scaling factor for the projection matrices at the end of each attention block and each feed-forward block. These weights are scaled by 1/nlayers1/\sqrt{n_{\text{layers}}}:

wproj∼N ⁣(0,0.02nlayers)w_{\text{proj}} \sim \mathcal{N}\!\left(0, \frac{0.02}{\sqrt{n_{\text{layers}}}}\right)

where:

  • nlayersn_{\text{layers}}: the total number of transformer layers in the model
  • 0.020.02: the base standard deviation shared with BERT
  • 1nlayers\frac{1}{\sqrt{n_{\text{layers}}}}: a residual scaling factor that prevents variance from accumulating across layers

This scaling addresses a specific problem with residual connections. In a residual network, each layer computes xℓ+1=xℓ+F(xℓ)x_{\ell+1} = x_\ell + F(x_\ell), adding the residual branch output F(xℓ)F(x_\ell) to the input. If each residual branch contributes variance σ2\sigma^2, then after LL layers the total variance grows to Lσ2L \sigma^2. Scaling the projection weights by 1/L1/\sqrt{L} reduces each branch's variance contribution to σ2/L\sigma^2/L, keeping the total variance constant at σ2\sigma^2.

This residual scaling trick has been adopted by many subsequent transformer architectures. Larger models like GPT-3 and various open-source LLMs use the same basic approach, sometimes with slightly different base standard deviations tuned to their specific architectures. The key insight transfers directly: when residual branches accumulate, scale the contributing weights to keep variance bounded.

In[21]:
Code
import torch.nn as nn


class BertStyleInit(nn.Module):
    """Demonstrates BERT-style initialization."""

    def __init__(self, d_model, n_layers):
        super().__init__()
        self.layers = nn.ModuleList(
            [nn.Linear(d_model, d_model) for _ in range(n_layers)]
        )
        self.layer_norms = nn.ModuleList(
            [nn.LayerNorm(d_model) for _ in range(n_layers)]
        )
        self._init_weights_bert()

    def _init_weights_bert(self):
        """BERT-style initialization: N(0, 0.02) for weights, 0 for biases."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

    def forward(self, x):
        for linear, norm in zip(self.layers, self.layer_norms):
            x = x + linear(norm(x))
        return x


class GPT2StyleInit(nn.Module):
    """Demonstrates GPT-2-style initialization with residual scaling."""

    def __init__(self, d_model, n_layers):
        super().__init__()
        self.n_layers = n_layers
        self.projections = nn.ModuleList(
            [nn.Linear(d_model, d_model) for _ in range(n_layers)]
        )
        self.layer_norms = nn.ModuleList(
            [nn.LayerNorm(d_model) for _ in range(n_layers)]
        )
        self._init_weights_gpt2()

    def _init_weights_gpt2(self):
        """GPT-2 initialization: N(0, 0.02) with 1/sqrt(n_layers) residual scaling."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)
        # Scale down output projection weights
        for proj in self.projections:
            proj.weight.data *= 1.0 / (self.n_layers**0.5)

    def forward(self, x):
        for proj, norm in zip(self.projections, self.layer_norms):
            x = x + proj(norm(x))
        return x


d_model = 256
n_layers = 12
Out[22]:
Console
BERT vs GPT-2 initialization comparison:
============================================================

BERT-style first layer weight statistics:
  Mean:  -0.000006
  Std:   0.0200
  Range: [-0.0845, 0.0882]

GPT-2-style first projection weight statistics (scaled by 1/sqrt(12)):
  Mean:  -0.000008
  Std:   0.005759  (expected: 0.02/sqrt(12) = 0.005774)
  Range: [-0.0232, 0.0236]

Output std after 12 layers:
  BERT:  1.5036
  GPT-2: 1.0453
Out[23]:
Visualization
Line plot showing output standard deviation across 12 transformer layers for BERT and GPT-2 initialization schemes.
Output standard deviation across 12 transformer layers for BERT-style and GPT-2-style initialization. BERT initialization causes output variance to grow as residual contributions accumulate across layers. GPT-2's residual scaling (dividing projection weights by the square root of the number of layers) compensates for this accumulation and keeps the output variance stable throughout the network depth.