PyTorch Transformer Model
The Transformer is the architecture behind GPT, BERT, T5, and virtually every state-of-the-art language model today. It replaced RNNs for sequence tasks by processing all positions in parallel using a mechanism called self-attention. PyTorch provides both low-level building blocks and high-level Transformer classes to build these models from scratch.
Why Transformers Beat RNNs
RNN: Processes tokens one at a time: word1 → word2 → word3 → ... Slow (sequential) — cannot parallelize across tokens Long-range memory degrades over many steps Transformer: Processes ALL tokens simultaneously using attention Fast — fully parallelizable across the sequence Every token can directly attend to every other token No vanishing gradient across sequence length
The Self-Attention Mechanism
Self-attention lets each position in a sequence look at every other position and decide how much to "attend" to it. Think of it as each word asking: "which other words in this sentence are most relevant to understanding me?"
Sentence: "The cat sat on the mat because it was tired" When processing "it": Attention to "cat" = HIGH (0.72) ← "it" refers to the cat Attention to "mat" = LOW (0.03) Attention to "was" = MED (0.18) ... Self-attention computes these weights automatically from the data.
Attention Computation
For each token, compute three vectors: Q (Query) = what this token is looking for K (Key) = what this token offers to others V (Value) = the actual content to retrieve Attention(Q, K, V) = softmax(Q·K^T / √d_k) · V Q·K^T: similarity between query and all keys / √d_k: scale to prevent exploding values softmax: convert similarities to probabilities (sum to 1) × V: weighted sum of value vectors
PyTorch Built-in Transformer Components
import torch
import torch.nn as nn
# Multi-Head Self-Attention
mha = nn.MultiheadAttention(
embed_dim=128, # size of each token embedding
num_heads=8, # number of attention heads
batch_first=True # input shape: [batch, seq, embed]
)
# Single Transformer Encoder Layer
enc_layer = nn.TransformerEncoderLayer(
d_model=128, # embedding dimension
nhead=8, # attention heads
dim_feedforward=512,# FFN hidden dimension
dropout=0.1,
batch_first=True
)
# Stack N encoder layers
encoder = nn.TransformerEncoder(enc_layer, num_layers=6)Building a Simple Text Classifier with Transformer
class TransformerClassifier(nn.Module):
def __init__(self, vocab_size, embed_dim, num_heads,
num_layers, num_classes, max_seq_len, dropout=0.1):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.pos_embed = nn.Embedding(max_seq_len, embed_dim)
enc_layer = nn.TransformerEncoderLayer(
d_model=embed_dim, nhead=num_heads,
dim_feedforward=embed_dim * 4,
dropout=dropout, batch_first=True
)
self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers)
self.classifier = nn.Linear(embed_dim, num_classes)
self.dropout = nn.Dropout(dropout)
def forward(self, x, src_key_padding_mask=None):
B, T = x.shape
positions = torch.arange(T, device=x.device).unsqueeze(0)
emb = self.embedding(x) + self.pos_embed(positions)
emb = self.dropout(emb)
out = self.encoder(emb, src_key_padding_mask=src_key_padding_mask)
cls_token = out[:, 0, :] # use first token's output for classification
return self.classifier(cls_token)
model = TransformerClassifier(
vocab_size=10000, embed_dim=128, num_heads=8,
num_layers=4, num_classes=5, max_seq_len=256
)
tokens = torch.randint(0, 10000, (16, 64)) # 16 sequences of 64 tokens
output = model(tokens)
print(output.shape) # torch.Size([16, 5])Positional Encoding
Unlike RNNs, Transformers process all tokens in parallel and have no built-in sense of order. Positional encoding injects position information into each token embedding so the model knows whether a word appears first, last, or in the middle.
Without positional encoding: "cat sat on mat" = "mat on sat cat" (same to the model) With positional encoding: Each position gets a unique signal added to its embedding Position 0 (cat) ≠ Position 3 (mat) even if same word
import math
class SinusoidalPositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=512, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(dropout)
pe = torch.zeros(max_len, d_model)
pos = torch.arange(max_len).unsqueeze(1).float()
div = torch.exp(torch.arange(0, d_model, 2).float()
* -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(pos * div) # even dims: sine
pe[:, 1::2] = torch.cos(pos * div) # odd dims: cosine
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
x = x + self.pe[:, :x.size(1)]
return self.dropout(x)Padding Mask
When batching sequences of different lengths, you pad shorter ones with zeros. Passing a padding mask tells the Transformer which positions are padding and should be ignored during attention.
tokens = torch.tensor([[1, 2, 3, 0, 0], # "hello world hi [PAD] [PAD]"
[4, 5, 6, 7, 0]]) # "this is a sentence [PAD]"
# True where padding (position should be ignored)
pad_mask = (tokens == 0) # shape: [2, 5]
# [[False, False, False, True, True ],
# [False, False, False, False, True ]]
output = model(tokens, src_key_padding_mask=pad_mask)Transformer Encoder vs Decoder
Encoder: Reads the input sequence Builds a rich contextual representation of each position Used in: BERT, text classification, NER, feature extraction Decoder: Generates output one token at a time Uses cross-attention to attend to encoder output Uses causal mask so each position only sees past tokens Used in: GPT, machine translation, summarization Full Encoder-Decoder: Used in: T5, BART, original "Attention Is All You Need" model
Using PyTorch's nn.Transformer
transformer = nn.Transformer(
d_model=128,
nhead=8,
num_encoder_layers=6,
num_decoder_layers=6,
dim_feedforward=512,
dropout=0.1,
batch_first=True
)
src = torch.rand(16, 20, 128) # encoder input: [B, src_seq, d_model]
tgt = torch.rand(16, 15, 128) # decoder input: [B, tgt_seq, d_model]
out = transformer(src, tgt)
print(out.shape) # torch.Size([16, 15, 128])Summary
The Transformer processes entire sequences in parallel using self-attention — each token attends to all others simultaneously. This eliminates the sequential bottleneck of RNNs and solves vanishing gradients over long sequences. The core components are multi-head attention, positional encoding, a feed-forward network per position, and layer normalization. Use nn.TransformerEncoder for classification and feature extraction tasks. Use nn.Transformer for sequence-to-sequence tasks. Pass a padding mask to ignore padded positions during attention. The Transformer architecture is the foundation of every modern large language model and vision transformer.