What is the FFN layer in a transformer and why is non-linearity between attention blocks needed?

The FFN layer (Feed-Forward Network) in a transformer is a two-layer fully connected neural network that is applied separately to each element of the sequence after the attention layer. Its task is to add additional nonlinear processing and enhance the expressive power of the model.

Nonlinearity between attention blocks is necessary for the model to approximate complex functions and dependencies. Attention is responsible for weighted summation of information across different positions, while the FFN with nonlinear activation (such as ReLU or GELU) helps the transformer model more complex patterns than just linear combinations.

An example structure of an FFN layer:

import torch
import torch.nn as nn

class FFN(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.activation = nn.GELU()
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.linear2(self.activation(self.linear1(x)))

Here, d_model is the dimension of input and output vectors, d_ff is the dimension of the inner layer, usually larger than d_model.

Thus, the FFN layer with nonlinearity provides deep nonlinear feature transformation, which improves training and the quality of the model.