Building the Transformer Encoder

Introduction

Welcome to the fourth lesson in Deconstructing the Transformer Architecture! We've reached a pivotal moment in our journey where theory meets practice, and individual components unite to form something greater than the sum of their parts. Over the past three lessons, you've mastered the intricate mechanics of multi-head attention, understood the crucial role of feed-forward networks and Add & Norm operations, and discovered how positional encodings inject sequence awareness into these powerful models. Now comes the exciting challenge: assembling these components into a complete TransformerEncoderLayer . If our previous lessons were about understanding the instruments in an orchestra, today we're conducting the symphony. We'll explore not just what goes where, but why the specific ordering of operations matters, how data flows through the encoder, and what makes this architecture so remarkably effective at learning complex patterns. By the end of this lesson, you'll have built a fully functional encoder layer and will understand the elegant choreography that makes modern language models possible. Let's dive in!

The Architecture's Hidden Symphony

The Transformer encoder layer isn't just a stack of components: it's a carefully choreographed sequence of transformations that turns simple token embeddings into rich, contextually aware representations. Understanding this choreography requires appreciating a fundamental design principle: the encoder alternates between global reasoning (through attention) and local computation (through feed-forward networks), with each operation wrapped in mechanisms that ensure stable, efficient learning. Let's take a look at the architecture diagram (image taken from the original "Attention is all you need" paper): The magic happens through a precise order of operations: input embeddings first pass through multi-head self-attention, where each token can gather information from every other token in the sequence. This global information exchange is immediately followed by a residual connection and layer normalization, creating what we call the first sub-layer. The partially transformed representations then flow through position-wise feed-forward networks that process each position independently, applying non-linear transformations to extract higher-level features. Again, this is wrapped with residual connections and normalization as the second sub-layer. Why this specific order? Attention must come first because it establishes the contextual relationships that the feed-forward network will then process. The residual connections ensure that no information is lost during transformation — crucial for deep networks where gradients might otherwise vanish. Layer normalization after each sub-layer keeps the activations in a stable range, preventing the exploding or vanishing gradient problems that plagued earlier deep architectures. This design pattern, repeated across multiple encoder layers, creates a powerful learning machine capable of capturing everything from simple word associations to complex semantic relationships.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Building the Encoder Layer Foundation

Let's start constructing our TransformerEncoderLayer by establishing its core structure and understanding how our previously built components come together:

Python

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1, activation='relu'):
        """Transformer Encoder Layer"""
        super(TransformerEncoderLayer, self).__init__()
        
        # Multi-Head Self-Attention
        self.self_attention = MultiHeadAttention(d_model, num_heads, dropout)
        
        # Position-wise Feed-Forward Network
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout, activation)
        
        # Add & Norm layers
        self.add_norm1 = AddNorm(d_model, dropout)
        self.add_norm2 = AddNorm(d_model, dropout)

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1, activation='relu'):
        """Transformer Encoder Layer"""
        super(TransformerEncoderLayer, self).__init__()
        
        # Multi-Head Self-Attention
        self.self_attention = MultiHeadAttention(d_model, num_heads, dropout)
        
        # Position-wise Feed-Forward Network
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout, activation)
        
        # Add & Norm layers
        self.add_norm1 = AddNorm(d_model, dropout)
        self.add_norm2 = AddNorm(d_model, dropout)

This initialization reveals the encoder layer's elegant simplicity: just four components working in harmony. The MultiHeadAttention module we built earlier handles the complex task of computing relationships between all token pairs simultaneously. The PositionwiseFeedForward network provides the non-linear processing power that transforms these relationships into useful features. What's particularly important is that we create two separate AddNorm instances. Each maintains its own layer normalization parameters, allowing the model to learn different normalization strategies for attention outputs versus feed-forward outputs — a subtle but crucial detail that improves training stability and model performance.

The parameter choices here matter too: d_ff is typically set to four times d_model in standard Transformers, creating an expansion-compression pattern in the feed-forward network that has proven effective across many tasks. The activation parameter lets us choose between ReLU (traditional) and GELU (often preferred in modern models), while dropout provides regularization at multiple points throughout the layer.

Stacking Layers for Deep Understanding

Real Transformer power comes from stacking multiple encoder layers, each refining the representations from the previous layer:

Python

class TransformerEncoder(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, d_ff, dropout=0.1, activation='relu'):
        """Stack of Transformer Encoder Layers"""
        super(TransformerEncoder, self).__init__()
        
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout, activation)
            for _ in range(num_layers)
        ])
        
        self.num_layers = num_layers
    
    def forward(self, x, mask=None):
        """Forward pass through encoder stack"""
        attention_weights = []
        
        for layer in self.layers:
            x, attn_weights = layer(x, mask)
            attention_weights.append(attn_weights)
        
        return x, attention_weights

class TransformerEncoder(nn.Module):
    def __init__(self, num_layers, d_model, num_heads, d_ff, dropout=0.1, activation='relu'):
        """Stack of Transformer Encoder Layers"""
        super(TransformerEncoder, self).__init__()
        
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(d_model, num_heads, d_ff, dropout, activation)
            for _ in range(num_layers)
        ])
        
        self.num_layers = num_layers
    
    def forward(self, x, mask=None):
        """Forward pass through encoder stack"""
        attention_weights = []
        
        for layer in self.layers:
            x, attn_weights = layer(x, mask)
            attention_weights.append(attn_weights)
        
        return x, attention_weights

Using nn.ModuleList ensures PyTorch properly tracks all parameters across layers for optimization and device management. Each layer in the stack has its own set of parameters, allowing different layers to specialize in different aspects of language understanding. Research has shown that lower layers often focus on syntactic features and local patterns, while higher layers capture more semantic and task-specific information.

The design choice to collect attention weights from all layers provides a window into the model's reasoning process. Different layers often exhibit distinct attention patterns: early layers might focus on adjacent words and syntactic relationships, while later layers can capture long-range semantic dependencies. This hierarchical processing is what enables Transformers to understand language at multiple levels of abstraction simultaneously.

Testing the Individual Encoder Layer

Let's start by validating our TransformerEncoderLayer implementation to ensure it correctly handles the core transformation pipeline we've built:

Python

def test_encoder_layer():
    """Test single Transformer Encoder Layer"""
    print("Testing Transformer Encoder Layer...")
    
    # Configuration
    batch_size = 2
    seq_len = 10
    d_model = 64
    num_heads = 8
    d_ff = 256
    
    # Create sample input
    torch.manual_seed(42)
    x = torch.randn(batch_size, seq_len, d_model)
    
    # Create encoder layer
    encoder_layer = TransformerEncoderLayer(d_model, num_heads, d_ff)
    
    # Forward pass
    output, attention_weights = encoder_layer(x)
    
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    print(f"Attention weights shape: {attention_weights.shape}")
    
    # Verify dimensions preserved
    assert output.shape == x.shape
    
    # Test with padding mask
    padding_mask = torch.ones(batch_size, seq_len, seq_len)
    padding_mask[:, :, -2:] = 0  # Mask last 2 positions in the key dimension
    
    masked_output, _ = encoder_layer(x, padding_mask)
    print(f"Masked output shape: {masked_output.shape}")

def test_encoder_layer():
    """Test single Transformer Encoder Layer"""
    print("Testing Transformer Encoder Layer...")
    
    # Configuration
    batch_size = 2
    seq_len = 10
    d_model = 64
    num_heads = 8
    d_ff = 256
    
    # Create sample input
    torch.manual_seed(42)
    x = torch.randn(batch_size, seq_len, d_model)
    
    # Create encoder layer
    encoder_layer = TransformerEncoderLayer(d_model, num_heads, d_ff)
    
    # Forward pass
    output, attention_weights = encoder_layer(x)
    
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    print(f"Attention weights shape: {attention_weights.shape}")
    
    # Verify dimensions preserved
    assert output.shape == x.shape
    
    # Test with padding mask
    padding_mask = torch.ones(batch_size, seq_len, seq_len)
    padding_mask[:, :, -2:] = 0  # Mask last 2 positions in the key dimension
    
    masked_output, _ = encoder_layer(x, padding_mask)
    print(f"Masked output shape: {masked_output.shape}")

Running this test produces the following output:

text

Testing Transformer Encoder Layer...
Input shape: torch.Size([2, 10, 64])
Output shape: torch.Size([2, 10, 64])
Attention weights shape: torch.Size([2, 8, 10, 10])
Masked output shape: torch.Size([2, 10, 64])

Testing Transformer Encoder Layer...
Input shape: torch.Size([2, 10, 64])
Output shape: torch.Size([2, 10, 64])
Attention weights shape: torch.Size([2, 8, 10, 10])
Masked output shape: torch.Size([2, 10, 64])

These results validate several critical aspects of our encoder layer: the transformation preserves input dimensions while creating rich internal representations, attention weights have the correct shape for 8 heads attending across all 10 sequence positions, and the masking functionality works correctly for handling padding tokens.

Testing the Complete Encoder Stack

Now let's validate our full encoder implementation by testing it within a complete pipeline that mimics real-world usage:

Python

def test_complete_encoder():
    """Test complete encoder with embeddings"""
    print("Testing Complete Encoder with Embeddings...")
    
    # Configuration
    vocab_size = 1000
    num_layers = 3
    d_model = 64
    num_heads = 8
    d_ff = 256
    seq_len = 12
    batch_size = 2
    
    # Create sample token sequences
    torch.manual_seed(42)
    token_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    
    # Create components
    token_embedding = TokenEmbedding(vocab_size, d_model)
    positional_encoding = PositionalEncoding(d_model, dropout=0.1)
    encoder = TransformerEncoder(num_layers, d_model, num_heads, d_ff)
    
    # Process through complete pipeline
    embeddings = token_embedding(token_ids)
    embeddings_with_pos = positional_encoding(embeddings)
    encoder_output, attention_weights = encoder(embeddings_with_pos)
    
    print(f"Token IDs shape: {token_ids.shape}")
    print(f"Encoder output shape: {encoder_output.shape}")
    print(f"Number of attention weight tensors: {len(attention_weights)}")

def test_complete_encoder():
    """Test complete encoder with embeddings"""
    print("Testing Complete Encoder with Embeddings...")
    
    # Configuration
    vocab_size = 1000
    num_layers = 3
    d_model = 64
    num_heads = 8
    d_ff = 256
    seq_len = 12
    batch_size = 2
    
    # Create sample token sequences
    torch.manual_seed(42)
    token_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    
    # Create components
    token_embedding = TokenEmbedding(vocab_size, d_model)
    positional_encoding = PositionalEncoding(d_model, dropout=0.1)
    encoder = TransformerEncoder(num_layers, d_model, num_heads, d_ff)
    
    # Process through complete pipeline
    embeddings = token_embedding(token_ids)
    embeddings_with_pos = positional_encoding(embeddings)
    encoder_output, attention_weights = encoder(embeddings_with_pos)
    
    print(f"Token IDs shape: {token_ids.shape}")
    print(f"Encoder output shape: {encoder_output.shape}")
    print(f"Number of attention weight tensors: {len(attention_weights)}")

This comprehensive test produces the following output:

text

Testing Complete Encoder with Embeddings...
Token IDs shape: torch.Size([2, 12])
Encoder output shape: torch.Size([2, 12, 64])
Number of attention weight tensors: 3

Testing Complete Encoder with Embeddings...
Token IDs shape: torch.Size([2, 12])
Encoder output shape: torch.Size([2, 12, 64])
Number of attention weight tensors: 3

The results confirm that our complete encoder implementation works seamlessly: raw token IDs are successfully transformed through the entire pipeline into rich contextual representations, the three-layer encoder stack produces three sets of attention weights as expected, and the final output maintains the correct dimensions while containing the hierarchical transformations from all encoder layers. This validates that our architectural components integrate perfectly to create a functioning Transformer encoder ready for downstream tasks.

Introduction

The Architecture's Hidden Symphony

Join the 1M+ learners on CodeSignal

Building the Encoder Layer Foundation

Orchestrating the Forward Pass

Completing the Transformation Pipeline

Stacking Layers for Deep Understanding

Testing the Individual Encoder Layer

Testing the Complete Encoder Stack

Conclusion and Next Steps