Transformers & Attention

Starting from Zero — A Physical Intuition

Before looking at equations, let's understand self-attention through a physical office filing system analogy:

• The Query (Q) — Your Sticky Note Search: Imagine you are working at a desk and need information about a specific task. You write down exactly what you're looking for on a sticky note. This search term is your Query.
• The Key (K) — The Drawer Labels: Around your office, there are filing cabinets. Each drawer has a label identifying its contents. These labels are the Keys. You compare your sticky note (Query) against every drawer label (Key) to find the best match.
• The Value (V) — The Document Inside: Once you find the cabinet drawer whose label (Key) matches your sticky note (Query), you open it and extract the actual files inside. These files are the Values.
• Self-Attention — Blending the Files: If your search matches multiple drawers, you don't just pull one document; you pull files from all matching drawers and compile a summary. The brightness of your match (dot product similarity) determines how much page space each file gets in your final summary (weighted average).

The problem attention solved

Before transformers, sequence models were RNNs / LSTMs: they read tokens one at a time, carrying a hidden state forward. Two problems killed them at scale:

• They can't parallelize. Token t needs the hidden state from token t−1, so training is inherently sequential — you can't use a GPU's thousands of cores on one sequence at once.
• Long-range memory decays. Information from token 1 has to survive being squeezed through hundreds of state updates to influence token 500. It fades (the vanishing-gradient problem, stretched over time).

The 2017 paper "Attention Is All You Need" threw out recurrence entirely. The idea: let every token look directly at every other token in a single step — no chain to walk, fully parallel, and any two tokens are one hop apart.

Self-attention: Query, Key, Value

Each token is projected (by learned weight matrices) into three vectors:

• Query (Q) — what am I looking for?
• Key (K) — what do I offer?
• Value (V) — what will I contribute if attended to?

A token attends to others by matching its Query against everyone's Keys, then pulling a weighted blend of their Values:

attention(Q, K, V) = softmax( Q · Kᵀ / √d ) · V

Step it through below. The pronoun "it" has a Query that matches "animal"'s Key, so after softmax it pulls mostly animal's Value — that's how a pronoun quietly resolves to its noun. Switch the query token to see what each one looks at:

Three things worth saying in an interview:

• The dot product is a similarity score — Q · K is large when a token is "relevant to me".
• The √d scaling keeps those dot products from growing with dimension; without it the softmax saturates and gradients vanish.
• Softmax makes it a weighted average — the weights sum to 1, so the output is a convex blend of Values. Nothing is recurrent; it's all matrix multiplies, which is exactly why a GPU loves it.

Multi-head attention

One attention computation can only capture one kind of relationship. So transformers run several heads in parallel, each with its own Q/K/V projections, then concatenate and project the results. One head might track subject→verb agreement, another coreference, another nearby word order. It's an ensemble of relationships learned at once — cheap, because the heads run simultaneously.

Positional encoding

Attention treats its input as a set, not a sequence — softmax(QKᵀ/√d)V is the same no matter how you shuffle the tokens. But "dog bites man" ≠ "man bites dog". So before the first layer we add a positional encoding to each token's embedding — a sinusoidal pattern (or a learned vector) that encodes where the token sits. Now position is part of the signal the model can attend to.

The transformer block

A transformer stacks N identical blocks; each block is attention + a feed-forward network, each wrapped in a residual connection + layer norm (the residuals are what let you stack dozens of layers without the gradient dying):

Two flavors:

• Encoder (e.g. BERT): each token attends to all tokens (bidirectional) — great for understanding/classification.
• Decoder (e.g. GPT): causal / masked self-attention — a token may only attend to tokens before it, never the future. That mask is exactly what lets a model generate left-to-right, one token at a time. Modern LLMs are decoder stacks.

Self-Attention in Code

Here is a complete, clean Python implementation of Scaled Dot-Product Attention with causal masking using NumPy:

import numpy as np
from typing import Tuple, Optional

def scaled_dot_product_attention(
    q: np.ndarray, 
    k: np.ndarray, 
    v: np.ndarray, 
    mask: Optional[np.ndarray] = None
) -> Tuple[np.ndarray, np.ndarray]:
    # q, k, v shape: (seq_len, d_k)
    d_k = q.shape[-1]
    
    # 1. Compute similarity scores
    scores = np.matmul(q, k.T) / np.sqrt(d_k)
    
    # 2. Apply causal mask (1 = allow, 0 = mask out)
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)
        
    # 3. Softmax along the last dimension (for weights summing to 1.0)
    exp_scores = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
    weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
    
    # 4. Multiply weights by Values
    output = np.matmul(weights, v)
    return output, weights

# Example: 3 tokens, 2-dimensional embeddings
# Tokens: ["The", "cat", "sat"]
Q = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]])
V = np.array([[10.0, 0.0], [0.0, 20.0], [5.0, 5.0]])

# Causal lower-triangular mask
causal_mask = np.tril(np.ones((3, 3)))

out, attn_weights = scaled_dot_product_attention(Q, K, V, causal_mask)
print("Attention Weights:\n", attn_weights)
print("Output:\n", out)

Think it through like the interview

Don't just state that Transformers are fast — derive the computational bottleneck of sequential loops.

Encoder vs. Decoder: BERT vs. GPT

Depending on how attention masking is set up, the transformer architecture branches into three primary configurations:

Encoder-Decoder & Cross-Attention

In Encoder-Decoder models (like T5):

1. The Encoder processes the input sequence bidirectionally to create rich context vectors.
2. The Decoder generates the target sequence one token at a time.
3. To connect them, the decoder uses Cross-Attention: the decoder's Queries ($Q$) look at the encoder's Keys ($K$) and Values ($V$). This lets the decoder decide which parts of the input sequence to focus on at each generation step.

Vision Transformers (ViT)

Though designed for text, transformers revolutionized computer vision via the Vision Transformer (ViT).

                  Vision Transformer (ViT) Ingestion:
  
  [ Image ] ──> Split into Patches ──> Flatten to Vectors ──> Linear Projection (Tokens)
                 (e.g., 16x16 pixels)    (e.g., 16x16x3 = 768)     + Class [CLS] Token
                                                                       + Position Embeddings
                                                                                │
                                                                                ▼
                                                                     Standard Transformer Encoder

How ViT Works

1. Image Patching: A 2D image is split into a grid of non-overlapping patches (e.g., $16 \times 16$ pixels).
2. Linear Projection: Each patch is flattened into a 1D vector and projected into a vector space of size $D$ (acting as a "visual token" embedding).
3. Position Embedding: 1D learnable position embeddings are added to the patch embeddings to retain spatial relationships.
4. CLS Token: A special learnable [CLS] (classification) token is prepended to the sequence.
5. Transformer Encoder: The sequence of tokens is processed by a standard transformer encoder (bidirectional self-attention). The final output state of the [CLS] token is passed to a classification head.

PyTorch Transformer Training Pipeline (From Scratch)

Here is a complete, runnable script showing how to construct a single Transformer Encoder Layer using PyTorch, run a forward pass, compute loss, and perform a backward gradient update step:

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleTransformerEncoder(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, nhead: int, num_layers: int):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        
        # PyTorch built-in Transformer Layer
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, 
            nhead=nhead, 
            dim_feedforward=d_model * 4,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.classifier = nn.Linear(d_model, vocab_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Embedding: (batch_size, seq_len) -> (batch_size, seq_len, d_model)
        embedded = self.embedding(x)
        
        # Self-Attention + FFN pass
        encoded = self.transformer(embedded)
        
        # Classify back to vocabulary tokens
        logits = self.classifier(encoded)
        return logits

# Initialize Model, Optimizer, and Loss
torch.manual_seed(42)
vocab_size = 1000
d_model = 64
nhead = 4
model = SimpleTransformerEncoder(vocab_size, d_model, nhead, num_layers=1)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

# Mock batch: batch_size=2, seq_len=5
x_inputs = torch.randint(0, vocab_size, (2, 5))
y_targets = torch.randint(0, vocab_size, (2, 5))

# 1. Forward Pass
logits = model(x_inputs)

# 2. Compute Loss (reshape to combine batch and sequence dimensions)
loss = criterion(logits.view(-1, vocab_size), y_targets.view(-1))

# 3. Backward Pass & Step
optimizer.zero_grad()
loss.backward()
optimizer.step()

print(f"Training Step Success! Initial Loss: {loss.item():.4f}")

Why transformers won

• Parallelizable — every token is processed at once, so training scales to internet-sized data and billions of parameters.
• Long-range by construction — any two tokens are one attention hop apart, no decaying state.
• One architecture, every modality — text, code, images (ViT), audio. It's the backbone of essentially every modern foundation model.

The catch: O(n²)

Attention compares every token to every token, so cost grows with the square of the sequence length: O(n²·d) time and O(n²) memory. Double the context, quadruple the work. This is the real reason context windows are finite and long-context is expensive. Mitigations interviewers love to hear:

• FlashAttention — an IO-aware exact implementation that avoids materializing the full n×n matrix in slow memory (a huge practical speedup, same math).
• Sparse / sliding-window attention — each token attends to a local window (or a few global tokens) instead of all n, trading some reach for linear-ish cost.
• KV cache — at inference, cache past Keys/Values so generating token t doesn't recompute attention over the whole prefix.

Common Interview Questions

Interactive Quiz

Practice

1. Attention by hand: given three tokens with 2-D Q/K/V vectors, compute the attention output for one query token — scores, softmax, weighted sum. Check the weights sum to 1.
2. Causal mask: write the mask matrix for a 4-token decoder so each token only attends to itself and earlier tokens. Why is it an upper-triangular −∞?
3. Run the Attention code: Execute the Python self-attention function above. Run it with and without the causal mask and inspect how the output vectors change when looking at past vs future context.
4. Cost estimate: a model has a 128k-token context. Relative to a 4k context, how much more attention compute and memory does one layer need? (Answer with the n² factor.)

Next: Large Language Models — decoder-only transformers, scaled up.