Transformer Block

The core component of an LLM is the transformer block. We’ve already encountered all of the main components of a transformer block, namely self-attention and the feed-forward network:

There are a couple of extra components here: LayerNorm and residual connections.

LayerNorm

This is a scaling mechanism which keeps the numbers well-behaved as they pass through the many layers of the LLM. It stops the numbers from blowing up unpredictably and helps to prevent instability while training the model.

It works by subtracting the mean of the vector and dividing by the standard deviation.

There is a built-in PyTorch implementation for this:

import torch.nn as nn


layer_norm = nn.LayerNorm(model_dim)
normalised_x = layer_norm(x)

Residual connections

After each stage, a residual connection carries the stage's input around the side and adds it back to the result. This preserves useful existing information while the stage adds what it has learned. During training, the same shortcuts also give the learning signal (called the gradient) a clearer route back through the model, helping even the earliest blocks learn effectively.

Without residual connections our transformer block would look like this:

import torch.nn as nn


class TransformerBlockWithoutResidualConnections(nn.Module):

    def __init__(self, model_dim, head_dim, num_heads, hidden_dim):
        super().__init__()
        self.attention_norm = nn.LayerNorm(model_dim)
        self.attention = MultiHeadAttention(model_dim, head_dim, num_heads)
        self.feed_forward_norm = nn.LayerNorm(model_dim)
        self.feed_forward = FeedForwardLayer(model_dim, hidden_dim)

    def forward(self, x):
        x = self.attention(self.attention_norm(x))
        x = self.feed_forward(self.feed_forward_norm(x))
        return x

With residual connections we get our implementation of our transformer:

import torch.nn as nn


class TransformerBlock(nn.Module):

    def __init__(self, model_dim, head_dim, num_heads, hidden_dim):
        super().__init__()
        self.attention_norm = nn.LayerNorm(model_dim)
        self.attention = MultiHeadAttention(model_dim, head_dim, num_heads)
        self.feed_forward_norm = nn.LayerNorm(model_dim)
        self.feed_forward = FeedForwardLayer(model_dim, hidden_dim)

    def forward(self, x):
        x = x + self.attention(self.attention_norm(x))
        x = x + self.feed_forward(self.feed_forward_norm(x))
        return x