Assembling the Transformer Model

Welcome to Bringing Transformers to Life: Training & Inference! Great work in getting to this point. Over the past two courses, you've journeyed from traditional sequence models to mastering the intricate components of the Transformer architecture. In our first course, you explored RNNs, LSTMs, and discovered how attention mechanisms revolutionized sequence modeling. In our second course, you meticulously deconstructed every piece of the Transformer, from multi-head attention and positional encodings to complete encoder and decoder layers.

Now, we enter an exciting new phase where theory meets practice. This course transforms your architectural knowledge into working, trainable models that can tackle real-world tasks. As we begin this first lesson of our third course, we'll take the crucial step of assembling all those carefully crafted components into a unified, fully functional Transformer model. By the end of this lesson, you'll have built a complete implementation that bridges the gap between isolated components and production-ready architectures.

Let's begin constructing our complete Transformer model by establishing its foundational structure. Notice how we bring together all the components we've built separately:

class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, num_heads=8, 
                 num_encoder_layers=6, num_decoder_layers=6, d_ff=2048, 
                 max_seq_len=5000, dropout=0.1, activation='relu'):
        """Complete Transformer model for sequence-to-sequence tasks"""
        super(Transformer, self).__init__()
        
        self.d_model = d_model
        
        # Embeddings
        self.src_embedding = TokenEmbedding(src_vocab_size, d_model)
        self.tgt_embedding = TokenEmbedding(tgt_vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_seq_len, dropout)

This initialization reveals the Transformer's modular design philosophy. We maintain separate embeddings for source and target vocabularies because they often represent different languages or domains; imagine translating from English (with its 170,000+ words) to Chinese (with its character-based system). The d_model parameter acts as the universal dimensionality that ensures all components speak the same mathematical language. By setting max_seq_len to 5000, we pre-compute positional encodings for sequences up to this length, striking a balance between memory efficiency and practical flexibility for most real-world applications.

Now, we integrate the encoder and decoder stacks, along with the crucial output projection layer:

        # Encoder and Decoder
        self.encoder = TransformerEncoder(num_encoder_layers, d_model, num_heads, d_ff, dropout, activation)
        self.decoder = TransformerDecoder(num_decoder_layers, d_model, num_heads, d_ff, dropout, activation)
        
        # Output projection
        self.output_projection = nn.Linear(d_model, tgt_vocab_size)
        
        self._init_parameters()

The beauty of our modular approach shines here: the TransformerEncoder and TransformerDecoder classes we painstakingly built in previous lessons now slot in seamlessly. The output projection layer deserves special attention — it's essentially asking, "Given everything the decoder has learned about the sequence, what's the probability of each word in my target vocabulary being next?" This linear transformation from d_model dimensions to tgt_vocab_size dimensions creates the logits that, after softmax, become our predicted token probabilities.

Any Machine Learning model needs robust initialization and utility functions to handle real-world data complexities:

    def _init_parameters(self):
        """Initialize model parameters"""
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
    
    def create_padding_mask(self, seq, pad_token=0):
        """
        Create a padding mask for variable-length sequences.
        """
        # (B, S) -> (B, 1, 1, S) to broadcast over heads and query positions
        return (seq != pad_token).unsqueeze(1).unsqueeze(2)
    
    def create_causal_mask(self, size):
        """
        Create a causal (look-ahead) mask for decoder self-attention.
        """
        # (S, S) lower-triangular -> (1, 1, S, S) to broadcast over batch and heads
        mask = torch.tril(torch.ones(size, size))
        return mask.unsqueeze(0).unsqueeze(1).bool()

Here:

• Xavier uniform initialization ensures our model starts training from a stable point — without it, random initialization could lead to vanishing or exploding gradients that make training impossible.
• The create_padding_mask function is crucial for handling variable-length sequences: in a batch containing "Hello world" and "Hi," the shorter sequence gets padded, and this mask ensures the model ignores those meaningless pad tokens. It returns a boolean tensor of shape (batch_size, 1, 1, seq_len), which is broadcastable to attention score tensors of shape (batch_size, num_heads, query_len, key_len).
• The create_causal_mask generates a lower triangular matrix that enforces the decoder's autoregressive property, preventing it from "cheating" by looking at future tokens during training. It returns a boolean tensor of shape (1, 1, size, size), where size is the target length. In practice, the decoder self-attention mask is formed by combining the target padding mask (batch_size, 1, 1, size) with the causal mask (1, 1, size, size) via logical AND, yielding a mask of shape (batch_size, 1, size, size).

The forward pass reveals how all components collaborate to transform input sequences into output predictions:

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        """Forward pass through complete Transformer"""
        # Embed and encode source
        src_embedded = self.positional_encoding(self.src_embedding(src))
        encoder_output, _ = self.encoder(src_embedded, src_mask)
        
        # Embed and decode target
        tgt_embedded = self.positional_encoding(self.tgt_embedding(tgt))
        decoder_output, _ = self.decoder(tgt_embedded, encoder_output, tgt_mask, src_mask)
        
        # Project to vocabulary
        output = self.output_projection(decoder_output)
        return output

This elegantly simple forward pass conceals sophisticated interactions. Source tokens undergo a three-stage transformation: embedding lookup converts integers to vectors, positional encoding adds sequence order information, and the encoder stack builds rich contextual representations. The decoder follows a parallel path but with a crucial difference: it receives the encoder's output as additional context through cross-attention mechanisms. The final projection transforms the decoder's abstract representations back into concrete vocabulary predictions, completing the sequence-to-sequence pipeline.

Let's validate our implementation through systematic testing that demonstrates the model's capabilities. We'll begin by setting up the model configuration and verifying its basic structure:

def main():    
    # Configuration
    src_vocab_size = 1000
    tgt_vocab_size = 1200
    d_model = 128
    num_heads = 8
    num_encoder_layers = 2
    num_decoder_layers = 2
    d_ff = 512
    batch_size = 2
    src_seq_len = 10
    tgt_seq_len = 8
    
    # Create model
    model = Transformer(
        src_vocab_size=src_vocab_size,
        tgt_vocab_size=tgt_vocab_size,
        d_model=d_model,
        num_heads=num_heads,
        num_encoder_layers=num_encoder_layers,
        num_decoder_layers=num_decoder_layers,
        d_ff=d_ff
    )
    
    # Create sample inputs
    torch.manual_seed(42)
    src_tokens = torch.randint(1, src_vocab_size, (batch_size, src_seq_len))
    tgt_tokens = torch.randint(1, tgt_vocab_size, (batch_size, tgt_seq_len))
    
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
    print(f"Source tokens shape: {src_tokens.shape}")
    print(f"Target tokens shape: {tgt_tokens.shape}")

This configuration uses deliberately modest dimensions to enable quick verification while maintaining architectural authenticity. The different vocabulary sizes (1000 vs. 1200) demonstrate the model's flexibility, essential for tasks like translating between languages with vastly different vocabulary sizes:

Model parameters: 1,362,096
Source tokens shape: torch.Size([2, 10])
Target tokens shape: torch.Size([2, 8])

Now, let's execute the forward pass and verify that our model produces the expected outputs with proper gradient flow:

    # Create masks
    src_mask = model.create_padding_mask(src_tokens)
    tgt_mask = model.create_causal_mask(tgt_seq_len)
    
    # Forward pass
    output = model(src_tokens, tgt_tokens, src_mask, tgt_mask)
    
    print(f"Model output shape: {output.shape}")
    print(f"Expected shape: ({batch_size}, {tgt_seq_len}, {tgt_vocab_size})")
    
    # Test gradient flow
    loss = output.mean()
    loss.backward()
    
    print("Gradient verification:")
    print(f"  Source embedding grad: {model.src_embedding.embedding.weight.grad.norm():.6f}")
    print(f"  Output projection grad: {model.output_projection.weight.grad.norm():.6f}")

The gradient flow verification is particularly crucial — it confirms that gradients propagate seamlessly from the output projection layer all the way back to the source embeddings. This end-to-end gradient flow ensures every parameter can learn during training:

Model output shape: torch.Size([2, 8, 1200])
Expected shape: (2, 8, 1200)
Gradient verification:
  Source embedding grad: 0.050715
  Output projection grad: 0.268537

With over 1.3 million parameters, our model has substantial capacity while remaining computationally manageable for educational exploration.

You've successfully assembled a complete, production-ready Transformer model that elegantly integrates every component from our architectural journey. This achievement represents more than technical implementation: you now possess a deep understanding of how token embeddings, positional encodings, attention mechanisms, and projection layers collaborate to enable powerful sequence-to-sequence transformations. Your model stands ready to tackle any task, from machine translation to text summarization, from dialogue generation to code completion.

The practice exercises ahead will solidify your understanding through hands-on experimentation with this complete architecture. As we progress through this course, we'll explore how to prepare datasets, implement efficient training loops, and deploy these models for real-world inference, transforming your theoretical mastery into practical expertise that can solve meaningful problems.