Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Feed-Forward Networks › Implementation: A Complete FFN Module

Michael BrenndoerferUpdated June 10, 20253 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 107Feed-Forward Networks › Implementation: A Complete FFN Module

Reading section 9 of 13 · Implementation: A Complete FFN Module · 3 minSection 9 of 13 · 3 min

Section 9 of 13 · 3 min read

Feed-Forward Networks › Implementation: A Complete FFN ModuleOne section · Sections in this chapter9 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Having traced through the mathematics by hand, we can now build a reusable FFN module that encapsulates everything we've learned. This implementation follows patterns used in production transformer libraries: it initializes weights using proper scaling (Xavier/Glorot initialization), supports optional biases, and handles both single vectors and batched sequences without any special casing.

naturally the class structure maps onto the mathematical formula. The __init__ method creates the four learnable parameter tensors ( $W_1$ , $b_1$ , $W_2$ , $b_2$ ). The __call__ method implements the formula $\sigma(xW_1 + b_1)W_2 + b_2$ in three readable lines. The num_parameters method computes the formula we derived earlier. The entire forward pass is just three matrix multiplications (including the bias additions). This reflects the fundamental simplicity of the FFN despite its expressive power.

In[28]:

Code

class FeedForwardNetwork:
    """
    Position-wise feed-forward network for transformer blocks.

    Implements: FFN(x) = activation(x @ W1 + b1) @ W2 + b2
    """

    def __init__(self, d_model, d_ff, activation=relu, use_bias=True):
        """
        Initialize the feed-forward network.

        Args:
            d_model: Input and output dimension
            d_ff: Hidden dimension (typically 4 * d_model)
            activation: Nonlinear activation function
            use_bias: Whether to include bias terms
        """
        self.d_model = d_model
        self.d_ff = d_ff
        self.activation = activation
        self.use_bias = use_bias

        # Xavier/Glorot initialization
        self.W1 = np.random.randn(d_model, d_ff) * np.sqrt(
            2.0 / (d_model + d_ff)
        )
        self.W2 = np.random.randn(d_ff, d_model) * np.sqrt(
            2.0 / (d_ff + d_model)
        )

        if use_bias:
            self.b1 = np.zeros(d_ff)
            self.b2 = np.zeros(d_model)
        else:
            self.b1 = None
            self.b2 = None

    def __call__(self, x):
        """
        Apply the feed-forward transformation.

        Args:
            x: Input tensor of shape (..., d_model)

        Returns:
            Output tensor of shape (..., d_model)
        """
        # First linear layer
        hidden = x @ self.W1
        if self.use_bias:
            hidden = hidden + self.b1

        # Activation
        hidden = self.activation(hidden)

        # Second linear layer
        output = hidden @ self.W2
        if self.use_bias:
            output = output + self.b2

        return output

    def num_parameters(self):
        """Return total parameter count."""
        params = self.d_model * self.d_ff + self.d_ff * self.d_model
        if self.use_bias:
            params += self.d_ff + self.d_model
        return params

In[29]:

Code

# Test the module
ffn_module = FeedForwardNetwork(d_model=512, d_ff=2048)

# Single vector
x_single_test = np.random.randn(512)
y_single_test = ffn_module(x_single_test)

# Batch of vectors (sequence)
x_batch_test = np.random.randn(16, 512)  # 16 tokens
y_batch_test = ffn_module(x_batch_test)

Out[30]:

Console

FeedForwardNetwork module test:

Configuration:
  d_model: 512
  d_ff: 2048
  Parameters: 2,099,712

Single vector:
  Input shape: (512,)
  Output shape: (512,)

Batch (sequence):
  Input shape: (16, 512)
  Output shape: (16, 512)

The module correctly handles both single vectors and batched sequences. When a 2D input with shape $(n, d_{\text{model}})$ is provided (a batch of $n$ token representations), the matrix multiplication automatically broadcasts across all $n$ positions simultaneously. This is the source of the FFN's computational efficiency: instead of calling the function once per token in a loop, the entire sequence is processed in a single pair of batch matrix multiplications.

The parameter count matches our theoretical formula: $2 \times 512 \times 2048 + 2048 + 512 = 2{,}099{,}200 + 2{,}560 = 2{,}101{,}760$ , close to 2.1 million parameters just for this single FFN layer. A 12-layer GPT-2 Small model has 12 such layers, plus attention layers, adding up to the 117 million total parameters commonly cited.