The Transformer Block

A skyscraper is built from one floor design, repeated many times. Every transformer works the same way. There's one block, repeated. Twelve times for the first GPT models, and even more today. So if you understand a single floor, you understand the whole building.

• One floor, repeated
• Every transformer does this
• Understand one block

Inside one block

A block has two parts. First, self attention, where words share information with each other. Second, a small feed forward network, applied to each word by itself, where each word thinks about what it's just heard. Around each of these, there's a skip connection, the very same idea as ResNet. And layer norm keeps the numbers steady. Then you stack many of these blocks. Each one refines the understanding of every word.

Building the block

In code, it's beautifully short. Here's the attention, and here's the feed forward network, which expands to four times the width, then back again. Forward has two lines that matter. X plus attention. X plus the feed forward. Those plus signs are the skip connections. Normalising before each part is the modern style, called pre norm. Run it. A single block, and the shape is unchanged, so we can stack as many as we like.

import torch
from torch import nn

class Block(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(),
                                 nn.Linear(4 * dim, dim))
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x, mask=None):
        y = self.norm1(x)
        x = x + self.attn(y, y, y, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.norm2(x))

block = Block(dim=64, heads=8)
print("Weights in one block:", sum(p.numel() for p in block.parameters()))
print("Shape:", list(block(torch.rand(2, 10, 64)).shape))

Weights in one block: 49984
Shape: [2, 10, 64]

Quick check

Check yourself. In a transformer block, what's the job of the feed forward network? A, it mixes information between different words. B, it lets each word process what it's learned from attention. C, it adds position. D, it masks the future. The answer is B. Attention is where words talk to each other. The feed forward part is where each word thinks by itself.

Summary