# Attention and the transformer block interview questions

> 20 attention and the transformer block interview questions for AI and ML engineer roles, each with a model answer.

Canonical: https://llmdojo.dev/interview/topic/attention
Last updated: 2026-10-23

20 questions: 3 junior, 11 mid-level and 6 senior. Taught in [Attention: How Tokens Use Their Context](https://llmdojo.dev/learn/model-foundations/attention).

## [Walk me through scaled dot-product attention. What are queries, keys and values, and why divide by the square root of d?](https://llmdojo.dev/interview/q/walk-me-through-scaled-dot-product-attention-what-are-queries-keys-and)

Level: Junior

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.

## [What's inside one transformer block, and where do most of its parameters sit?](https://llmdojo.dev/interview/q/what-s-inside-one-transformer-block-and-where-do-most-of-its)

Level: Junior

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

**Diagram: Inside a transformer block**

1. Each token enters as its embedding plus a position signal, so the model knows the order of the tokens.
2. Attention: every token looks at itself and the tokens before it, and pulls in information from them. (A layer norm first rescales each vector to keep the numbers stable.)
3. The attention result is added back onto the token's own vector: a residual connection. The block refines each vector instead of replacing it.
4. A feed-forward network then processes each token's vector on its own: attention mixes information between tokens, the MLP transforms it within a token.
5. Another residual add, and one block is done. The model stacks many identical blocks with different learned weights; each block's output feeds the next.
6. After the last block, each position's final vector is turned into a score for every token in the vocabulary. Turning those scores into text is the next lesson.

Structure of a pre-norm decoder block as used by GPT-2: layer norm before attention and before the feed-forward network, each followed by a residual add.

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.

## [What does the causal mask do, and what does it imply for a decoder model reading 'The ladder did not reach the window because it was too short'?](https://llmdojo.dev/interview/q/what-does-the-causal-mask-do-and-what-does-it-imply-for)

Level: Mid-level

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.

**Diagram: The causal mask**

1. Rows are the token looking (its query), columns the token looked at (its key). Each cell is a score: the query's dot product with the key, divided by √d.
2. The causal mask: before the softmax, every score where the key comes after the query is set to minus infinity.
3. Softmax turns those cells into exactly 0, and the remaining weights in each row still add up to 1.
4. Token 1 sees only itself, token 2 sees tokens 1 and 2, and so on. Adding tokens later never changes an earlier row, which is why the KV cache is exact.

The causal mask as taught in lesson 3, on the attention lab's three tokens “The river bank”.

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.

## [A teammate shows you a GPT-2 head that puts most of a pronoun's attention on the right noun and concludes that the head resolves pronouns. How do you respond?](https://llmdojo.dev/interview/q/a-teammate-shows-you-a-gpt-2-head-that-puts-most-of)

Level: Mid-level

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.

## [Product wants to double the context you support, from 1,000 to 2,000 tokens. From the attention side, what gets more expensive, and what families of fixes would you weigh?](https://llmdojo.dev/interview/q/product-wants-to-double-the-context-you-support-from-1-000-to)

Level: Senior

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.

## [Implement batched multi-head self-attention with a causal mask.](https://llmdojo.dev/interview/q/implement-batched-multi-head-self-attention-with-a-causal-mask)

Level: Mid-level

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

```python
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.

## [A pre-norm transformer block trains badly. What do you check around the norms and residuals?](https://llmdojo.dev/interview/q/a-pre-norm-transformer-block-trains-badly-what-do-you-check-around)

Level: Senior

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**.

**Diagram: Inside a transformer block**

1. Each token enters as its embedding plus a position signal, so the model knows the order of the tokens.
2. Attention: every token looks at itself and the tokens before it, and pulls in information from them. (A layer norm first rescales each vector to keep the numbers stable.)
3. The attention result is added back onto the token's own vector: a residual connection. The block refines each vector instead of replacing it.
4. A feed-forward network then processes each token's vector on its own: attention mixes information between tokens, the MLP transforms it within a token.
5. Another residual add, and one block is done. The model stacks many identical blocks with different learned weights; each block's output feeds the next.
6. After the last block, each position's final vector is turned into a score for every token in the vocabulary. Turning those scores into text is the next lesson.

Structure of a pre-norm decoder block as used by GPT-2: layer norm before attention and before the feed-forward network, each followed by a residual add.

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.

## [How does RMSNorm differ from LayerNorm, and why do recent models use it?](https://llmdojo.dev/interview/q/how-does-rmsnorm-differ-from-layernorm-and-why-do-recent-models-use)

Level: Mid-level

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.

## [Implement scaled dot-product attention in PyTorch, then tell me where it can go wrong numerically.](https://llmdojo.dev/interview/q/implement-scaled-dot-product-attention-in-pytorch-then-tell-me-where-it)

Level: Mid-level

```python
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.

## [Implement self-attention, give its time and memory complexity, and describe the efficient variants.](https://llmdojo.dev/interview/q/implement-self-attention-give-its-time-and-memory-complexity-and-describe-the)

Level: Mid-level

```python
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.

## [Using only plain Python loops, implement a numerically stable sigmoid, a stable softmax and single-head scaled dot-product attention.](https://llmdojo.dev/interview/q/using-only-plain-python-loops-implement-a-numerically-stable-sigmoid-a-stable)

Level: Junior

```python

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.

## [Quick-fire: explain self-attention, LoRA, Adam versus SGD, and the Vision Transformer.](https://llmdojo.dev/interview/q/quick-fire-explain-self-attention-lora-adam-versus-sgd-and-the-vision)

Level: Mid-level

**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.

## [With no starter code, implement a GPT-style transformer decoder from scratch in about 30 minutes.](https://llmdojo.dev/interview/q/with-no-starter-code-implement-a-gpt-style-transformer-decoder-from-scratch)

Level: Mid-level

```python

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.

## [In a recommender that uses cross-attention between a user's behavior sequence and a candidate item, what exactly are the query, key and value?](https://llmdojo.dev/interview/q/in-a-recommender-that-uses-cross-attention-between-a-user-s-behavior)

Level: Senior

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.

## [Write pseudocode for single-head scaled dot-product self-attention with masks and tensor shapes, and for the cross-entropy loss.](https://llmdojo.dev/interview/q/write-pseudocode-for-single-head-scaled-dot-product-self-attention-with-masks)

Level: Mid-level

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

```python
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()
```

**Diagram: The causal mask**

1. Rows are the token looking (its query), columns the token looked at (its key). Each cell is a score: the query's dot product with the key, divided by √d.
2. The causal mask: before the softmax, every score where the key comes after the query is set to minus infinity.
3. Softmax turns those cells into exactly 0, and the remaining weights in each row still add up to 1.
4. Token 1 sees only itself, token 2 sees tokens 1 and 2, and so on. Adding tokens later never changes an earlier row, which is why the KV cache is exact.

The causal mask as taught in lesson 3, on the attention lab's three tokens “The river bank”.

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.

## [Compare RNNs, LSTMs and Transformers for sequence modeling, which Transformer variant you used and why, and how model predictive control fits in.](https://llmdojo.dev/interview/q/compare-rnns-lstms-and-transformers-for-sequence-modeling-which-transformer-variant-you)

Level: Senior

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.

## [Implement scaled dot-product multi-head attention and a Transformer block from scratch, including the backward pass for attention, without autograd.](https://llmdojo.dev/interview/q/implement-scaled-dot-product-multi-head-attention-and-a-transformer-block-from)

Level: Senior

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

```python
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**.

## [Compare convolutional networks and Transformers on inductive bias, data and compute needs, and long-range dependencies. Where do hybrids fit?](https://llmdojo.dev/interview/q/compare-convolutional-networks-and-transformers-on-inductive-bias-data-and-compute-needs)

Level: Mid-level

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.

## [Walk me through Transformer attention: causal and padding masks, cross-attention, why a KV cache helps at inference, and how mixed-precision training works.](https://llmdojo.dev/interview/q/walk-me-through-transformer-attention-causal-and-padding-masks-cross-attention-why)

Level: Senior

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.

**Diagram: The causal mask**

1. Rows are the token looking (its query), columns the token looked at (its key). Each cell is a score: the query's dot product with the key, divided by √d.
2. The causal mask: before the softmax, every score where the key comes after the query is set to minus infinity.
3. Softmax turns those cells into exactly 0, and the remaining weights in each row still add up to 1.
4. Token 1 sees only itself, token 2 sees tokens 1 and 2, and so on. Adding tokens later never changes an earlier row, which is why the KV cache is exact.

The causal mask as taught in lesson 3, on the attention lab's three tokens “The river bank”.

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.

## [Implement sinusoidal or learned positional encodings in PyTorch, explain the equations and tensor shapes, and say what you'd see if positional information were missing.](https://llmdojo.dev/interview/q/implement-sinusoidal-or-learned-positional-encodings-in-pytorch-explain-the-equations-and)

Level: Mid-level

```python

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.

## Related
- [All interview questions](https://llmdojo.dev/interview)