Project Milestone: Custom Transformer Architecture Setup

Previously in this course, we explored individual building blocks like Implementing Multi-Head Attention: A Deep Dive into Transformers and Residual Connections and Gradient Stability in Deep Learning. Now, it’s time to move from isolated modules to a unified system.

In this Project Milestone, we define the blueprint for our custom Transformer. We will implement a structured configuration system, construct the base model wrapper, and apply rigorous weight initialization to ensure our training starts on stable ground.

Defining Model Configuration Parameters

In production, hard-coding hyperparameters is a recipe for technical debt. We need a centralized, serializable configuration schema. Using a dataclass is the standard for PyTorch projects; it provides type safety and makes it trivial to save/load model metadata alongside your weights.

PYTHON
from dataclasses import dataclass

@dataclass
class TransformerConfig:
    vocab_size: int = 50257
    n_layers: int = 12
    n_heads: int = 12
    n_embd: int = 768
    block_size: int = 1024
    dropout: float = 0.1
    layer_norm_eps: float = 1e-5
    # Add project-specific metadata
    model_name: str = "custom-transformer-v1"

By decoupling these parameters, you enable easy experiment tracking. You can pass a dictionary from a YAML file directly into this dataclass, ensuring your architecture remains flexible as we scale up.

Implementing the Base Transformer Class

Your base class should act as an orchestrator. It shouldn't contain the logic for attention or feed-forward networks—those should be imported as modular components. Instead, it manages the embedding lookup, the stack of Transformer blocks, and the final output projection.

PYTHON
import torch
import torch.nn as nn

class CustomTransformer(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        
        self.transformer = nn.ModuleDict({
            "wte": nn.Embedding(config.vocab_size, config.n_embd),
            "wpe": nn.Embedding(config.block_size, config.n_embd),
            "h": nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)]),
            "ln_f": nn.LayerNorm(config.n_embd, eps=config.layer_norm_eps),
        })
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        # Initialize weights
        self.apply(self._init_weights)

    def forward(self, idx):
        b, t = idx.size()
        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)
        
        x = self.transformer.wte(idx) + self.transformer.wpe(pos)
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        return self.lm_head(x)

Notice the use of nn.ModuleDict and nn.ModuleList. These ensure that PyTorch correctly registers all sub-parameters for optimization and checkpointing.

Weight Initialization and Normalization

Effective training depends heavily on how you set your initial weights. Per our earlier discussion on Advanced Weight Initialization Strategies for Deep Learning, we must apply specific scaling factors to ensure signal variance remains stable across deep architectures.

PYTHON
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            # Standard initialization with gain for residual paths
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.LayerNorm):
            torch.nn.init.zeros_(module.bias)
            torch.nn.init.ones_(module.weight)

Always initialize your LayerNorm weights to 1 and biases to 0; this provides the layer with an identity-mapping starting point, which prevents the network from collapsing during the first few training steps.

Hands-on Exercise

1. Extend the Config: Add a tie_weights boolean to your TransformerConfig. If true, the wte (embedding) and lm_head should share the same weight matrix to reduce parameter count.
2. Implement Weight Tying: In your CustomTransformer.__init__, if config.tie_weights is true, set self.lm_head.weight = self.transformer.wte.weight.
3. Verification: Print model.named_parameters() and verify that the number of parameters decreases significantly when tie_weights is enabled.

Common Pitfalls

• Forgetting to call apply(): If you define an _init_weights method but don't call self.apply(self._init_weights) in your constructor, PyTorch will use default Kaiming initialization, which is often suboptimal for Transformers.
• Device Mismatch: Always ensure your positional embeddings are registered to the same device as the input tensors. Using torch.arange(..., device=idx.device) is a clean way to handle this without explicit .to(device) calls.
• LayerNorm Placement: Ensure you are using the correct normalization technique. If you're building a modern architecture, consider replacing standard LayerNorm with RMSNorm as discussed in Normalization Techniques at Scale: Implementing RMSNorm.

Recap

We have successfully scaffolded our production project by defining a clean configuration schema, assembling the Transformer backbone, and applying robust initialization. This structure allows us to iterate on model depth and width without modifying the core logic.

Up next: Tokenization Strategies for LLMs — where we will build the vocabulary and text-processing pipeline required to feed data into this architecture.