04 · Intro to Transformers ⚡

5 min read

The architecture that dethroned the LSTM and powers every large language model you've heard of. Its secret is a deceptively simple idea: let every token look directly at every other token.

1. The 'Why'

LSTMs (sub-module 03) carry memory through a sequence, but they have two stubborn limitations. First, they're sequential by nature — token t can't be processed until token t-1 is done — which makes them slow to train and unable to fully exploit the parallel hardware you learned to love in Module 01. Second, even with gating, information from far-back tokens has to survive a long chain of updates to influence the present, so very long-range dependencies remain hard. For years these were accepted costs of doing sequence modeling. Then, in 2017, the paper "Attention Is All You Need" proposed throwing out recurrence entirely, and the field changed overnight.

The Transformer's core innovation is self-attention: instead of passing memory step by step, every token directly computes how much it should pay attention to every other token in the sequence, and builds its new representation as a weighted blend of them all. A pronoun can attend straight to the noun it refers to, however far away, in a single step — and because every position is computed independently, the whole sequence processes in parallel. This combination of long-range reach and parallelism is why Transformers scale to enormous models and datasets in ways RNNs never could, and it's the foundation of GPT, BERT, and modern LLMs. This sub-module gives you the intuition and shows you PyTorch's built-in nn.Transformer building blocks — a first look, not a deep dive, but enough to demystify what's under the hood of today's AI.

2. Core Concepts

Self-attention (Query, Key, Value). Each token produces a query, a key, and a value vector. A token's output is a weighted sum of all tokens' values, where the weights come from matching its query against every key. In plain terms: each token asks "who here is relevant to me?" and blends in the answers.

Multi-head attention. Attention is run several times in parallel ("heads"), each learning to focus on different kinds of relationships (syntax, coreference, etc.), then combined. nn.MultiheadAttention implements this.

Positional encoding. Because attention has no inherent sense of order (it sees a set of tokens), Transformers add position information to the embeddings so the model knows token order — recovering what recurrence gave for free.

Encoder vs. decoder. The encoder builds rich representations of an input (great for classification/understanding — BERT-style). The decoder generates a sequence one token at a time using masked self-attention (GPT-style). nn.Transformer provides both; nn.TransformerEncoderLayer is the common starting block.

It still starts with embeddings. Everything from sub-modules 01–02 still applies: tokenize, embed, (add positions), then attention. Transformers replace the recurrent middle, not the text plumbing.

3. Code in Action

Self-attention with `nn.MultiheadAttention`

import torch
import torch.nn as nn

attn = nn.MultiheadAttention(embed_dim=16, num_heads=4, batch_first=True)

x = torch.randn(2, 5, 16)          # (batch=2, seq_len=5, embed_dim=16)
# Self-attention: query, key, value are all the same sequence
out, weights = attn(x, x, x)
print(out.shape)                   # torch.Size([2, 5, 16]) -> attended representations
print(weights.shape)               # torch.Size([2, 5, 5]) -> who attends to whom

A Transformer encoder block/stack

import torch
import torch.nn as nn

# One encoder layer = self-attention + feed-forward, with norms and residuals built in
layer = nn.TransformerEncoderLayer(
    d_model=16,          # embedding/model dimension
    nhead=4,             # number of attention heads
    dim_feedforward=64,  # size of the internal MLP
    batch_first=True,
)
encoder = nn.TransformerEncoder(layer, num_layers=2)   # stack 2 such layers

x = torch.randn(2, 5, 16)
encoded = encoder(x)
print(encoded.shape)     # torch.Size([2, 5, 16]) -> contextualized token vectors

Adding positional information (the missing order)

import torch
import torch.nn as nn

seq_len, d_model = 5, 16
token_embed = nn.Embedding(20, d_model, padding_idx=0)
pos_embed   = nn.Embedding(seq_len, d_model)            # a learnable position table

ids = torch.tensor([[3, 4, 6, 7, 0]])                   # (batch, seq_len)
positions = torch.arange(seq_len).unsqueeze(0)          # [[0,1,2,3,4]]
x = token_embed(ids) + pos_embed(positions)             # combine meaning + order
# `x` is now ready to feed into the TransformerEncoder above.
print(x.shape)                                          # torch.Size([1, 5, 16])

4. Common Pitfalls

Pitfall 1 — Forgetting positional encoding. Self-attention is order-agnostic: without position information, "dog bites man" and "man bites dog" look identical to the model. You must add positional encodings/embeddings to the token embeddings before the Transformer layers.

Pitfall 2 — Omitting the padding mask. Attention will happily attend to <pad> tokens, contaminating every position's output. Pass a key_padding_mask (or src_key_padding_mask) marking pad positions so they're ignored — the mask from sub-module 01 comes back here.

Pitfall 3 — Confusing encoder and decoder / their masks. An encoder sees the whole sequence; a decoder must use a causal mask so each position can't peek at future tokens (essential for generation). Using the wrong one — or no causal mask in a generative model — leaks the answer and breaks training. Match the block and mask to your task.

5. Further Reading & Watch List

⬅️ Prev: 03 · Recurrent Layers & LSTMs · 🏁 Module complete! Back to the module index or on to Module 09 — Ecosystem: Lightning.