Feed Forward Neural Network

Feed-Forward Neural Networks in Transformers

The Transformer model, introduced in the paper “Attention Is All You Need,” has revolutionized the field of natural language processing (NLP) and other sequence-to-sequence tasks. While its architecture primarily relies on multi-head self-attention mechanisms to capture contextual relationships, another critical component equally contributes to its success: the Feed-Forward Neural Network (FFNN), also known as the Position-wise Feed-Forward Network.

This FFNN sub-layer, present in both the encoder and decoder stacks of the Transformer, plays a vital role in transforming the representations learned by the attention mechanisms, enabling the model to capture more complex patterns within the input data sequence. This chapter explains the FFNN sub-layer, its role within the Transformer architecture, and provides a conceptual Python implementation.

The Feed-Forward Neural Network (FFNN) Sub-layer in Transformers

In the Transformer architecture, each encoder and decoder layer contains two main sub-layers: a multi-head self-attention mechanism and a position-wise feed-forward network. The FFNN sub-layer typically follows the attention sub-layer (and its associated Add & Norm step). Its input is the output of the preceding layer normalization step, which has a dimensionality of $d_{model}$ (e.g., $d_{model} = 512$).

A diagram illustrating a single block of the Transformer’s encoder would show the input first passing through a Multi-Head Attention sub-layer, followed by an Add & Norm (residual connection and layer normalization) step. The output of this Add & Norm layer then feeds into the Feed-Forward Network (FFNN) sub-layer. The FFNN itself is also followed by another Add & Norm step before producing the final output of that encoder/decoder block. This structure highlights the FFNN’s sequential position after the attention mechanism in processing information.

Key characteristics of the FFNN sub-layer:

• Fully Connected: It is a standard fully connected feed-forward network.
• Position-wise: The FFNN is applied to each position (i.e., each token’s representation in the sequence) separately and identically. This means that while the network is applied independently to each position, the same set of weights ($W_1, b_1, W_2, b_2$) is used across all positions in a given layer. This is different from a convolutional network that might share weights across local regions, or a recurrent network that processes tokens sequentially.
• Structure: It consists of two linear transformations with an activation function in between. The most commonly used activation function is ReLU (Rectified Linear Unit), though GELU (Gaussian Error Linear Unit) is also popular in modern Transformers.
• Dimensionality: The input and output dimensionality of the FFNN sub-layer is $d_{model}$ (e.g., 512). The inner layer typically has a larger dimensionality, often denoted as $d_{ff}$ (e.g., $d_{ff} = 2048$, which is $4 imes d_{model}$). This expansion and contraction allow the model to learn more complex transformations.
• Identical in Encoder and Decoder: Similar FFNN structures are used in both encoder and decoder layers, though their weights are learned independently.

Mathematical Representation of the FFNN Sub-layer

Let $x$ be the input to the FFNN sub-layer from a single position in the sequence (with dimension $d_{model}$). The FFNN operation can be mathematically represented as:

1. First Linear Transformation

$$ ext{Layer}_1(x) = x W_1 + b_1 $$

Here, $W_1$ is a weight matrix of shape $(d_{model}, d_{ff})$, and $b_1$ is a bias vector of shape $(d_{ff})$. The output of this layer has dimension $d_{ff}$.

2. Activation Function (e.g., ReLU)

The ReLU activation function introduces non-linearity, enabling the network to learn more complex patterns.

$$ ext{Activation}( ext{Layer}_1(x)) = max(0, x W_1 + b_1) $$

3. Second Linear Transformation

$$ ext{FFNN}(x) = (max(0, x W_1 + b_1)) W_2 + b_2 $$

Here, $W_2$ is a weight matrix of shape $(d_{ff}, d_{model})$, and $b_2$ is a bias vector of shape $(d_{model})$. The final output of the FFNN sub-layer for position $x$ has dimension $d_{model}$.

Importance of the FFNN Sub-layer in Transformers

The FFNN sub-layer plays a crucial role in the Transformer’s ability to model complex data:

Non-linear Transformation and Feature Enrichment

The self-attention mechanism primarily captures linear relationships and dependencies between different tokens in the sequence. The FFNN introduces non-linearity, allowing the model to learn more complex transformations of the token representations. The expansion to $d_{ff}$ and subsequent contraction back to $d_{model}$ can be seen as projecting the representations into a higher-dimensional space where more intricate features can be learned, and then projecting them back.

Position-wise Processing

While the self-attention mechanism considers relationships between different positions in the input sequence, the FFNN sub-layer processes each position’s representation independently. This allows the model to apply a consistent transformation to each token based on its individual features (as modified by the attention layer), further refining these representations. It acts as a per-token feature extractor/transformer.

Increased Model Capacity

The FFNN layers contribute significantly to the total number of parameters in the Transformer model. This increased capacity allows the model to learn from vast amounts of data and capture nuanced patterns necessary for high performance on complex NLP tasks.

Conceptual Python Implementation of a Feed-Forward Network

Below is a step-by-step conceptual Python implementation using NumPy to demonstrate how the Position-wise Feed-Forward Network sub-layer might operate within the Transformer architecture. Note that actual Transformer implementations use deep learning frameworks like PyTorch or TensorFlow for efficiency and automatic differentiation.

Step 1: Define the FeedForwardNN Class

We’ll define a class for our FFNN. In a Transformer, $d_{model}$ is the model’s embedding dimension (e.g., 512), and $d_{ff}$ is the inner-layer dimension (e.g., 2048).

import numpy as np

class FeedForwardNN:
    def __init__(self, d_model, d_ff):
        # Initialize weights and biases for the two linear transformations
        # Weights are often initialized with small random values (e.g., Xavier/Glorot initialization)
        # For simplicity, using np.random.randn scaled by 0.01
        self.W1 = np.random.randn(d_model, d_ff) * 0.01
        self.b1 = np.zeros((1, d_ff))  # Bias for the first layer

        self.W2 = np.random.randn(d_ff, d_model) * 0.01
        self.b2 = np.zeros((1, d_model)) # Bias for the second layer

    def relu(self, x):
        return np.maximum(0, x)

    def forward(self, x):
        # x shape: (batch_size, seq_len, d_model)
        # The FFNN is applied position-wise, so we can use batched matrix multiplication.

        # First linear transformation
        # (batch_size, seq_len, d_model) @ (d_model, d_ff) -> (batch_size, seq_len, d_ff)
        layer1_output = np.dot(x, self.W1) + self.b1

        # ReLU activation
        activation_output = self.relu(layer1_output)

        # Second linear transformation
        # (batch_size, seq_len, d_ff) @ (d_ff, d_model) -> (batch_size, seq_len, d_model)
        output = np.dot(activation_output, self.W2) + self.b2
        return output

Step 2: Create Example Input Data

In a Transformer, inputs are typically tensors of shape (batch_size, seq_len, d_model). Let’s create some example input data.

# Example input dimensions
batch_size = 2   # Number of sequences in a batch
seq_len = 10     # Length of each sequence (number of tokens)
d_model = 512    # Dimensionality of the model (input/output dimension of FFNN)
d_ff = 2048      # Dimensionality of the inner layer of FFNN

# Random input tensor with shape (batch_size, seq_len, d_model)
# This would be the output from the preceding (attention + add_norm) layer.
x_input = np.random.rand(batch_size, seq_len, d_model)

Step 3: Instantiate and Use the FFNN

Finally, we instantiate the FeedForwardNN class and perform a forward pass with our example input tensor.

# Create FFNN instance
ffnn = FeedForwardNN(d_model, d_ff)

# Perform forward pass
output = ffnn.forward(x_input)

# Print output shape (should be the same as input shape: (batch_size, seq_len, d_model))
print("Input shape:", x_input.shape)
print("Output shape:", output.shape)

Complete Implementation Example

Combining the above steps provides the complete conceptual implementation:

import numpy as np

class FeedForwardNN:
    def __init__(self, d_model, d_ff):
        self.W1 = np.random.randn(d_model, d_ff) * 0.01
        self.b1 = np.zeros((1, d_ff))
        self.W2 = np.random.randn(d_ff, d_model) * 0.01
        self.b2 = np.zeros((1, d_model))

    def relu(self, x):
        return np.maximum(0, x)

    def forward(self, x):
        layer1_output = np.dot(x, self.W1) + self.b1
        activation_output = self.relu(layer1_output)
        output = np.dot(activation_output, self.W2) + self.b2
        return output

# Example parameters
batch_size = 2
seq_len = 10
d_model = 512
d_ff = 2048    # Typically 4 * d_model

x_input = np.random.rand(batch_size, seq_len, d_model)

# Instantiate and run FFNN
ffnn = FeedForwardNN(d_model, d_ff)
output = ffnn.forward(x_input)

print("Input shape:", x_input.shape)
print("Output shape:", output.shape)

# For further learning, explore:
# - tf.keras.layers.Dense or torch.nn.Linear for actual implementations.
# - GELU activation function often used in modern Transformers.
# - Layer Normalization and Residual Connections around the FFNN.

Output

Running the above script will print the shape of the input and output tensors, verifying that the FFNN processes the data while maintaining the $d_{model}$ dimension at each position.

Input shape: (2, 10, 512)
Output shape: (2, 10, 512)

Conclusion

The Position-wise Feed-Forward Network (FFNN) is an essential component within each layer of the Transformer architecture. It enhances the model’s ability to capture complex patterns and relationships in the input data sequence by applying non-linear transformations independently to each token’s representation.

By complementing the global context aggregation of the self-attention mechanism with these local, position-wise transformations, the FFNN contributes significantly to the Transformer’s power and effectiveness in a wide range of NLP and sequence modeling tasks. Understanding its structure and function is key to fully appreciating the intricacies of the Transformer model.