Position-wise Feed-Forward Networks (FFN)

The Position-wise Feed-Forward Network (FFN) is an essential component found within each encoder and decoder layer. While attention mechanisms are designed to process relationships and integrate information across sequence positions, the FFN provides additional computational depth and non-linearity, operating independently on each position within the sequence.

Think of the attention sub-layer as the communication hub, allowing tokens to exchange information. The FFN then acts as a processing unit for each token individually, transforming the information received through attention.

Structure and Operation

The FFN is typically a simple fully connected feed-forward network consisting of two linear transformations with a non-linear activation function in between. The standard choice for the activation function in the original Transformer paper is the Rectified Linear Unit (ReLU).

The input to the FFN at a specific position $t$ is the output vector $z_t$ from the preceding sub-layer (either the self-attention or cross-attention, after the Add & Norm step, depending on the specific layer configuration like Post-LN or Pre-LN). The computation proceeds as follows:

1. First Linear Transformation: The input vector $z_t$ (with dimension $d_{model}$) is projected to a higher-dimensional space, typically $d_{ff}$. This expansion is achieved via a weight matrix $W_1$ (shape $d_{model} \times d_{ff}$) and a bias vector $b_1$ (shape $d_{ff}$). $$\text{Linear}_1(z_t) = z_t W_1 + b_1$$
2. Non-linear Activation: The result of the first transformation is passed through the ReLU activation function (or sometimes other activations like GeLU in more recent models). ReLU introduces non-linearity, allowing the network to learn more complex functions. $$\text{ReLU}(x) = \max(0, x)$$
3. Second Linear Transformation: The activated output is projected back down to the original model dimension $d_{model}$ using a second weight matrix $W_2$ (shape $d_{ff} \times d_{model}$) and bias vector $b_2$ (shape $d_{model}$). $$\text{Linear}_2(x) = x W_2 + b_2$$

Combining these steps, the complete FFN operation for a single position $t$ can be expressed as:

$$FFN(z_t) = \text{ReLU}(z_t W_1 + b_1) W_2 + b_2$$

Or, more generally using $f$ for the activation function:

$$FFN(z_t) = f(z_t W_1 + b_1) W_2 + b_2$$

Position-wise Application

A critical aspect of this network is the "position-wise" application. While the same FFN (meaning the exact same weight matrices $W_1, W_2$ and biases $b_1, b_2$) is used across all positions in the sequence, it is applied independently to the vector representation at each position.

If the input sequence representation after the attention sub-layer is $Z = [z_1, z_2, ..., z_n]$, where $n$ is the sequence length, the FFN computes:

$$FFN_{output} = [FFN(z_1), FFN(z_2), ..., FFN(z_n)]$$

This contrasts sharply with the attention mechanism, which explicitly models interactions between different positions. The FFN processes each position's representation in isolation, allowing the model to learn complex transformations of individual token representations informed by the context gathered via attention.

The diagram illustrates how the same two linear layers (with shared weights $W_1, b_1$ and $W_2, b_2$) and ReLU activation are applied independently to the input vector ($z_t$) at each sequence position.

Dimensionality and Implementation

In the original Transformer paper "Attention Is All You Need", the model dimension $d_{model}$ was 512, and the inner-layer dimension of the FFN, $d_{ff}$, was set to 2048. This fourfold expansion ($d_{ff} = 4 \times d_{model}$) is a common heuristic, although variations exist. The expansion allows the FFN to project the representation into a higher-dimensional space where complex patterns might be more easily learned, before projecting it back to the standard model dimension.

From an implementation perspective, applying the same linear transformations independently at each position is computationally efficient. It can be implemented using 1x1 convolutions. If you view the sequence of vectors $Z$ (shape: sequence length $\times$ $d_{model}$) as an "image" with height 1 and width equal to sequence length, then the linear transformations of the FFN correspond to convolutions with kernel size 1x1, input channels $d_{model}$, and output channels $d_{ff}$ (for the first layer) or $d_{model}$ (for the second layer). This formulation allows deep learning frameworks to leverage highly optimized convolution implementations for parallel processing across the sequence length.

The FFN sub-layer, along with the attention sub-layer, forms the core computational block within each encoder and decoder layer. Understanding its structure and position-wise operation is necessary for grasping how Transformers process information at both the sequence-interaction level (attention) and the individual token level (FFN). This sub-layer is followed by the residual connection and layer normalization, which we will discuss next.