Transformers and Self-Attention

Imagine a meeting with 10 people. An LSTM-style meeting: person 1 speaks, whispers to person 2, person 2 whispers to person 3 — by person 10, the original message is distorted. A self-attention meeting: every person simultaneously reads every other person's written statement and decides how much to pay attention to each one when forming their own response. No information degrades. No sequential bottleneck.

The attention score between two tokens is their relevance — how much should token A look at token B when computing its contextual meaning? "Bank" should look at "river" with high attention weight. "Bank" should look at "steep" with lower weight. These weights are learned during training.

import numpy as np

# ── Self-attention from scratch — every operation visible ─────────────

def softmax(x, axis=-1):
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True)

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (seq_len, d_k)
    K: (seq_len, d_k)
    V: (seq_len, d_v)
    Returns: output (seq_len, d_v), attention_weights (seq_len, seq_len)
    """
    d_k = Q.shape[-1]

    # Step 1: Compute raw attention scores — how much does each Q match each K?
    scores = Q @ K.T / np.sqrt(d_k)   # (seq_len, seq_len)

    # Step 2: Apply mask (optional — for causal/decoder attention)
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)

    # Step 3: Softmax — convert scores to weights (each row sums to 1)
    weights = softmax(scores, axis=-1)  # (seq_len, seq_len)

    # Step 4: Weighted sum of values
    output = weights @ V                # (seq_len, d_v)

    return output, weights

# ── Concrete example: 5-token sentence ────────────────────────────────
np.random.seed(42)
seq_len, d_model, d_k = 5, 8, 4

# Token embeddings (in practice these come from an embedding table)
X = np.random.randn(seq_len, d_model)

# Projection matrices — learned during training
Wq = np.random.randn(d_model, d_k) * 0.1
Wk = np.random.randn(d_model, d_k) * 0.1
Wv = np.random.randn(d_model, d_k) * 0.1

# Project to Q, K, V
Q = X @ Wq   # (5, 4)
K = X @ Wk   # (5, 4)
V = X @ Wv   # (5, 4)

output, attn_weights = scaled_dot_product_attention(Q, K, V)

print(f"Input X shape:    {X.shape}")
print(f"Q, K, V shapes:   {Q.shape}")
print(f"Output shape:     {output.shape}")
print(f"Weights shape:    {attn_weights.shape}")
print(f"
Attention weights (each row sums to 1.0):")
print(attn_weights.round(3))
print(f"
Row sums: {attn_weights.sum(axis=1).round(6)}  ← all 1.0")

# ── Why √dₖ scaling matters ───────────────────────────────────────────
print("
Effect of scaling on softmax:")
raw_scores   = np.array([1.0, 2.0, 3.0, 4.0])
unscaled     = softmax(raw_scores)
scaled_4     = softmax(raw_scores / np.sqrt(4))
scaled_64    = softmax(raw_scores / np.sqrt(64))

print(f"No scaling  (d_k=1):  {unscaled.round(4)}  max={unscaled.max():.4f}")
print(f"Scale √4:             {scaled_4.round(4)}  max={scaled_4.max():.4f}")
print(f"Scale √64:            {scaled_64.round(4)}  max={scaled_64.max():.4f}")
print("Large d_k → large dot products → softmax saturates → vanishing gradients")
print("Dividing by √dₖ keeps variance stable regardless of d_k")

import numpy as np
import torch
import torch.nn as nn

# ── Multi-head attention from scratch ─────────────────────────────────
class MultiHeadAttentionScratch:
    def __init__(self, d_model, n_heads):
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k     = d_model // n_heads

        np.random.seed(42)
        scale = np.sqrt(1.0 / d_model)
        # One projection matrix per head per role
        self.Wq = [np.random.randn(d_model, self.d_k) * scale for _ in range(n_heads)]
        self.Wk = [np.random.randn(d_model, self.d_k) * scale for _ in range(n_heads)]
        self.Wv = [np.random.randn(d_model, self.d_k) * scale for _ in range(n_heads)]
        self.Wo = np.random.randn(d_model, d_model) * scale   # output projection

    def softmax(self, x):
        e = np.exp(x - x.max(axis=-1, keepdims=True))
        return e / e.sum(axis=-1, keepdims=True)

    def attention(self, Q, K, V):
        scores  = Q @ K.T / np.sqrt(self.d_k)
        weights = self.softmax(scores)
        return weights @ V, weights

    def forward(self, X):
        heads = []
        all_weights = []
        for i in range(self.n_heads):
            Q = X @ self.Wq[i]
            K = X @ self.Wk[i]
            V = X @ self.Wv[i]
            head_out, weights = self.attention(Q, K, V)
            heads.append(head_out)
            all_weights.append(weights)

        # Concatenate all heads → (seq_len, d_model)
        concat = np.concatenate(heads, axis=-1)
        # Final linear projection
        output = concat @ self.Wo
        return output, all_weights

np.random.seed(42)
d_model, n_heads, seq_len = 32, 4, 6
X = np.random.randn(seq_len, d_model)

mha    = MultiHeadAttentionScratch(d_model, n_heads)
out, weights = mha.forward(X)

print(f"Multi-head attention (d_model={d_model}, heads={n_heads}):")
print(f"  Input shape:    {X.shape}")
print(f"  Output shape:   {out.shape}  ← same as input")
print(f"  d_k per head:   {d_model // n_heads}")
print(f"  Weight shapes:  {weights[0].shape} × {n_heads} heads")

print("
Attention weight patterns per head (row 0 = token 0 attending to all):")
for i, w in enumerate(weights):
    print(f"  Head {i+1}: {w[0].round(3)}")
print("Each head learns a different attention pattern")

# ── PyTorch nn.MultiheadAttention ─────────────────────────────────────
mha_pt = nn.MultiheadAttention(
    embed_dim=32,
    num_heads=4,
    dropout=0.0,
    batch_first=True,
)
X_pt   = torch.FloatTensor(X).unsqueeze(0)   # (1, seq, d_model)
out_pt, attn_pt = mha_pt(X_pt, X_pt, X_pt)  # Q=K=V=X for self-attention
print(f"
PyTorch MultiheadAttention output: {tuple(out_pt.shape)}")

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

# ── Positional encoding — inject position info since attention is order-agnostic ──
class PositionalEncoding(nn.Module):
    """
    Self-attention has no notion of token order — "cat sat mat" and
    "mat sat cat" produce the same attention scores without positional encoding.
    Sinusoidal PE adds a unique position signal to each token embedding.
    """
    def __init__(self, d_model, max_seq_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        # Compute sinusoidal encoding once
        pe   = torch.zeros(max_seq_len, d_model)
        pos  = torch.arange(max_seq_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 dimensions
        pe[:, 1::2] = torch.cos(pos * div)   # odd dimensions
        self.register_buffer('pe', pe.unsqueeze(0))  # (1, max_len, d_model)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

# ── Single Transformer encoder block ──────────────────────────────────
class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model=128, n_heads=4, d_ff=512, dropout=0.1):
        super().__init__()
        # Self-attention sub-layer
        self.attention  = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout, batch_first=True,
        )
        # Feed-forward sub-layer
        self.ff         = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),              # GELU is standard in modern transformers
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        # Layer norms — Pre-LN style (norm before sub-layer, not after)
        self.norm1      = nn.LayerNorm(d_model)
        self.norm2      = nn.LayerNorm(d_model)
        self.dropout    = nn.Dropout(dropout)

    def forward(self, x, src_key_padding_mask=None):
        # ── Self-attention with residual ──────────────────────────────
        # Pre-LN: normalise first, then attend, then add residual
        normed   = self.norm1(x)
        attn_out, _ = self.attention(
            normed, normed, normed,
            key_padding_mask=src_key_padding_mask,
        )
        x = x + self.dropout(attn_out)   # residual connection

        # ── Feed-forward with residual ────────────────────────────────
        x = x + self.dropout(self.ff(self.norm2(x)))
        return x

# ── Stack multiple blocks — a full Transformer encoder ────────────────
class TransformerEncoder(nn.Module):
    def __init__(self, vocab_size, d_model=128, n_heads=4,
                 n_layers=4, d_ff=512, max_seq=512, dropout=0.1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pos_enc   = PositionalEncoding(d_model, max_seq, dropout)
        self.layers    = nn.ModuleList([
            TransformerEncoderBlock(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])
        self.norm      = nn.LayerNorm(d_model)

    def forward(self, x, padding_mask=None):
        # x: (batch, seq_len) — token indices
        x = self.embedding(x)    # (batch, seq, d_model)
        x = self.pos_enc(x)
        for layer in self.layers:
            x = layer(x, src_key_padding_mask=padding_mask)
        return self.norm(x)      # (batch, seq, d_model)

# ── Shape check ───────────────────────────────────────────────────────
torch.manual_seed(42)
model    = TransformerEncoder(vocab_size=5000, d_model=128, n_heads=4, n_layers=2)
x_tokens = torch.randint(1, 5000, (8, 32))   # batch=8, seq_len=32
out      = model(x_tokens)

total = sum(p.numel() for p in model.parameters())
print(f"Transformer Encoder (d=128, heads=4, layers=2):")
print(f"  Input:  {tuple(x_tokens.shape)}")
print(f"  Output: {tuple(out.shape)}  ← contextual token representations")
print(f"  Params: {total:,}")

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

# ── Causal mask — the key difference between BERT and GPT ─────────────
def make_causal_mask(seq_len):
    """
    Upper triangular mask: token i cannot attend to token j > i.
    PyTorch MultiheadAttention uses True = IGNORE (confusing but correct).
    """
    mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
    return mask

seq = 6
mask = make_causal_mask(seq)
print("Causal mask (True = blocked):")
print(mask.int().numpy())
print("Token 0 can only attend to: token 0")
print("Token 3 can attend to: tokens 0, 1, 2, 3")
print("Token 5 can attend to: all tokens 0-5")

# ── GPT-style decoder block ────────────────────────────────────────────
class GPTBlock(nn.Module):
    def __init__(self, d_model=128, n_heads=4, d_ff=512, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout, batch_first=True,
        )
        self.ff   = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.GELU(),
            nn.Dropout(dropout), nn.Linear(d_ff, d_model),
        )
        self.norm1   = nn.LayerNorm(d_model)
        self.norm2   = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        seq_len   = x.size(1)
        # Causal mask — prevents attending to future tokens
        causal    = make_causal_mask(seq_len).to(x.device)
        normed    = self.norm1(x)
        attn_out, _ = self.attention(normed, normed, normed, attn_mask=causal)
        x = x + self.dropout(attn_out)
        x = x + self.dropout(self.ff(self.norm2(x)))
        return x

class MiniGPT(nn.Module):
    def __init__(self, vocab_size=1000, d_model=128, n_heads=4,
                 n_layers=4, d_ff=512, max_seq=256):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb   = nn.Embedding(max_seq, d_model)   # learnable PE (GPT style)
        self.blocks    = nn.ModuleList([
            GPTBlock(d_model, n_heads, d_ff) for _ in range(n_layers)
        ])
        self.norm      = nn.LayerNorm(d_model)
        self.lm_head   = nn.Linear(d_model, vocab_size, bias=False)
        # Weight tying: token embedding and lm_head share weights
        self.lm_head.weight = self.token_emb.weight

    def forward(self, idx):
        # idx: (batch, seq_len) token indices
        B, T     = idx.shape
        tok_emb  = self.token_emb(idx)
        pos_emb  = self.pos_emb(torch.arange(T, device=idx.device))
        x        = tok_emb + pos_emb
        for block in self.blocks:
            x = block(x)
        x        = self.norm(x)
        logits   = self.lm_head(x)   # (B, T, vocab_size)
        return logits

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0):
        """Autoregressive generation — one token at a time."""
        for _ in range(max_new_tokens):
            logits  = self(idx)[:, -1, :]   # last token's logits
            logits  = logits / temperature
            probs   = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx     = torch.cat([idx, next_id], dim=1)
        return idx

torch.manual_seed(42)
gpt   = MiniGPT(vocab_size=1000, d_model=64, n_heads=4, n_layers=2)
x     = torch.randint(1, 1000, (2, 16))
logits = gpt(x)
total  = sum(p.numel() for p in gpt.parameters())

print(f"
MiniGPT (d=64, heads=4, layers=2, vocab=1000):")
print(f"  Input:   {tuple(x.shape)}")
print(f"  Logits:  {tuple(logits.shape)}  ← (batch, seq, vocab)")
print(f"  Params:  {total:,}")

# Generate 5 tokens from a seed
seed     = torch.randint(1, 1000, (1, 4))
generated = gpt.generate(seed, max_new_tokens=5)
print(f"
Generation: {seed[0].tolist()} → {generated[0].tolist()}")

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data import Dataset, DataLoader
import warnings
warnings.filterwarnings('ignore')

torch.manual_seed(42)
np.random.seed(42)

# ── Simulate Stripe dispute classification ───────────────────────────
# In production: use HuggingFace transformers
# pip install transformers
# from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

# Here: demonstrate the fine-tuning pattern with a minimal Transformer
DISPUTE_CLASSES = [
    'fraudulent_charge',
    'service_not_received',
    'wrong_amount',
    'duplicate_charge',
]
N_CLASSES = len(DISPUTE_CLASSES)
VOCAB_SIZE = 500
MAX_LEN    = 32

class DisputeDataset(Dataset):
    def __init__(self, n=800):
        np.random.seed(42)
        self.labels   = np.random.randint(0, N_CLASSES, n)
        # Each class: different vocabulary distribution (simulates real text)
        self.sequences = []
        for label in self.labels:
            seq = np.random.randint(1, VOCAB_SIZE, MAX_LEN)
            # Add class-specific signal to first 5 tokens
            seq[:5] = label * 50 + np.random.randint(1, 50, 5)
            self.sequences.append(seq)

    def __len__(self): return len(self.labels)

    def __getitem__(self, i):
        return (
            torch.LongTensor(self.sequences[i]),
            self.labels[i],
        )

train_ds = DisputeDataset(640)
val_ds   = DisputeDataset(160)
train_ld = DataLoader(train_ds, batch_size=32, shuffle=True)
val_ld   = DataLoader(val_ds,   batch_size=32)

# ── BERT-style classifier — encoder + [CLS] token classification head ──
class BERTClassifier(nn.Module):
    """
    BERT uses a special [CLS] token prepended to every input.
    The final hidden state of [CLS] is used for classification.
    It aggregates information from the entire sequence via attention.
    """
    def __init__(self, vocab_size=VOCAB_SIZE, d_model=64, n_heads=4,
                 n_layers=2, n_classes=N_CLASSES, dropout=0.1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size + 1, d_model, padding_idx=0)
        self.pos_emb   = nn.Embedding(MAX_LEN + 1, d_model)
        self.layers    = nn.ModuleList([
            self._make_block(d_model, n_heads, dropout)
            for _ in range(n_layers)
        ])
        self.norm      = nn.LayerNorm(d_model)
        # Classification head — applied to [CLS] token representation
        self.classifier = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, n_classes),
        )
        self.cls_token = nn.Parameter(torch.randn(1, 1, d_model))

    def _make_block(self, d_model, n_heads, dropout):
        return nn.ModuleDict({
            'attn':  nn.MultiheadAttention(d_model, n_heads,
                         dropout=dropout, batch_first=True),
            'ff':    nn.Sequential(
                         nn.Linear(d_model, d_model*4), nn.GELU(),
                         nn.Dropout(dropout), nn.Linear(d_model*4, d_model)),
            'norm1': nn.LayerNorm(d_model),
            'norm2': nn.LayerNorm(d_model),
            'drop':  nn.Dropout(dropout),
        })

    def forward(self, x):
        B, T  = x.shape
        tok   = self.embedding(x)
        pos   = self.pos_emb(torch.arange(T, device=x.device))
        x_emb = tok + pos

        # Prepend [CLS] token
        cls   = self.cls_token.expand(B, -1, -1)
        x_emb = torch.cat([cls, x_emb], dim=1)   # (B, T+1, d)

        h = x_emb
        for block in self.layers:
            n   = block['norm1'](h)
            a, _ = block['attn'](n, n, n)
            h   = h + block['drop'](a)
            h   = h + block['drop'](block['ff'](block['norm2'](h)))
        h = self.norm(h)

        # Use [CLS] token (position 0) for classification
        cls_repr = h[:, 0, :]
        return self.classifier(cls_repr)

model     = BERTClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0.01)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

print("Fine-tuning BERT-style classifier on Stripe disputes:")
print(f"{'Epoch':>6} {'Loss':>10} {'Val acc':>10}")
print("─" * 30)

for epoch in range(1, 21):
    model.train()
    total_loss = 0
    for Xb, yb in train_ld:
        optimizer.zero_grad()
        loss = criterion(model(Xb), yb)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()
    scheduler.step()

    if epoch % 4 == 0:
        model.eval()
        correct = 0
        with torch.no_grad():
            for Xb, yb in val_ld:
                correct += (model(Xb).argmax(1) == yb).sum().item()
        acc = correct / len(val_ds)
        print(f"  {epoch:>4}  {total_loss/len(train_ld):>10.4f}  {acc:>10.4f}")

# ── A capacity-planning calculator, not a bug report ───────────────────
# For a batch of requests at a given context length, roughly how much
# memory does the attention matrix alone require? This is the kind of
# back-of-envelope math an infra team runs before committing to a
# maximum context length in a product spec.

def attention_matrix_gb(batch_size, n_heads, seq_len, bytes_per_val=4):
    # One (seq_len x seq_len) score matrix per head, per batch item
    cells = batch_size * n_heads * seq_len * seq_len
    return cells * bytes_per_val / (1024 ** 3)

print(f"{'Context length':>15} {'Attention memory (batch=8, 8 heads)':>38}")
print("-" * 56)
for seq_len in [512, 2_048, 8_192, 32_768, 131_072]:
    mem_gb = attention_matrix_gb(batch_size=8, n_heads=8, seq_len=seq_len)
    print(f"  {seq_len:>13,}  {mem_gb:>36.2f} GB")

print()
print("Going from 2,048 to 8,192 tokens is a 4x jump in sequence length")
print("and a 16x jump in attention memory alone -- before weights,")
print("activations, or gradients. This is exactly why 'just raise the")
print("context window' is a GPU-tier and cost decision for an infra")
print("team, not a one-line config change.")

Try prompting a hosted pretrained model first — no training at all. If accuracy or format control is not good enough, fine-tune a small open pretrained model on task-specific labelled data, often with LoRA to keep the cost and turnaround low. If long documents need to be handled, evaluate retrieval-augmented generation against simply widening the context window before assuming more context is the answer — retrieval is frequently both cheaper and more accurate, since it avoids paying attention's quadratic cost for context the model will not actually need on most requests. Training a Transformer from scratch is essentially never the first option evaluated.

Self-attention computes Q·Kᵀ for every pair of tokens using only their content vectors — there is nothing in the raw operation that depends on where a token sits in the sequence. Shuffle "the cat sat on the mat" into "mat the on sat cat the" and, without positional encoding, self-attention produces the exact same set of pairwise scores rearranged — it is permutation-equivariant, not order-aware. That is precisely why every Transformer adds a positional signal (sinusoidal in the original paper, learned in BERT and GPT) into the token embedding before the first attention layer. "Sees everything" is not the same as "knows the order of what it sees" — attention treats the input as a set until position is injected explicitly.

Multi-head attention with h heads and single-head attention with one d_model-sized head use the same total number of parameters — the projection matrices are simply split into h smaller pieces (d_k = d_model / h per head) instead of one large piece. Nothing about the head count adds parameters or FLOPs to the QKV projections themselves. What multiple heads buy you is not more capacity but more independent subspaces: head 1 might learn to track subject-verb agreement while head 4 tracks coreference, each computing its own attention pattern over the same tokens. A single head is mathematically forced to average all these relationship types into one attention distribution per token — multi-head lets them coexist.

The scaling solves a specific, provable problem: when Q and K entries are drawn from a distribution with unit variance, the dot product Q·K summed over d_k dimensions has variance proportional to d_k itself — so as d_k grows, raw attention scores grow with it. Large scores push softmax into its saturated region, where the largest score gets a weight near 1 and every other gets a weight near 0, and the gradient of softmax there is close to zero everywhere. Dividing by √d_k exactly cancels the variance growth — a sum of d_k independent unit-variance terms, divided by √d_k, has variance 1 regardless of d_k — keeping softmax in a well-conditioned regime no matter how large the head dimension gets. It is a derived correction, not an arbitrary constant.

An attention weight of 0.61 from "bank" to "river" tells you the softmax-normalised dot product between those two vectors at that layer — it does not tell you that "river" caused the model's downstream prediction, nor does it account for the value vectors those weights multiply, the residual stream carrying information around attention entirely, or the many other layers and heads contributing to the final output. Research on attention-as-explanation has repeatedly found you can often edit attention weights substantially without changing the model's output, and find alternative attention patterns that produce the same prediction — the opposite of what a true causal explanation should allow. Treat attention weights as a debugging signal and a useful visualisation, not a certified explanation of model behaviour.

Parallelisable and cheap are different properties. An RNN does O(n) sequential steps, each O(1) relative to sequence length, for O(n) total compute and O(1) memory per step. Self-attention computes a full (seq_len × seq_len) score matrix, which is O(n²) in both compute and memory — for n=4096 that is 16 million score pairs before a single head's output is even computed, and it scales quadratically from there. Attention wins on wall-clock time because O(n²) work spread across thousands of GPU cores in parallel finishes faster than O(n) work done one step at a time — but for long enough sequences, the quadratic memory footprint of the attention matrix becomes the actual bottleneck, which is exactly why Flash Attention, sparse attention, and linear-attention variants exist.

An RNN processes tokens strictly in sequence — token 3's hidden state is computed from token 2's hidden state, so the order tokens arrive in is baked into the computation itself; you cannot get the same hidden state from a shuffled input. Self-attention has no such mechanism: it computes Q·Kᵀ between every pair of tokens using only their content, independent of position, so it is mathematically permutation-equivariant — reorder the input tokens and you get the same set of pairwise scores in a different order, not a different computation. To recover any notion of sequence order, the Transformer has to inject it explicitly, adding a positional vector (sinusoidal, or a learned embedding per position) to each token's representation before attention ever runs. Without it, "the dog bit the man" and "the man bit the dog" would look identical to self-attention.

Self-attention computes a full pairwise score matrix of shape (seq_len, seq_len), so both compute and memory scale as O(n²) in sequence length — double the sequence and you quadruple the cost, not double it. For short sequences (a few hundred tokens) this is a non-issue and the massive parallelism across GPU cores makes attention far faster in wall-clock time than an RNN's O(n) sequential steps. But for long-context use cases — long documents, long codebases, long conversations — the quadratic term dominates: at 32k tokens the score matrix alone has a billion entries per head. That's the direct motivation behind Flash Attention (same math, tiled to avoid materialising the full matrix in memory), sparse/local attention patterns, and linear-attention approximations — all attempts to get attention's quality without paying the full O(n²) bill.

A single attention head computes exactly one attention distribution per query token — one weighted average over all the other tokens' values. But a sentence carries multiple simultaneous relationship types at once: which word is the grammatical subject of which verb, which pronoun refers to which noun, which words are simply near each other. Forcing all of that into one softmax distribution means the model has to compromise — blend several different relevance signals into a single weighting, which loses information. Multi-head attention splits the d_model dimension into h smaller subspaces (d_k = d_model/h each) and runs attention independently in each — same total parameter count and compute as one big head, just partitioned — so different heads are free to specialise on different relationship types, and their outputs are concatenated and linearly projected back to d_model at the end.

As the key/query dimension d_k grows, the dot product Q·K — a sum of d_k terms — grows in variance proportionally to d_k, assuming roughly unit-variance inputs. Without correction, larger d_k produces larger-magnitude raw scores, and large scores push softmax toward a near-one-hot output where one token gets almost all the weight and the rest get almost none. In that saturated regime the softmax gradient is close to zero almost everywhere, which stalls learning — the model can't smoothly adjust which tokens to attend to because the loss landscape is nearly flat. Dividing by √d_k exactly compensates for the variance growth from summing d_k terms, keeping the score distribution's variance roughly constant regardless of head dimension, so softmax stays in a well-behaved, differentiable range no matter how wide the attention heads are.

Both are built from the identical Transformer block — multi-head self-attention plus a feedforward network, wrapped in residual connections and LayerNorm. The entire difference is the attention mask and the resulting pretraining objective. BERT applies no mask, so every token attends to every token in both directions — bidirectional context, ideal for understanding tasks like classification or NER where you want the fullest possible context before making a judgment. GPT applies a causal (upper-triangular) mask, so token i can only attend to tokens ≤ i — a hard requirement for autoregressive generation, since predicting token t+1 from tokens 1..t would leak the answer if the model could see token t+1 while training. That single masking choice cascades into everything else: BERT is pretrained with masked language modelling (fill in the blank, needs both directions), GPT with next-token prediction (needs only the past) — the mask is the cause, the pretraining objective and downstream use case are the effects.