Transformers in PyTorch

Real-World Scenario

A company needs to classify customer support tickets into categories. Training a Transformer from scratch would take far more labeled tickets than they have, so the practical route is to fine-tune a pretrained model such as DistilBERT on the tickets they do have. Doing that well — and debugging it when predictions look wrong — is much easier once you have seen what attention, masks, and positions actually compute.

Scaled Dot-Product Attention by Hand

Each token’s query is compared with every token’s key; the softmax of those scores says how much of each token’s value to take.

import math
import torch
import torch.nn.functional as F

torch.manual_seed(0)
Q = torch.randn(1, 3, 4)            # (batch, tokens, d_k)
K = torch.randn(1, 3, 4)
V = torch.tensor([[[1., 0., 0., 0.], [0., 1., 0., 0.], [0., 0., 1., 0.]]])

scores = Q @ K.transpose(-2, -1) / math.sqrt(Q.size(-1))   # (1, 3, 3)
weights = scores.softmax(dim=-1)
out = weights @ V

print(weights.squeeze(0).round(decimals=3))
print(weights.sum(dim=-1))                                  # each row sums to 1
print(out.squeeze(0).round(decimals=3))
print(torch.allclose(out, F.scaled_dot_product_attention(Q, K, V), atol=1e-6))

$ python attention.py
tensor([[0.4880, 0.3510, 0.1610],
        [0.3950, 0.2930, 0.3120],
        [0.4540, 0.3870, 0.1580]])
tensor([[1., 1., 1.]])
tensor([[0.4880, 0.3510, 0.1610, 0.0000],
        [0.3950, 0.2930, 0.3120, 0.0000],
        [0.4540, 0.3870, 0.1580, 0.0000]])
True

Row i of the weights says how much token i attends to each of the three tokens, and every row sums to 1. The values were chosen as one-hot vectors, so the output is literally the attention weights copied into the first three columns — in general, each output row is a weighted average of the value rows. Dividing by √d_k keeps the scores from growing with the vector size, which would push the softmax towards one-hot. PyTorch’s fused F.scaled_dot_product_attention computes the same thing.

Padding Masks — and a NaN Trap

Padding tokens must not receive attention. Setting their scores to -inf before the softmax gives them exactly zero weight.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
x = torch.randn(1, 4, 8)                         # 4 tokens; the last two are padding
keep = torch.tensor([[True, True, False, False]])

scores = x @ x.transpose(-2, -1)
masked = scores.masked_fill(~keep[:, None, :], float("-inf"))
print(masked.softmax(dim=-1)[0].round(decimals=3))

nothing = torch.full((1, 4), float("-inf"))      # a row where every key is masked
print(nothing.softmax(dim=-1))

$ python masking.py
tensor([[0.9990, 0.0010, 0.0000, 0.0000],
        [0.0590, 0.9410, 0.0000, 0.0000],
        [0.3000, 0.7000, 0.0000, 0.0000],
        [0.3070, 0.6930, 0.0000, 0.0000]])
tensor([[nan, nan, nan, nan]])

Every row puts all its weight on the two real tokens, including the rows for the padding positions — those outputs are ignored later anyway. The last line is the trap: if every key in a row is masked, softmax computes 0/0 and returns NaN, which then spreads through every later layer. A batch that contains an all-padding sequence is enough to trigger it.

Multi-Head Attention

Multi-head attention splits the model dimension into several smaller heads that attend independently, then recombines them.

import torch
import torch.nn as nn

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.d_k = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)     # Q, K, V projections in one matmul
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, L, D = x.shape
        q, k, v = self.qkv(x).view(B, L, 3, self.h, self.d_k).permute(2, 0, 3, 1, 4)
        print("per-head q:", tuple(q.shape))           # (batch, heads, tokens, d_k)
        ctx = nn.functional.scaled_dot_product_attention(q, k, v)
        return self.out(ctx.transpose(1, 2).reshape(B, L, D))

mine = MultiHeadSelfAttention(d_model=128, n_heads=4)
print("output:    ", tuple(mine(torch.randn(8, 32, 128)).shape))

builtin = nn.MultiheadAttention(128, 4, batch_first=True)
count = lambda m: sum(p.numel() for p in m.parameters())
print("params:    ", count(mine), count(builtin))

$ python mha.py
per-head q: (8, 4, 32, 32)
output:     (8, 32, 128)
params:     66048 66048

With d_model=128 and 4 heads, each head works on 32-dimensional queries, keys, and values: (batch, heads, tokens, d_k). Computing Q, K, and V with one Linear(128, 384) and reshaping is equivalent to three separate projections and faster. The parameter count matches nn.MultiheadAttention exactly — 4 × 128² + 4 × 128 = 66,048 — because it is the same computation.

Attention Ignores Word Order

import torch
import torch.nn as nn

torch.manual_seed(0)
layer = nn.TransformerEncoderLayer(d_model=16, nhead=2, batch_first=True, dropout=0.0)
layer.eval()
emb = nn.Embedding(10, 16)

dog_bites_man = torch.tensor([[3, 7, 5]])
man_bites_dog = torch.tensor([[5, 7, 3]])

with torch.no_grad():
    a = layer(emb(dog_bites_man))
    b = layer(emb(man_bites_dog))
print(torch.allclose(a[0, 0], b[0, 2], atol=1e-6))       # token 3 gets the same vector
print(torch.allclose(a.mean(1), b.mean(1), atol=1e-6))   # pooled sentences are identical

$ python order_blind.py
True
True

“Dog bites man” and “man bites dog” contain the same tokens in a different order. A Transformer layer gives token 3 exactly the same output vector in both sentences, and averaging over the tokens produces the same sentence vector. Self-attention treats its input as a set; nothing in it knows which token came first.

Positions Make Order Learnable

To see what that costs, train the same small classifier on a task that depends on order — does token 2 appear before token 3? — with and without learned position embeddings. The model and data generator, saved as order_model.py:

import torch
import torch.nn as nn

class Classifier(nn.Module):
    def __init__(self, use_positions, vocab=12, d=32, max_len=8):
        super().__init__()
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(max_len, d) if use_positions else None
        layer = nn.TransformerEncoderLayer(d, nhead=4, dim_feedforward=64, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)
        self.head = nn.Linear(d, 2)

    def forward(self, ids):
        x = self.tok(ids)
        if self.pos is not None:
            x = x + self.pos(torch.arange(ids.size(1)))       # learned position vectors
        return self.head(self.encoder(x).mean(dim=1))

def make_data(n, g):
    ids = torch.stack([torch.randperm(10, generator=g)[:8] + 2 for _ in range(n)])
    first_a = (ids == 2).float().argmax(1) + 8 * (ids != 2).all(1)
    first_b = (ids == 3).float().argmax(1) + 8 * (ids != 3).all(1)
    return ids, (first_a < first_b).long()                   # 1 if token 2 precedes token 3

The training script trains one model of each kind on the same data:

import torch
import torch.nn as nn
from order_model import Classifier, make_data

g = torch.Generator().manual_seed(0)
(X, y), (Xt, yt) = make_data(4000, g), make_data(1000, g)
for use_positions in (False, True):
    torch.manual_seed(0)
    model = Classifier(use_positions)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    for epoch in range(15):
        for i in range(0, 4000, 64):
            opt.zero_grad()
            nn.functional.cross_entropy(model(X[i:i + 64]), y[i:i + 64]).backward()
            opt.step()
    model.eval()
    with torch.no_grad():
        acc = (model(Xt).argmax(1) == yt).float().mean().item()
    print(f"positions={use_positions!s:5}  test accuracy {acc:.3f}")

$ python order_task.py
positions=False  test accuracy 0.695
positions=True   test accuracy 1.000

With positions, the model solves the task perfectly. Without them it reaches about 69% — and that is not random luck but the best an order-blind model can do. Each sequence draws 8 of 10 possible tokens, so in about 38% of sequences token 2 or 3 is missing, and presence alone decides the label; in the remaining 62% the model can only guess. 0.378 + 0.622 × 0.5 ≈ 0.689. Adding self.pos(...) to the token embeddings is the entire fix.

The Hugging Face API

Hugging Face Transformers wraps pretrained models and their tokenizers behind a common API. This machine could not reach the Hugging Face Hub, so the setup script below builds a small local tokenizer and a small DistilBERT configuration in the same format that from_pretrained("distilbert-base-uncased") would download. Saved as tickets.py:

TEXTS = [
    "my payment failed and i was charged twice",
    "how do i reset my password",
    "the app crashes when i open it",
    "i need to cancel my subscription",
    "when will my order arrive",
]
LABELS = [0, 1, 2, 3, 4]
LABEL_NAMES = ["billing", "account", "technical", "subscription", "shipping"]

from pathlib import Path
from transformers import BertTokenizerFast, DistilBertConfig
from tickets import TEXTS

# Offline stand-in for AutoTokenizer.from_pretrained("distilbert-base-uncased")
words = sorted({w for t in TEXTS for w in t.split()})
Path("tiny").mkdir(exist_ok=True)
Path("tiny/vocab.txt").write_text("\n".join(["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + words))
tokenizer = BertTokenizerFast("tiny/vocab.txt")
tokenizer.save_pretrained("tiny")

config = DistilBertConfig(vocab_size=tokenizer.vocab_size, dim=64, n_layers=2, n_heads=2,
                          hidden_dim=128, num_labels=5)
config.save_pretrained("tiny")

batch = tokenizer(["how do i reset my password", "refund please"], padding=True)
print(tokenizer.convert_ids_to_tokens(batch["input_ids"][1]))
print(batch["attention_mask"])

$ python hf_setup.py
['[CLS]', '[UNK]', '[UNK]', '[SEP]', '[PAD]', '[PAD]', '[PAD]', '[PAD]']
[[1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 0, 0, 0, 0]]

The tokenizer adds [CLS] and [SEP] around every text and pads shorter texts in the batch with [PAD], while attention_mask marks which positions are real — the same padding mask as above. refund and please are not in this tiny vocabulary and become [UNK]. A real WordPiece vocabulary would split unfamiliar words into known pieces instead.

Training a sequence classifier is an ordinary PyTorch loop:

import torch
import transformers
from transformers import AutoModelForSequenceClassification, AutoConfig, AutoTokenizer
from tickets import TEXTS, LABELS

transformers.logging.disable_progress_bar()
tokenizer = AutoTokenizer.from_pretrained("tiny")
torch.manual_seed(0)
# With Hub access: AutoModelForSequenceClassification.from_pretrained(
#     "distilbert-base-uncased", num_labels=5)
model = AutoModelForSequenceClassification.from_config(AutoConfig.from_pretrained("tiny"))
print(f"{type(model).__name__}: {sum(p.numel() for p in model.parameters()):,} parameters")

batch = tokenizer(TEXTS, padding=True, return_tensors="pt")
labels = torch.tensor(LABELS)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

model.train()
for step in range(1, 61):
    out = model(**batch, labels=labels)           # passing labels makes the model compute the loss
    out.loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    if step % 20 == 0:
        print(f"step {step}  loss {out.loss.item():.4f}  logits {tuple(out.logits.shape)}")

model.save_pretrained("ticket_classifier")
tokenizer.save_pretrained("ticket_classifier")

$ python hf_finetune.py
DistilBertForSequenceClassification: 106,245 parameters
step 20  loss 1.5336  logits (5, 5)
step 40  loss 1.0760  logits (5, 5)
step 60  loss 0.4041  logits (5, 5)

from_config builds the architecture with random weights; from_pretrained — shown in the comment — would start from DistilBERT’s pretrained weights, and nothing else in the loop changes. Passing labels makes the model compute the cross-entropy loss itself and return it as out.loss, alongside out.logits. Pretrained fine-tuning uses a much smaller learning rate, typically 2e-5 to 5e-5, so as not to erase what pretraining learned; this randomly initialized model needs a larger one. save_pretrained writes the weights and config in the format from_pretrained reads back:

import torch
import transformers
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from tickets import LABEL_NAMES

transformers.logging.disable_progress_bar()
tokenizer = AutoTokenizer.from_pretrained("ticket_classifier")
model = AutoModelForSequenceClassification.from_pretrained("ticket_classifier").eval()

for text in ["i was charged twice", "the app crashes on open", "refund please"]:
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        probs = model(**inputs).logits.softmax(dim=-1)[0]
    unknown = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]).count("[UNK]")
    print(f"{text!r:28} -> {LABEL_NAMES[probs.argmax()]:12s} {probs.max():.2f}  ({unknown} unknown)")

$ python hf_predict.py
'i was charged twice'        -> billing      0.67  (0 unknown)
'the app crashes on open'    -> technical    0.65  (1 unknown)
'refund please'              -> billing      0.42  (2 unknown)

The model reloads from the directory without any class definitions. Its predictions follow the words it memorized: “charged” and “twice” point to billing, “app” and “crashes” to technical. “refund please” becomes [CLS], two [UNK]s the model never saw during training, and [SEP] — and it still picks “billing” (0.42, about twice the 0.2 of a uniform guess), an answer with no evidence behind it. A pretrained model with a full vocabulary, fine-tuned on hundreds of real tickets, is what fixes both problems.

Practice

import math
import torch
import torch.nn.functional as F

torch.manual_seed(0)
q, k, v = torch.randn(3, 1, 4, 8).unbind(0)          # 4 tokens

scores = q @ k.transpose(-2, -1) / math.sqrt(8)
future = torch.triu(torch.ones(4, 4, dtype=torch.bool), diagonal=1)
weights = scores.masked_fill(future, float("-inf")).softmax(-1)
print(weights[0].round(decimals=3))

print(torch.allclose(weights @ v, F.scaled_dot_product_attention(q, k, v, is_causal=True),
                     atol=1e-6))

$ python p1_causal.py
tensor([[1.0000, 0.0000, 0.0000, 0.0000],
        [0.4930, 0.5070, 0.0000, 0.0000],
        [0.6220, 0.0320, 0.3460, 0.0000],
        [0.2030, 0.6230, 0.1020, 0.0710]])
True

import torch

row = torch.full((1, 4), float("-inf"))
print(row.softmax(-1))

safe = torch.full((1, 4), torch.finfo(torch.float32).min)   # very negative, but finite
print(safe.softmax(-1))

$ python p2_safe_mask.py
tensor([[nan, nan, nan, nan]])
tensor([[0.2500, 0.2500, 0.2500, 0.2500]])

import torch
import torch.nn as nn

torch.manual_seed(0)
emb = nn.Embedding(20, 16, padding_idx=0)
layer = nn.TransformerEncoderLayer(16, nhead=2, batch_first=True, dropout=0.0).eval()

a = torch.tensor([[4, 9, 7, 0, 0]])
b = torch.tensor([[4, 9, 7, 13, 18]])            # same real tokens, junk where padding was
pad = torch.tensor([[False, False, False, True, True]])

with torch.no_grad():
    no_mask = torch.allclose(layer(emb(a))[0, :3], layer(emb(b))[0, :3], atol=1e-5)
    masked = torch.allclose(layer(emb(a), src_key_padding_mask=pad)[0, :3],
                            layer(emb(b), src_key_padding_mask=pad)[0, :3], atol=1e-5)
print("without mask, real tokens unaffected by padding:", no_mask)
print("with mask,    real tokens unaffected by padding:", masked)

$ python p3_padding_mask.py
without mask, real tokens unaffected by padding: False
with mask,    real tokens unaffected by padding: True