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.

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.

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