Transformers and Modern NLP Fundamentals

The Transformer architecture, introduced in "Attention Is All You Need" (2017), is the foundation of every modern large language model: GPT, BERT, LLaMA, Claude, Gemini. Understanding how attention works is not optional for an ML engineer in 2025 - it is the central mechanism you will encounter in NLP, vision, multimodal, and generative AI work.

The Core Motivation: Attention as Content-Based Routing

In an RNN, every token must route information through a sequential chain. Information from the first token must pass through 512 intermediate states to reach the last token - the longer the sequence, the harder it is to preserve.

Attention solves this by letting every token directly access every other token. Instead of routing information through a chain, the model learns to route it by content matching: for each token, ask "which other tokens are relevant to understanding this one?" and directly pull information from them.

The Attention Mechanism

The scaled dot-product attention takes three matrices: Query (Q), Key (K), and Value (V):

Attention(Q, K, V) = softmax(QKᵀ / √dₖ) · V

In Python:

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

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: [batch, heads, seq_len, d_k]
    K: [batch, heads, seq_len, d_k]
    V: [batch, heads, seq_len, d_v]
    """
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # scores: [batch, heads, seq_len_q, seq_len_k]

    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    weights = F.softmax(scores, dim=-1)  # attention weights sum to 1 per query position
    return torch.matmul(weights, V), weights

Intuition: Q·Kᵀ measures similarity between each query-key pair (dot product = cosine similarity when normalized). Softmax converts similarities to a probability distribution over positions. The output is a weighted sum of values - each token receives a blend of all other tokens' values, weighted by their relevance.

The √dₖ scaling prevents the dot products from growing too large as dimension increases (large dot products → softmax becomes very peaked → vanishing gradients during training).

Multi-Head Attention

A single attention head learns one type of relationship. Multi-head attention runs h attention heads in parallel, each on a projected subspace, then concatenates the results:

pythonCopy
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model)

    def split_heads(self, x):
        B, T, D = x.shape
        return x.view(B, T, self.num_heads, self.d_k).transpose(1, 2)
        # [B, num_heads, T, d_k]

    def forward(self, x, mask=None):
        Q = self.split_heads(self.W_q(x))
        K = self.split_heads(self.W_k(x))
        V = self.split_heads(self.W_v(x))

        attn_out, _ = scaled_dot_product_attention(Q, K, V, mask)
        # attn_out: [B, num_heads, T, d_k]

        # Concatenate heads and project
        B, H, T, d_k = attn_out.shape
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, H * d_k)
        return self.W_o(attn_out)

Each head can learn different relationships: head 1 might learn syntactic dependencies, head 2 might learn coreference, head 3 might learn positional proximity.

Positional Encoding

Unlike RNNs, attention is permutation-invariant - it doesn't know which token came first. Positional encoding adds a position signal to each token's embedding.

pythonCopy
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        position = torch.arange(max_len).unsqueeze(1)  # [max_len, 1]
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)  # even dimensions: sine
        pe[:, 1::2] = torch.cos(position * div_term)  # odd dimensions: cosine
        self.register_buffer('pe', pe.unsqueeze(0))    # [1, max_len, d_model]

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

Modern models use Rotary Position Embedding (RoPE) or ALiBi instead - these learned or relative position schemes generalize better to sequences longer than seen during training. But sinusoidal encoding remains useful for understanding the concept.

The Transformer Block

Each Transformer layer consists of Multi-Head Attention + Feed-Forward Network, both with residual connections and layer normalization:

pythonCopy
class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads)
        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, mask=None):
        # Pre-norm (used by modern LLMs): normalize before attention, not after
        x = x + self.dropout(self.attn(self.norm1(x), mask))
        x = x + self.dropout(self.ff(self.norm2(x)))
        return x

The residual connection x = x + f(x) is crucial: it provides a direct gradient path through the network and allows the model to be much deeper (GPT-3 has 96 layers).

Pre-norm (normalize before the sub-layer) vs. post-norm (the original paper): modern LLMs use pre-norm because it has more stable training at large scale.

Encoder vs. Decoder vs. Encoder-Decoder

Encoder-only (BERT): every position attends to every other position (bidirectional). Used for classification, named entity recognition, question answering. Cannot generate text.

Decoder-only (GPT, LLaMA): causal attention - each position can only attend to earlier positions. Used for text generation, language modeling. The dominant architecture for modern LLMs.

Encoder-decoder (T5, original Transformer): encoder processes input with full attention; decoder attends to encoder output and generates output autoregressively. Used for translation, summarization, instruction-following.

pythonCopy
# Causal mask for decoder-only models
def make_causal_mask(seq_len):
    # Lower triangular matrix: position i can see positions 0..i
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, seq_len]

Complexity and Efficiency

The attention mechanism has O(n²) time and memory complexity with respect to sequence length - doubling the sequence length quadruples the memory. For a sequence of 4,096 tokens with 8 attention heads, the attention matrix is 4096×4096×8 = 134M values just for attention weights.

# Memory growth: quadratic with sequence length
seq_len = 1024  → attention matrix: 1024 * 1024 * 8 heads = 8M values
seq_len = 4096  → 134M values
seq_len = 32768 → 8.6B values (requires Grouped Query Attention or Flash Attention)

Flash Attention: reorders computation to avoid materializing the full attention matrix in memory - same result, O(n) memory. Essential for training models with long contexts.

Using a Pretrained Transformer

You will rarely train a Transformer from scratch. The standard workflow:

pythonCopy
from transformers import AutoTokenizer, AutoModel
import torch

model_name = 'bert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

texts = ["The model failed at inference time.", "Training loss diverged at epoch 3."]
encoded = tokenizer(texts, padding=True, truncation=True,
                    max_length=128, return_tensors='pt')

with torch.no_grad():
    outputs = model(**encoded)

# For classification: use the [CLS] token representation
cls_embeddings = outputs.last_hidden_state[:, 0, :]  # [batch, 768]

# For token-level tasks (NER): use all token representations
token_embeddings = outputs.last_hidden_state  # [batch, seq_len, 768]

Common Mistakes and Bad Instincts

Using max_length much larger than your actual sequences. Transformers use O(n²) memory. If your sequences are 64 tokens long on average, setting max_length=512 wastes 64x memory.

Not masking padding tokens in attention. Without a padding mask, the model attends to padding positions and treats them as meaningful content. Always pass attention_mask from the tokenizer to the model.

Forgetting that GPT/decoder models don't have an encoder. You cannot get a "sentence embedding" from a decoder-only model the same way you can from BERT's [CLS] token. Decoder models pool across all positions or use the last token's representation.

Training a Transformer from scratch on a small dataset. Transformers have weak inductive biases compared to CNNs or LSTMs - they need large amounts of data to outperform simpler models. If you have < 100K labeled examples, fine-tune a pretrained model instead.

Mixing up pre-norm and post-norm in custom architectures. If you are stacking Transformer blocks and the loss diverges at scale, check whether you are using pre-norm (more stable) or post-norm (requires careful LR warm-up).

Where to Go Next

• Module 21 (Transfer Learning and Fine-Tuning) covers how to fine-tune pretrained Transformers (BERT, LLaMA) for specific tasks using modern PEFT methods.
• Module 22 (Working with LLMs) covers using LLM APIs, prompt engineering, and context windows for building production AI systems.
• The standalone post attention-mechanism-explained goes deeper on the mathematics of attention, including cross-attention, KV cache, and grouped query attention.