Building a Transformer Encoder Block

Hello! Welcome back to our journey into Audio AI.

In the last lesson, we did a deep dive into the implementation of the multi-head self-attention mechanism, the computational heart of the Transformer. You built a CausalSelfAttention module in PyTorch, mastering how queries, keys, and values interact to allow parts of a sequence to communicate.

Today, we'll zoom out one level. An attention layer, while powerful, is just one part of a larger, repeatable structure. Our goal is to assemble the complete Transformer Encoder Block, the fundamental building block of models like BERT, wav2vec 2.0, and the encoder part of Whisper.

Your learning outcome for this lesson is to: Assemble a Transformer encoder block combining multi-head attention and a feed-forward network with residual connections.

By the end of this lesson, you will understand how attention (for communication) and feed-forward networks (for computation) work together, and why residual connections and layer normalization are the essential "glue" that allows us to stack these blocks into a deep, trainable network.

1. Anatomy of a Transformer Encoder Block

A Transformer encoder block is composed of two main sub-layers:

1. A Multi-Head Self-Attention Layer: This is the communication layer. As we saw, it allows every position in the sequence to gather information from every other position (in a non-causal setting).
2. A Position-wise Feed-Forward Network: This is the computation layer. It processes the output of the attention layer at each position independently.

Crucially, each of these sub-layers is wrapped with two additional operations: a residual connection and layer normalization. This is often referred to as "Add & Norm".

The data flow for a single block is:

1. Apply Multi-Head Self-Attention to the input.
2. Add the output of the attention layer to the original input (Residual Connection) and then normalize the result (Layer Normalization).
3. Pass the result through the Feed-Forward Network.
4. Add the output of the FFN to its input (another Residual Connection) and then normalize the result (Layer Normalization).

Let's break down the new components.

2. Position-wise Feed-Forward Network (FFN)

After the self-attention layer has aggregated information across the sequence, the feed-forward network provides an additional layer of processing for each position independently. You can think of self-attention as communication and the FFN as computation or "thinking" about the information that was just gathered.

It's a simple two-layer MLP (Multi-Layer Perceptron):

1. A linear layer that expands the embedding dimension (e.g., from d_model to 4 * d_model).
2. A ReLU activation function.
3. A linear layer that projects the dimension back down (e.g., from 4 * d_model to d_model).

This transformation is applied identically at each position in the sequence.

3. Residual Connections and Layer Normalization

These two techniques are critical for successfully training deep Transformer models.

Residual Connections ("Add")
A residual connection, or skip connection, simply takes the input to a layer (x) and adds it to the output of the layer (sublayer(x)). The operation is x + sublayer(x).

Why is this so important?

• Gradient Flow: In very deep networks, gradients can shrink as they are backpropagated through many layers (the vanishing gradient problem). The addition operation acts like a "gradient superhighway," allowing gradients to flow directly from the output back to the input, making optimization much more stable.
• Identity Mapping: It allows the network to easily learn an identity function. If a layer is not useful, the network can learn to make its output close to zero, and the input will pass through unchanged. This makes it easier to add more layers without harming performance.

Layer Normalization ("Norm")
Layer Normalization standardizes the features for each sequence element independently. It calculates the mean and variance across the embedding dimension (d_model) for a single token at a single position in the batch and uses them to normalize that token's feature vector.

This differs from Batch Normalization, which normalizes across the batch dimension. LayerNorm is preferred in Transformers because:

• Its calculations are independent of other examples in the batch, making it work well with variable sequence lengths and small batch sizes (which are common due to the high memory usage of Transformers).
• It helps stabilize the training dynamics by keeping the activations within a consistent range.

In modern Transformers, it's common to apply LayerNorm before the sub-layer (Pre-LN) rather than after (Post-LN, as in the original paper), as this tends to lead to more stable training. The flow becomes x + sublayer(norm(x)).

4. Assembling the EncoderBlock in PyTorch

Now that we understand all the individual pieces, we can assemble them into a complete EncoderBlock module.

The structure in PyTorch will be:

1. __init__:
   • An instance of a MultiHeadAttention module (like the one we studied in the last lesson).
   • An nn.Sequential module for the FeedForward network.
   • Two nn.LayerNorm modules.
   • Two nn.Dropout modules for regularization.
2. forward(x):
   • Pass x through the first Add & Norm wrapper around the self-attention layer. Note the Pre-LN formulation: x = x + dropout(self_attn(norm1(x))).
   • Pass the result through the second Add & Norm wrapper around the feed-forward network: x = x + dropout(ffn(norm2(x))).

The following tutorial provides a clean, well-commented implementation of this block.

To see this built live, the following video provides a clear walkthrough.

5. From a Single Block to a Full Encoder

The power of the Transformer architecture comes from stacking these encoder blocks on top of each other. The full encoder is simply a sequence of N identical EncoderBlocks.

The output of the first block becomes the input to the second, the output of the second becomes the input to the third, and so on. This deep hierarchy allows the model to build progressively more complex and abstract representations of the input sequence.

The implementation typically uses an nn.ModuleList to hold the stack of EncoderBlock instances.

import torch
import torch.nn as nn
from typing import List




# Assume MultiheadAttention and FeedForward classes are defined as discussed
# For simplicity, we'll use stubs here.

class MultiHeadAttention(nn.Module):
    def __init__(self, input_dim, embed_dim, num_heads):
        super().__init__()



        # In a real implementation, this would contain the logic from the previous lesson
        self.proj = nn.Linear(input_dim, embed_dim)
    def forward(self, x, mask=None):
        return self.proj(x)

class FeedForward(nn.Module):
    def __init__(self, input_dim, dim_feedforward):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, dim_feedforward),
            nn.ReLU(),
            nn.Linear(dim_feedforward, input_dim)
        )
    def forward(self, x):
        return self.net(x)




# --- This is our main class for this lesson ---
class TransformerEncoderBlock(nn.Module):
    def __init__(self, input_dim: int, num_heads: int, dim_feedforward: int, dropout: float = 0.1):
        """
        Args:
            input_dim: Dimensionality of the input (d_model).
            num_heads: Number of attention heads.
            dim_feedforward: Hidden dimension of the FFN.
            dropout: Dropout probability.
        """
        super().__init__()
        



        # Self-Attention sub-layer (using a simplified MultiHeadAttention for this example)
        self.self_attn = MultiHeadAttention(input_dim, input_dim, num_heads)
        



        # Feed-Forward sub-layer
        self.feed_forward = FeedForward(input_dim, dim_feedforward)
        



        # Normalization layers (for Pre-LN)
        self.norm1 = nn.LayerNorm(input_dim)
        self.norm2 = nn.LayerNorm(input_dim)
        



        # Dropout for regularization
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:



        # 1. Self-Attention sub-layer with Pre-LN
        x_norm1 = self.norm1(x)
        attn_output = self.self_attn(x_norm1, mask=mask)
        x = x + self.dropout1(attn_output) # Residual connection




        # 2. Feed-Forward sub-layer with Pre-LN
        x_norm2 = self.norm2(x)
        ff_output = self.feed_forward(x_norm2)
        x = x + self.dropout2(ff_output) # Residual connection
        
        return x




# --- Stacking the blocks to create the full encoder ---
class TransformerEncoder(nn.Module):
    def __init__(self, num_layers: int, **block_args):
        super().__init__()
        self.layers = nn.ModuleList([TransformerEncoderBlock(**block_args) for _ in range(num_layers)])

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        for layer in self.layers:
            x = layer(x, mask=mask)
        return x




# Example usage:
num_layers = 6
input_dim = 512
num_heads = 8
dim_feedforward = 2048 # 4 * input_dim
batch_size = 4
seq_len = 100

encoder = TransformerEncoder(num_layers, input_dim=input_dim, num_heads=num_heads, dim_feedforward=dim_feedforward)
dummy_input = torch.randn(batch_size, seq_len, input_dim)
output = encoder(dummy_input)

print(f"Input shape: {dummy_input.shape}")
print(f"Output shape: {output.shape}")



# Expected output:
# Input shape: torch.Size([4, 100, 512])
# Output shape: torch.Size([4, 100, 512])

This demonstrates that the encoder block and the full encoder stack are shape-preserving, a key property that allows them to be stacked deeply.

Conclusion

In this lesson, you have successfully assembled a complete Transformer encoder block. You've seen how the individual components we've discussed fit together to form a powerful, repeatable unit for sequence processing.

Key Takeaways:

• A Transformer encoder block consists of two sub-layers: multi-head self-attention and a position-wise feed-forward network.
• The attention layer handles communication across the sequence, while the FFN handles computation at each position.
• Residual connections (Add) are essential for stable gradient flow in deep stacks of these blocks.
• Layer normalization (Norm) stabilizes training dynamics and is typically applied before each sub-layer in modern implementations (Pre-LN).
• The full encoder is simply a stack of these identical blocks, created efficiently using nn.ModuleList.

Preview of the Next Lesson:
We've now mastered the encoder block. Next time, we'll turn our attention to the other half of the original Transformer architecture. We will assemble a Transformer decoder block, including masked self-attention, cross-attention, and a feed-forward network. This will introduce the concept of cross-attention, where a layer attends to the output of a different component (the encoder), a crucial mechanism for sequence-to-sequence tasks.