xbe.at

LLMs and generative AI / LLM architecture and theory

LLM Initialization Techniques in Python

4 min read, 14 sections

Contents
  1. Introduction to LLM Initialization
  2. Xavier/Glorot Initialization
  3. He Initialization
  4. Orthogonal Initialization
  5. Scaled Initialization
  6. Embedding Layer Initialization
  7. Positional Encoding Initialization
  8. Layer Normalization Initialization
  9. Attention Weights Initialization
  10. Feed-Forward Network Initialization
  11. Custom Initialization Function
  12. Initialization with Pre-trained Embeddings
  13. Checking Initialization
  14. Additional Resources

Introduction to LLM Initialization

LLM initialization refers to the process of setting initial values for the model's parameters before training. Proper initialization is crucial for efficient training and better model performance.

import torch
import torch.nn as nn

class SimpleTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_layers)
        self.fc = nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer(x, x)
        return self.fc(x)

model = SimpleTransformer(vocab_size=10000, d_model=512, nhead=8, num_layers=6)

Xavier/Glorot Initialization

Xavier initialization is designed to keep the scale of gradients roughly the same in all layers. It's particularly useful for networks with symmetric activation functions.

def xavier_init(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

model.apply(xavier_init)

He Initialization

He initialization is similar to Xavier, but it's designed for ReLU activation functions. It helps prevent the vanishing gradient problem in deep networks.

def he_init(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
        if m.bias is not None:
            nn.init.zeros_(m.bias)

model.apply(he_init)

Orthogonal Initialization

Orthogonal initialization sets the weight matrices to be orthogonal, which can help with gradient flow in deep networks.

def orthogonal_init(m):
    if isinstance(m, nn.Linear):
        nn.init.orthogonal_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

model.apply(orthogonal_init)

Scaled Initialization

Scaled initialization adjusts the scale of initial weights based on the layer's size, which can be particularly useful for very deep networks.

def scaled_init(m):
    if isinstance(m, nn.Linear):
        in_dim = m.weight.size(1)
        std = 1 / math.sqrt(in_dim)
        nn.init.normal_(m.weight, mean=0, std=std)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

model.apply(scaled_init)

Embedding Layer Initialization

For LLMs, proper initialization of the embedding layer is crucial. Here's an example of initializing embeddings using a normal distribution.

def init_embedding(embedding_layer):
    nn.init.normal_(embedding_layer.weight, mean=0, std=0.02)

init_embedding(model.embedding)

Positional Encoding Initialization

Positional encodings help the model understand the order of tokens. Here's an example of sinusoidal positional encoding initialization.

def init_positional_encoding(d_model, max_len=5000):
    pe = torch.zeros(max_len, d_model)
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe

positional_encoding = init_positional_encoding(d_model=512)

Layer Normalization Initialization

Layer normalization is commonly used in transformers. Here's how to initialize its parameters.

def init_layer_norm(layer_norm):
    nn.init.ones_(layer_norm.weight)
    nn.init.zeros_(layer_norm.bias)

for module in model.modules():
    if isinstance(module, nn.LayerNorm):
        init_layer_norm(module)

Attention Weights Initialization

Proper initialization of attention weights can help with training stability. Here's an example for multi-head attention.

def init_attention_weights(attn_layer):
    nn.init.xavier_uniform_(attn_layer.in_proj_weight)
    if attn_layer.in_proj_bias is not None:
        nn.init.zeros_(attn_layer.in_proj_bias)
    nn.init.xavier_uniform_(attn_layer.out_proj.weight)
    if attn_layer.out_proj.bias is not None:
        nn.init.zeros_(attn_layer.out_proj.bias)

for module in model.modules():
    if isinstance(module, nn.MultiheadAttention):
        init_attention_weights(module)

Feed-Forward Network Initialization

The feed-forward networks in transformer layers also need proper initialization. Here's an example:

def init_feedforward(ffn):
    nn.init.xavier_uniform_(ffn.weight)
    if ffn.bias is not None:
        nn.init.zeros_(ffn.bias)

for module in model.modules():
    if isinstance(module, nn.Linear):
        init_feedforward(module)

Custom Initialization Function

Sometimes, you might want to combine different initialization techniques. Here's an example of a custom initialization function:

def custom_init(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Embedding):
        nn.init.normal_(m.weight, mean=0, std=0.02)

model.apply(custom_init)

Initialization with Pre-trained Embeddings

For transfer learning, you might want to initialize the embedding layer with pre-trained embeddings:

def load_pretrained_embeddings(embedding_layer, pretrained_embeddings):
    assert embedding_layer.num_embeddings == pretrained_embeddings.shape[0]
    assert embedding_layer.embedding_dim == pretrained_embeddings.shape[1]
    embedding_layer.weight.data._(torch.from_numpy(pretrained_embeddings))

# Assuming you have pretrained_embeddings as a numpy array
pretrained_embeddings = np.load('path_to_pretrained_embeddings.npy')
load_pretrained_embeddings(model.embedding, pretrained_embeddings)

Checking Initialization

After applying initialization techniques, it's important to verify that the weights are properly set:

def check_initialization(model):
    for name, param in model.named_parameters():
        if 'weight' in name:
            print(f"{name}: mean = {param.data.mean():.4f}, std = {param.data.std():.4f}")
        elif 'bias' in name:
            print(f"{name}: mean = {param.data.mean():.4f}")

check_initialization(model)

Additional Resources

  1. "Understanding the difficulty of training deep feedforward neural networks" by Xavier Glorot and Yoshua Bengio (http://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf)
  2. "Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification" by Kaiming He et al. (https://arxiv.org/abs/1502.01852)
  3. "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks" by Andrew M. Saxe et al. (https://arxiv.org/abs/1312.6120)
  4. "Improving Neural Networks by Preventing Co-adaptation of Feature Detectors" by Geoffrey E. Hinton et al. (https://arxiv.org/abs/1207.0580)

More in LLM architecture and theory

  1. Many-Shot vs Few-Shot In-Context Learning with Python9 min

    In-context learning (ICL) is a crucial aspect of Large Language Models (LLMs). This presentation explores the differences between many-shot and few-shot ICL...

  2. Knowledge Graphs as Powerful Evaluation Tools for LLM11 min

    Knowledge graphs are powerful tools for organizing and representing complex information. When combined with Large Language Models (LLMs), they can...

  3. Masked Grouped Causal Tracing in LLMs with Python10 min

    Masked Grouped Causal Tracing (MGCT) is an advanced technique for analyzing and interpreting the internal mechanisms of Large Language Models (LLMs). This...

  4. Introduction to Large Language Models12 min

    Tokenization is the fundamental process of converting raw text into numerical tokens that language models can process. This implementation demonstrates a...

  5. mPLUG-Owl3 AI Vision and Understanding10 min

    mPLUG-Owl3 represents a significant advancement in AI's ability to comprehend visual information. This model extends beyond simple image recognition, aiming...

  6. Improving Long Context LLMs with Writing in the Margins11 min

    Writing in the Margins (WiM) is an innovative inference pattern designed to address the Lost-in-the-Middle problem in long context Language Models (LLMs)....