Feed-Forward Networks and AddNorm

Introduction

Welcome back to Deconstructing the Transformer Architecture! In our second lesson, we're diving into the other essential components that make each Transformer block so powerful: the Position-wise Feed-Forward Network and the critical Add & Norm operations. While our previous lesson explored Multi-Head Attention and how multiple attention heads can capture diverse relationships across different representation subspaces, today we'll discover the complementary mechanisms that complete the Transformer's computational prowess. While attention mechanisms handle the complex relationships between positions in a sequence, position-wise feed-forward networks provide the computational power to transform these attended representations. Think of attention as gathering the right information, and the feed-forward network as processing that information to extract meaningful patterns. Combined with residual connections and layer normalization, these components form the complete building blocks that enable Transformers to learn deep, stable representations. This lesson will guide you through implementing both the feed-forward networks and the Add & Norm components that make deep Transformer training possible.

Understanding Position-wise Feed-Forward Networks

\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2

Join the 1M+ learners on CodeSignal

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

Implementing the Position-wise Feed-Forward Network

Let's implement the core PositionwiseFeedForward network with configurable activation functions and proper weight initialization:

Python

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1, activation='relu'):
        """
        Position-wise Feed-Forward Network
        
        Args:
            d_model: Model dimension
            d_ff: Feed-forward dimension (typically 4 * d_model)
            dropout: Dropout rate
            activation: Activation function ('relu' or 'gelu')
        """
        super(PositionwiseFeedForward, self).__init__()
        
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        
        # Choose activation function
        if activation.lower() == 'relu':
            self.activation = F.relu
        elif activation.lower() == 'gelu':
            self.activation = F.gelu
        else:
            raise ValueError(f"Unsupported activation: {activation}")
        
        self._init_weights()

The constructor establishes the essential FFN architecture with two linear transformations. The first layer expands from d_model to d_ff dimensions, providing the network with increased representational capacity in the intermediate layer. The second layer projects back to d_model, ensuring dimensional consistency with the rest of the Transformer. The dropout layer provides regularization to prevent overfitting, and the configurable activation function allows experimentation with different non-linearities.

The choice of d_ff is crucial: it's typically set to $4 \times d_{model}$ in standard Transformer implementations. This 4x expansion provides sufficient capacity for complex transformations while maintaining computational efficiency. The expansion-contraction pattern creates a bottleneck architecture that forces the network to learn compressed, meaningful representations in the intermediate layer.

The choice of activation function significantly impacts the FFN's behavior and training dynamics. Let's complete the implementation with proper weight initialization:

Python

    def _init_weights(self):
        """Initialize weights"""
        nn.init.xavier_uniform_(self.linear1.weight)
        nn.init.xavier_uniform_(self.linear2.weight)
        nn.init.constant_(self.linear1.bias, 0)
        nn.init.constant_(self.linear2.bias, 0)
    
    def forward(self, x):
        """Forward pass through position-wise FFN"""
        x = self.linear1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.linear2(x)
        return x

The forward pass implements the classic feed-forward pattern: linear transformation, activation, dropout, and final linear projection. Xavier uniform initialization ensures stable gradient flow during training by setting initial weights based on the number of input and output connections. The bias terms are initialized to zero, which is a common practice for linear layers in deep networks.

The activation function choice between ReLU and GELU represents different approaches to non-linearity. ReLU provides simple, efficient computation with sparse activation patterns, while GELU offers smoother gradients and has shown superior performance in many Transformer applications. GELU (Gaussian Error Linear Unit) approximates the activation pattern $x \cdot \Phi(x)$ , where $\Phi(x)$ is the cumulative distribution function of the standard normal distribution, providing a more nuanced activation profile than the hard threshold of ReLU.

Residual Connections and Layer Normalization

The Add & Norm component implements two critical techniques that enable stable training of deep Transformer networks: residual connections and layer normalization:

Python

class AddNorm(nn.Module):
    def __init__(self, d_model, dropout=0.1):
        """
        Add & Norm layer: residual connection + layer normalization
        
        Args:
            d_model: Model dimension
            dropout: Dropout rate
        """
        super(AddNorm, self).__init__()
        self.layer_norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, sublayer_output):
        """
        Apply residual connection and layer normalization
        
        Args:
            x: Original input (residual connection)
            sublayer_output: Output from sublayer (attention or FFN)
        
        Returns:
            Normalized output after residual connection
        """
        # Add residual connection and apply dropout
        output = x + self.dropout(sublayer_output)
        
        # Apply layer normalization
        return self.layer_norm(output)

The AddNorm module encapsulates two fundamental techniques that solve different training challenges. Residual connections address the vanishing gradient problem by providing direct pathways for gradients to flow through the network. The mathematical operation $\text{output} = x + \text{sublayer}(x)$ ensures that even if the sublayer learns to output zeros, the original input can still pass through unchanged.

Layer normalization normalizes activations across the feature dimension for each individual example, stabilizing training dynamics and reducing internal covariate shift. Unlike batch normalization, layer normalization doesn't depend on batch statistics, making it more suitable for sequence processing where batch sizes might vary. The eps parameter is crucial for numerical stability, preventing division by zero when variance is small. While PyTorch's default is 1e-5, using 1e-6 provides better stability in deep networks and mixed precision training. The learnable parameters gamma and beta are initialized to 1 and 0, allowing the network to adapt normalization strength per feature dimension.

The combination of these techniques enables the training of very deep Transformer networks with hundreds of layers.

Combining Components in the Transformer Pattern

Now let's see how these components work together in the complete Transformer encoder layer pattern and examine the practical benefits:

Python

def test_transformer_block():
    """Test complete Transformer block pattern"""
    print("Testing Transformer Block Pattern...")
    
    batch_size, seq_len, d_model = 2, 8, 64
    num_heads, d_ff = 8, 256
    
    # Create components
    torch.manual_seed(42)
    x = torch.randn(batch_size, seq_len, d_model)
    
    mha = MultiHeadAttention(d_model, num_heads)
    ffn = PositionwiseFeedForward(d_model, d_ff)
    add_norm1 = AddNorm(d_model)
    add_norm2 = AddNorm(d_model)
    
    # Apply Transformer encoder layer pattern
    attn_output, _ = mha(x, x, x)
    x1 = add_norm1(x, attn_output)
    
    ffn_output = ffn(x1)
    x2 = add_norm2(x1, ffn_output)

def test_transformer_block():
    """Test complete Transformer block pattern"""
    print("Testing Transformer Block Pattern...")
    
    batch_size, seq_len, d_model = 2, 8, 64
    num_heads, d_ff = 8, 256
    
    # Create components
    torch.manual_seed(42)
    x = torch.randn(batch_size, seq_len, d_model)
    
    mha = MultiHeadAttention(d_model, num_heads)
    ffn = PositionwiseFeedForward(d_model, d_ff)
    add_norm1 = AddNorm(d_model)
    add_norm2 = AddNorm(d_model)
    
    # Apply Transformer encoder layer pattern
    attn_output, _ = mha(x, x, x)
    x1 = add_norm1(x, attn_output)
    
    ffn_output = ffn(x1)
    x2 = add_norm2(x1, ffn_output)

This implementation demonstrates the canonical Transformer encoder layer pattern: Multi-Head Attention followed by Add & Norm, then Position-wise Feed-Forward Network followed by another Add & Norm. Each sublayer (attention and FFN) is wrapped with a residual connection and layer normalization, creating the characteristic two-step pattern that defines each Transformer block.

The sequential application shows how information flows through the block: first, the attention mechanism gathers relevant information from across the sequence; then, the Add & Norm operation stabilizes and normalizes this output while preserving the original input through the residual connection. Next, the FFN processes this normalized representation, and finally, another Add & Norm operation stabilizes the output while again preserving information flow through the residual connection.

When we run the complete test suite, we get the following output that demonstrates the effectiveness of our implementation:

text

Testing Position-wise Feed-Forward Network...
Input shape: torch.Size([2, 8, 64])
FFN ReLU output shape: torch.Size([2, 8, 64])
FFN GELU output shape: torch.Size([2, 8, 64])
Testing Transformer Block Pattern...
Final output shape: torch.Size([2, 8, 64])
Input stats - mean: 0.0088, std: 1.0050
Output stats - mean: 0.0000, std: 1.0005

Testing Position-wise Feed-Forward Network...
Input shape: torch.Size([2, 8, 64])
FFN ReLU output shape: torch.Size([2, 8, 64])
FFN GELU output shape: torch.Size([2, 8, 64])
Testing Transformer Block Pattern...
Final output shape: torch.Size([2, 8, 64])
Input stats - mean: 0.0088, std: 1.0050
Output stats - mean: 0.0000, std: 1.0005

This output reveals several important characteristics of our implementation. First, the shape preservation confirms that both the FFN and the complete Transformer block maintain consistent dimensionality throughout processing. More importantly, the statistics comparison shows the effect of layer normalization: while the input has a small positive mean, the output is centered around zero, and the standard deviation remains close to 1.0, indicating that layer normalization is successfully stabilizing the activations and preventing activation drift that could destabilize training.

We've successfully implemented the complete PositionWiseFeed-Forward Network and Add & Norm components that form the other half of each Transformer block! These components work in harmony with the MultiHeadAttention mechanism you built previously to create powerful, trainable deep architectures. The FFN provides the non-linear computational power, while the Add & Norm operations ensure stable gradient flow and consistent activation magnitudes throughout the network.

The combination of these elements demonstrates the elegant engineering behind Transformers: attention mechanisms for relationship modeling, feed-forward networks for non-linear processing, and residual connections with layer normalization for training stability. In our next lesson, we'll explore how Transformers handle sequence order through positional encodings, completing our understanding of the core architectural components that make these models so effective across diverse NLP tasks. Now, let's get ready for some practice!

class PositionwiseFeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout=0.1, activation='relu'): """ Position-wise Feed-Forward Network Args: d_model: Model dimension d_ff: Feed-forward dimension (typically 4 * d_model) dropout: Dropout rate activation: Activation function ('relu' or 'gelu') """ super(PositionwiseFeedForward, self).__init__() self.linear1 = nn.Linear(d_model, d_ff) self.linear2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) # Choose activation function if activation.lower() == 'relu': self.activation = F.relu elif activation.lower() == 'gelu': self.activation = F.gelu else: raise ValueError(f"Unsupported activation: {activation}") self._init_weights()

def _init_weights(self): """Initialize weights""" nn.init.xavier_uniform_(self.linear1.weight) nn.init.xavier_uniform_(self.linear2.weight) nn.init.constant_(self.linear1.bias, 0) nn.init.constant_(self.linear2.bias, 0) def forward(self, x): """Forward pass through position-wise FFN""" x = self.linear1(x) x = self.activation(x) x = self.dropout(x) x = self.linear2(x) return x

class AddNorm(nn.Module): def __init__(self, d_model, dropout=0.1): """ Add & Norm layer: residual connection + layer normalization Args: d_model: Model dimension dropout: Dropout rate """ super(AddNorm, self).__init__() self.layer_norm = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, sublayer_output): """ Apply residual connection and layer normalization Args: x: Original input (residual connection) sublayer_output: Output from sublayer (attention or FFN) Returns: Normalized output after residual connection """ # Add residual connection and apply dropout output = x + self.dropout(sublayer_output) # Apply layer normalization return self.layer_norm(output)

Introduction

Understanding Position-wise Feed-Forward Networks

Join the 1M+ learners on CodeSignal

Implementing the Position-wise Feed-Forward Network

Activation Functions and Weight Initialization

Residual Connections and Layer Normalization

Combining Components in the Transformer Pattern

Output Discussion

Conclusion and Next Steps