Attention and the transformer block interview questions

Score each token's query against every visible key, divide by √d, softmax, and weight the values. The query is what a token looks for, the key what it offers, the value what it hands over. Dot products grow with d, and large scores push softmax onto one token, where gradients get tiny.

Two sub-blocks, each layer norm, then the sublayer, then a residual add: first multi-head attention, then an MLP.

Attention moves information between tokens; the MLP transforms each token on its own and holds most parameters. In GPT-2 small's first block, the MLP has 4,722,432 against attention's 2,362,368, and the model stacks 12 blocks.

It hides later tokens: before softmax, every score whose key comes after the query becomes minus infinity, so those weights are exactly 0 and the rest still sum to 1.

So a decoder disambiguates only from the past: whether the ladder sentence ends "too short" or "too high", GPT-2 small's attention rows up to and including " it" were identical in all 12 layers and 12 heads. That also makes the KV cache exact.

It's a lead to test, not a conclusion. Head 4.3 puts 0.659 on " ladder" from " it", but " did", " not" and " reach" also look mostly at " ladder", and the subject always sat in the first two tokens. That can't separate "finds the subject" from "likes an early position".

Attention weights show where a head reads from, not what the values carry or how 144 heads combine. I'd test sentences with the noun in different positions.

Attention compute grows quadratically. Each head scores an n × n grid, so doubling n roughly quadruples that work; the paper puts self-attention per layer at n² · d. GPT-2's 1,024-token window is about a million scores per head. The KV cache grows linearly, with context and batch size, and with learned positions the window is a hard cap.

Fixes trade flexibility for speed and memory: a local attention window, keys and values shared across heads, a cache in fewer bits, or architectures that scale linearly. I'd measure latency and memory at both lengths first.

Project once, split into heads, attend in each head, merge.

def mhsa(x, w_qkv, w_o, h):            # x: (B, T, d)
    B, T, d = x.shape
    q, k, v = (x @ w_qkv).split(d, -1)
    heads = lambda t: t.view(B, T, h, -1).transpose(1, 2)
    q, k, v = heads(q), heads(k), heads(v)  # (B, h, T, d/h)
    s = q @ k.transpose(-2, -1) / (d // h) ** 0.5
    future = torch.ones(T, T, dtype=torch.bool,
                        device=x.device).triu(1)
    s = s.masked_fill(future, float("-inf"))
    y = s.softmax(-1) @ v
    return y.transpose(1, 2).reshape(B, T, d) @ w_o

Interviewers watch for the mask before softmax, scaling by the head size, not d, and the transpose before the final reshape.

Pre-norm should read x + attn(norm(x)), then x + mlp(norm(x)): normalise each sublayer's input and add its output to the un-normalised stream.

Common bugs: adding to the normed tensor, so the residual path stops being an identity; normalising the sublayer's output instead of its input; one norm shared where two are meant; and no final norm before the output layer. A quick test: zero both output projections, and the block must return x exactly.

LayerNorm subtracts the mean, divides by the standard deviation, then applies a learned scale and shift. RMSNorm skips the mean: it divides by the root mean square and keeps only the scale. That's less computation with comparable quality in the paper's experiments, so models like Llama use it, in the same pre-norm positions.

def sdpa(q, k, v, mask=None):  # mask True = may attend
    s = q @ k.transpose(-2, -1) / q.size(-1) ** 0.5
    if mask is not None:
        s = s.masked_fill(~mask, float("-inf"))
    s = s - s.amax(-1, keepdim=True)
    p = s.exp()
    return (p / p.sum(-1, keepdim=True)) @ v

Scale by √d so scores don't saturate softmax. Subtract the row max before exp so nothing overflows; torch's softmax does this internally. A fully masked row gives NaN, since max minus max is inf minus inf, so guard it. In FP16, compute scores and softmax in FP32. FlashAttention does the same blockwise with a running max.

def self_attn(x, wq, wk, wv):        # x: (T, d)
    q, k, v = x @ wq, x @ wk, x @ wv
    s = q @ k.T / q.shape[-1] ** 0.5  # (T, T)
    s = s - s.max(-1, keepdims=True)
    p = np.exp(s)
    return (p / p.sum(-1, keepdims=True)) @ v

Time is O(T² d) and the score matrix takes O(T²) memory, against O(T d²) for projections, so long sequences are attention-dominated. FlashAttention stays exact but tiles the work so the T-by-T matrix never reaches HBM. Sliding-window and sparse patterns make it linear in T; linear and low-rank attention approximate; MQA and GQA shrink the KV cache.

import math

def sigmoid(x):
    if x >= 0:
        return 1 / (1 + math.exp(-x))
    e = math.exp(x)  # x is negative: no overflow
    return e / (1 + e)

def softmax(xs):
    m = max(xs)  # shift so the largest exp is 1
    es = [math.exp(x - m) for x in xs]
    s = sum(es)
    return [e / s for e in es]

def attention(Q, K, V):
    d = len(Q[0])
    out = []
    for q in Q:
        scores = []
        for k in K:
            dot = 0.0
            for a, b in zip(q, k):
                dot += a * b
            scores.append(dot / math.sqrt(d))
        w = softmax(scores)
        row = [0.0] * len(V[0])
        for wi, v in zip(w, V):
            for j in range(len(v)):
                row[j] += wi * v[j]
        out.append(row)
    return out

The trap is overflow: exp of a large number blows up. Branch sigmoid on the sign, and subtract the max before softmax, which leaves the result unchanged. Attention is then dot products, scaled by √d, a softmax per query, and a weighted sum of values.

Self-attention: each token scores its query against every key, softmaxes, and sums the values by weight.

LoRA: freeze the weights, train a low-rank update.

Adam tracks running averages of each parameter's gradient and squared gradient, giving per-parameter step sizes: the transformer default. SGD uses one global rate and less memory; it often generalises well in vision.

ViT cuts an image into patches, embeds each as a token, and runs a standard transformer encoder.

import torch, torch.nn as nn, torch.nn.functional as F

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.h = h
        self.ln1 = nn.LayerNorm(d)
        self.ln2 = nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.mlp = nn.Sequential(
            nn.Linear(d, 4 * d), nn.GELU(),
            nn.Linear(4 * d, d))

    def attn(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)
        q, k, v = (t.view(B, T, self.h, -1)
                   .transpose(1, 2) for t in (q, k, v))
        s = q @ k.transpose(-2, -1) / k.shape[-1] ** 0.5
        m = torch.ones(T, T, dtype=torch.bool,
                       device=x.device).tril()
        s = s.masked_fill(~m, float("-inf"))
        y = s.softmax(-1) @ v
        y = y.transpose(1, 2).reshape(B, T, C)
        return self.proj(y)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        return x + self.mlp(self.ln2(x))

class Decoder(nn.Module):
    def __init__(self, vocab, d, h, n, max_len):
        super().__init__()
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(max_len, d)
        self.blocks = nn.ModuleList(
            Block(d, h) for _ in range(n))
        self.ln = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)

    def forward(self, idx, targets=None):
        pos = torch.arange(idx.shape[1],
                           device=idx.device)
        x = self.tok(idx) + self.pos(pos)
        for b in self.blocks:
            x = b(x)
        logits = self.head(self.ln(x))
        if targets is None:
            return logits
        loss = F.cross_entropy(logits.flatten(0, 1),
                               targets.flatten())
        return logits, loss

Interviewers watch the causal mask, head reshapes that keep batch and heads apart, pre-norm residuals, and targets shifted by one token outside the model.

The candidate item is the query; the user's behaviour sequence supplies the keys and values. Each past item is embedded from its ID, category and context like time gap, then projected into a key and a value; the candidate's embedding is projected to the query.

Softmax over query-key scores weights the history, so the output is a candidate-specific interest vector, the target attention DIN introduced. Cost grows with history length per candidate, so long histories are truncated or filtered first.

State shapes as you go: x is T by d, scores T by T. Mask before the softmax, and normalise stably.

def attention(x, Wq, Wk, Wv, mask):   # x: (T, d)
    q, k, v = x @ Wq, x @ Wk, x @ Wv  # (T, d_k), (T, d_v)
    s = q @ k.T / np.sqrt(k.shape[-1])  # (T, T)
    s = np.where(mask, s, -1e9)       # mask: True = keep
    s = s - s.max(-1, keepdims=True)
    p = np.exp(s)
    p /= p.sum(-1, keepdims=True)
    return p @ v                      # (T, d_v)

def cross_entropy(logits, y, ignore=-100):  # (N, V), (N,)
    keep = y != ignore
    z, t = logits[keep], y[keep]
    m = z.max(-1, keepdims=True)
    lse = m[:, 0] + np.log(np.exp(z - m).sum(-1))
    return (lse - z[np.arange(len(t)), t]).mean()

Cross-entropy is log-sum-exp minus the target logit, averaged over non-ignored tokens. Interviewers check the scale by the key width, masking scores rather than zeroing probabilities, and that padding labels are skipped.

RNNs carry one hidden state step by step: cheap per step, but sequential, and gradients vanish over long spans. LSTMs add gates and a cell state so information survives longer. Transformers attend over all steps in parallel, train faster and capture long range, at cost quadratic in length; for streaming prediction I'd pick a causal decoder.

MPC is different: it uses a dynamics model to optimise actions over a short horizon, applies the first and replans. A learned sequence model can be that dynamics model or forecaster, while MPC enforces constraints.

Have the forward pass cache what the backward pass needs: the softmax probabilities P.

def attn_fwd(Q, K, V):              # (T, d) each
    S = Q @ K.T / np.sqrt(Q.shape[-1])
    S = S - S.max(-1, keepdims=True)
    P = np.exp(S)
    P /= P.sum(-1, keepdims=True)
    return P @ V, P

def attn_bwd(dO, Q, K, V, P):
    dV = P.T @ dO
    dP = dO @ V.T
    dS = P * (dP - (dP * P).sum(-1, keepdims=True))
    dS /= np.sqrt(Q.shape[-1])
    return dS @ K, dS.T @ Q, dV     # dQ, dK, dV

The key line is the softmax backward: each row's gradient is P times dP minus its P-weighted mean. Multi-head attention reshapes into heads around this, and a block adds residuals, LayerNorm and an MLP, each with its own backward. Prove it with a finite-difference gradient check.

Convolutions build in locality and translation equivariance: each layer sees a small neighbourhood so they learn well from less data, but long-range links need depth or pooling. Self-attention connects every position to every other in one layer with few built-in assumptions, so it needs more data or pretraining, and its cost grows quadratically with sequence length. Hybrids take both: a convolutional stem feeding transformer layers, or convolution stages early and attention later, as in CoAtNet.

Masks set scores to minus infinity before the softmax. The causal mask hides future positions; the padding mask hides pad keys, and pads are excluded from the loss.

In cross-attention, queries come from the decoder, keys and values from the encoder.

At decode time, past keys and values never change, so the KV cache stores them; each step computes projections for one new token instead of re-running the prefix.

Mixed precision runs matmuls in FP16 or BF16, keeps FP32 master weights and optimizer state, and with FP16 scales the loss so small gradients don't underflow.

import math
import torch

def sinusoidal_pe(T, d):
    pos = torch.arange(T).float()[:, None]   # (T, 1)
    i = torch.arange(0, d, 2).float()        # (d/2,)
    freq = torch.exp(-math.log(10000.0) * i / d)
    pe = torch.zeros(T, d)                   # d even
    pe[:, 0::2] = torch.sin(pos * freq)
    pe[:, 1::2] = torch.cos(pos * freq)
    return pe

# x: (B, T, d); x = x + sinusoidal_pe(T, d).to(x)
# learned: x = x + nn.Embedding(max_len, d)(pos_ids)

Dimension pair 2i gets sin and cos of pos divided by 10000 to the power 2i/d. A learned table is simpler but can't go past max_len.

Without positions, attention is permutation-equivariant: an encoder sees a bag of words, so shuffled inputs give shuffled outputs and order-dependent tasks fail. A causal decoder recovers some position from its mask.