Feed Forward Neural Network

Feed Forward Neural Network (FFNN)

A FFNN is a neural network in which information flows only in the forward direction, without recurrence or feedback connections (hence feed forward). In a Transformer, the FFNN is used as the position-wise transformation component of each Transformer block. For a representation of token: $x\in {\mathbb{R}}^d$ the conventional 2-layer FFNN transforms it as:

$$f(x)=W_2\sigma (W_{1}x+b_1)+b_2$$

$W_1,W_2\to$ learned weights
$b_1,b_2\to$ biases
$\sigma \to$ non-linear activation function (like ReLU, tanh, …) \

Dimensions transform like this: ${\mathbb{R}}^{d_{model}}\to {\mathbb{R}}^{ff}\to {\mathbb{R}}^{d_{model}}$ , where $d_{ff}>>d_{model}$. So, FFNN can temporarily move the token representation into a substantially larger feature space and then down-project them back to their original dimension.

Projecting the hidden representation into a higher-dimensional space provides additional representational capacity while the non-linear activation function allows the network to model functions/relationships that can’t be represented purely by a linear transformation.

FFNN is applied independently to each token. So for a sequence of $N$ tokens the same FFNN is applied $N$ times and there is no direct communication between different positions inside this sublayer (hence position wise).

Stacking linear transformations without an activation function would still result in an overall linear transformation. Activation function prevents this collapse and enables the model to learn more expressive mappings.

Implementation:

class FeedForwardBlock(nn.Module):
    
    def __init__(self, d_model: int, d_ff: int, dropout: float):
        super().__init__()
        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)
    
    def forward(self, x):
        ## (B,S,d) --> (B,S,d_ff) --> (B,S,d)
        return self.linear_2(self.dropout(torch.relu(self.linear_1(x))))

Gated Linear Units (GLUs)

FFNN applies the same transformation to every input and hence is static. GLU introduces an input-dependent gating mechanism the FFNN transformation adaptive to the input token representation.

General GLU:

$$GLU(x)=\underset{gate}{\underbrace{\sigma (xW+b)}}\odot \underset{candidatefeatures}{\underbrace{(xV+c)}}$$

$\odot$ denotes point-wise multiplication. Since $\sigma$ produces value in $(0,1)$, it can suppress or preserve different features to different magnitude. So one branch decides what information is produced while the other determines how much of that information should be allowed through.

Swish Gated Linear Unit - SwiGLU

$$\text{SwishGLU}={\text{Swish}}_{\beta }(xW+b)\odot (xV+c)$$ $${\text{Swish}}_{\beta }(x)=x\sigma (\beta x)=\frac{x}{1+e^{-\beta x}}$$

All GLU variants use 3 weight matrices $(W,V,W_o)$ instead of 2 as we saw in FFNN. To keep the overall FFNN parameter count comparable we reduce the intermediate hidden dimension by $2/3$ ($d_{ff}\to \frac{2}{3}{d}_{ff}$).

Resources:

• Why AI Spent 60 Years Hunting for the Right Activation