Attention and Transformers — from scratch

Think of attention as a *retrieval system*. You have a query (what you're looking for), a set of keys (what's indexable), and a set of values (the data behind the keys). For each query, compute a similarity score against every key, softmax to get weights, and read out the weighted sum of values.

attention(Q, K, V) = softmax(Q K^T / √d_k) V

In code (one head, one batch):

import torch
import torch.nn.functional as F

Q, K, V = torch.randn(8, 64), torch.randn(8, 64), torch.randn(8, 64)
d_k = 64
scores = (Q @ K.T) / (d_k ** 0.5)         # (8, 8) similarity matrix
weights = scores.softmax(dim=-1)           # rows sum to 1
out = weights @ V                          # (8, 64)

The √d_k scaling keeps the softmax from saturating when d_k is large. The output is a *new* representation where each token has gathered information from every other token, weighted by relevance. This is the entire mechanism. Everything else in a transformer is multi-head plumbing, normalization, and feed-forward layers.

One attention head learns one relational pattern (e.g. 'attend to the subject of the sentence'). Multi-head attention runs h heads in parallel, each with its own projection matrices, and concatenates them. This lets the model attend to different *kinds* of relationships simultaneously.

class MultiHeadAttention(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        assert d % h == 0
        self.h, self.d_k = h, d // h
        self.qkv = nn.Linear(d, 3*d, bias=False)
        self.out = nn.Linear(d, d, bias=False)

    def forward(self, x, mask=None):
        B, T, D = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.h, self.d_k).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]      # (B, h, T, d_k)
        scores = (q @ k.transpose(-2, -1)) / (self.d_k ** 0.5)
        if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = scores.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, T, D)
        return self.out(out)

This is the entire attention layer of GPT, BERT, T5, LLaMA — modulo position encoding, normalization choices, and activation flavor. Read it twice; it pays for years.

Attention can look anywhere by default. Two important restrictions:

• Causal (decoder) mask: token t can only attend to tokens 0..t. Required for autoregressive language modeling (GPT family). Implemented by masking the upper triangle of the score matrix to -inf before softmax:

``python mask = torch.tril(torch.ones(T, T)) scores = scores.masked_fill(mask == 0, float('-inf')) ``

• Padding mask: ignore positions filled with <pad> tokens. Set those columns to -inf so they receive zero attention weight.
• Bidirectional: no mask. Used in encoder-only models (BERT) and the encoder side of encoder-decoder models (T5).

This single switch — masked vs. unmasked — is the difference between BERT and GPT. The architecture is otherwise nearly identical.

Attention is permutation-invariant: attention(Q, K, V) doesn't care about token order. We have to inject position somehow.

• Sinusoidal (original transformer): add fixed sin/cos curves of different frequencies to the input embeddings. No learned params; extrapolates to longer sequences in theory.
• Learned absolute (BERT, original GPT): a nn.Embedding(max_len, d) table added to token embeddings. Doesn't extrapolate past max_len.
• Rotary Position Embeddings (RoPE) (LLaMA, Mistral, GPT-NeoX, most modern LLMs): rotate Q and K in 2D subspaces by an angle proportional to position. Has nice properties for extrapolation when combined with techniques like NTK-aware scaling.
• ALiBi (BLOOM, MPT): no encoding — add a linear bias to attention scores based on distance. Strong extrapolation behavior.
• NoPE (decoder-only experiments): some recent work shows decoder-only models learn position implicitly without explicit encodings.

In 2026, RoPE is the de-facto choice for any new LLM, with ALiBi a respectable alternative when extreme context-length extrapolation matters.

class TransformerBlock(nn.Module):
    def __init__(self, d, h, ffn_mult=4, dropout=0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(d)
        self.attn = MultiHeadAttention(d, h)
        self.norm2 = nn.LayerNorm(d)
        self.ffn = nn.Sequential(
            nn.Linear(d, ffn_mult * d),
            nn.GELU(),
            nn.Linear(ffn_mult * d, d),
        )
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        x = x + self.drop(self.attn(self.norm1(x), mask))
        x = x + self.drop(self.ffn(self.norm2(x)))
        return x

A full GPT or BERT is N (12-96) of these blocks stacked, sandwiched between an embedding layer (token + position) at the input and an LM-head linear at the output. Modern variants swap:

• LayerNorm → RMSNorm (LLaMA, simpler and slightly faster).
• GELU → SwiGLU (LLaMA, Mistral): Linear(d, ffn_mult*d) * silu(Linear(d, ffn_mult*d)). Better quality per parameter.
• Standard multi-head → Grouped-Query Attention (GQA) or Multi-Query Attention (MQA): share K and V across multiple Q heads to slash KV-cache memory at inference. Essential for serving big models.

LLaMA-3 7B is ~32 of these blocks with d=4096, h=32, ffn=14336, RMSNorm, SwiGLU, GQA, RoPE. A ~30-line block plus a few dozen lines of orchestration is what powers a 7-billion-parameter model. The complexity is in the *training data* and *scale*, not the architecture.

• Encoder-only (BERT family): bidirectional attention, trained with masked-language-modeling (predict masked tokens). Best for classification, embedding, retrieval. Examples: BERT, RoBERTa, DeBERTa, ModernBERT. Output: per-token representations + [CLS] summary.

• Decoder-only (GPT family): causal attention, trained with next-token prediction. Best for generation, in-context learning, chat. Examples: GPT-2/3/4, LLaMA-1/2/3, Mistral, Gemma, Qwen. The dominant 2026 architecture.

• Encoder-decoder (T5 family): encoder sees full input bidirectionally, decoder generates autoregressively while cross-attending to the encoder. Best for translation, summarization, structured transformations. Examples: T5, BART, FLAN-T5, mT5. Niche but excellent for fine-tuning on input→output transformation tasks.

Rule of thumb in 2026: if you can pose the task as 'continue this sequence,' use a decoder-only. If you need a fixed-size embedding, use an encoder-only. Encoder-decoder is occasionally still the best fit for translation or structured rewriting.