Transformer Architecture Interview Prep

Attention Variants — KV-Cache Size vs Quality Trade-off (the decode-memory lever)

Sample interview questions

1. What was the key architectural shift introduced by Attention Is All You Need?
   • A. It replaced attention with recurrence.
   • B. It replaced recurrent sequence modeling with stacked self-attention and feed-forward blocks. ✓
   • C. It removed the need for positional information.
   • D. It made all sequence models decoder-only.

   Option B is correct. The paper's core contribution was to remove recurrence from the sequence model and replace it with self-attention plus feed-forward blocks. That change made training more parallel and set up the modern transformer family.

2. Why did the transformer need positional encodings?
   • A. Because self-attention is permutation-equivariant and otherwise cannot tell token order. ✓
   • B. Because positional encodings reduce model size.
   • C. Because they replace the feed-forward network.
   • D. Because they make the model causal by default.

   Option A is correct. Self-attention by itself does not encode token order, so the model needs positional information to distinguish different sequences with the same tokens in a different order.

3. What does multi-head attention buy you in the original transformer?
   • A. It lets the model attend to multiple relationship patterns in parallel. ✓
   • B. It removes the need for residual connections.
   • C. It turns the model into an RNN.
   • D. It eliminates the need for scaling the dot product.

   Option A is correct. Multi-head attention splits the attention computation into several heads so the model can focus on different patterns at once, such as syntax, long-range dependencies, or positional relations.

4. Why was the transformer easier to scale than recurrent models?
   • A. It processed all tokens in parallel during training. ✓
   • B. It required no optimization.
   • C. It used no learned parameters.
   • D. It only worked on short sequences.

   Option A is correct. By removing recurrence, the transformer allowed more parallel computation during training, which made it much easier to scale than sequential recurrent models.

5. What is the main long-sequence downside of self-attention?
   • A. It is O(N^2) in sequence length. ✓
   • B. It only works with images.
   • C. It cannot be trained with GPUs.
   • D. It requires handcrafted features.

   Option A is correct. Standard self-attention scales quadratically with sequence length, which is why long-context systems need special serving and attention-efficiency tricks.

6. In scaled dot-product attention, why are the query-key logits divided by the square root of the key dimension d_k before the softmax?
   • A. To reduce the parameter count of the query and key projection matrices.
   • B. Because the dot product of two d_k-dimensional vectors has variance that grows with d_k; without the 1/sqrt(d_k) scale, large logits push the softmax into a saturated region where the gradient to non-dominant keys vanishes. ✓
   • C. Because softmax alone cannot make the attention weights sum to one.
   • D. Because it lowers the asymptotic complexity of attention from O(N^2) to O(N log N).

   Option B is correct. If the components of the query and key are roughly independent with unit variance, their dot product over d_k dimensions has variance on the order of d_k, so its magnitude grows like sqrt(d_k). Feeding large-magnitude logits into the softmax pushes it toward a one-hot distribution: one key dominates, and the gradient with respect to every other key is driven toward zero, so the layer learns slowly or not at all. Dividing by sqrt(d_k) renormalizes the logit variance back to roughly unit scale, keeping the softmax in its informative, well-conditioned range. Option C is wrong — softmax always produces a normalized distribution regardless of scaling; the scale controls sharpness, not normalization. Option D is wrong — the division is an elementwise constant and does not change the O(N^2) cost of forming the N×N score matrix. Option A is wrong — scaling has nothing to do with parameter count. Interview tell: a candidate who can state "variance grows with d_k, softmax saturates, gradients vanish" and name the fix is demonstrating they understand the mechanism, not reciting "for stability."

7. A production LLM server runs near 100% GPU compute while processing the prompt but well below peak while generating tokens. What explains the asymmetry, and what should you optimize during generation?
   • A. Prefill is memory-bound and decode is compute-bound, so optimize raw FLOPs during decode.
   • B. Prefill processes all prompt tokens in parallel (compute-bound, near-peak FLOPs); decode emits one token at a time and must re-read the entire KV cache every step, making it memory-bandwidth-bound — so the lever is reducing KV-cache bytes moved via batching, paged KV caches, GQA, and KV-cache quantization. ✓
   • C. The server is misconfigured; a healthy deployment runs at 100% compute in both phases.
   • D. Decode is slow only because of inter-GPU network latency, not anything about attention.

   Option B is correct. Prefill consumes the whole prompt at once, so it is a large parallel matrix workload that saturates the GPU's arithmetic units — compute-bound. Decode is autoregressive: each new token requires reading every layer's cached keys and values for all prior tokens, doing very little arithmetic per byte loaded. That low arithmetic intensity makes decode memory-bandwidth-bound — the GPU stalls waiting on HBM, not on math. The optimization target during decode is therefore bytes moved, not FLOPs: (1) batching amortizes the fixed cost of streaming the model weights across many concurrent sequences; (2) PagedAttention (vLLM) removes cache fragmentation so larger batches fit; (3) GQA/MQA shrink the per-token KV footprint; (4) KV-cache quantization (e.g. FP8/INT8) halves or quarters the bytes read; (5) speculative decoding amortizes memory traffic over several tokens per verification pass. Option A inverts the two phases. Option C is wrong — the compute gap is intrinsic to autoregressive decode, not a misconfiguration. Option D mistakes a single-GPU memory-bandwidth limit for a networking issue.

8. Llama-2 70B uses grouped-query attention (GQA) with 8 key-value heads alongside 64 query heads. What problem does GQA solve, and how does it relate to MHA and MQA?
   • A. GQA raises accuracy by adding more parameters than multi-head attention.
   • B. GQA shrinks the KV cache (and thus decode memory bandwidth) by sharing each KV head across a group of query heads, sitting between MHA — one KV head per query head, largest cache, best quality — and MQA — a single shared KV head, smallest cache, more quality loss — recovering most of MHA quality at a fraction of the cache. ✓
   • C. GQA removes the need for positional encodings in long-context models.
   • D. GQA reformulates attention so its cost is linear in sequence length.

   Option B is correct. KV-cache size is proportional to the number of key-value heads, and (per the decode analysis) that cache is the memory-bandwidth bottleneck during generation. Multi-head attention (MHA) keeps one KV head per query head — maximum quality but the largest cache. Multi-query attention (MQA) collapses to a single shared KV head — an H-fold cache reduction, but it can cost quality and is sometimes unstable to train. GQA interpolates: partition the query heads into G groups that each share one KV head, giving an H/G-fold reduction (here 64 query heads / 8 KV heads = 8× fewer KV heads than full MHA) while empirically retaining near-MHA quality. This is why GQA is the default in modern open-weight models (Llama 2/3 70B, Mistral, Gemma): it directly attacks the decode bottleneck with minimal quality cost. Option A is wrong — GQA reduces KV parameters/cache, it does not add them. Options C and D describe positional encodings and linear-attention variants, which are unrelated mechanisms.

9. Why did modern LLMs largely move from sinusoidal/learned absolute positional encodings to rotary positional embeddings (RoPE), and how do long-context methods like YaRN relate?
   • A. RoPE removes the need for any positional signal by making attention permutation-invariant.
   • B. RoPE encodes relative position by rotating the query and key vectors by an angle proportional to their position, so relative offsets appear directly inside the QK dot product; this extrapolates to longer contexts better than absolute schemes, and methods like NTK-aware scaling, position interpolation, and YaRN rescale RoPE frequencies to extend the usable context beyond the trained length with little or no fine-tuning. ✓
   • C. RoPE is a learned lookup table that, like learned absolute embeddings, hard-caps at the training-time maximum sequence length.
   • D. RoPE is specific to encoder-only models such as BERT and is not used in decoder-only LLMs.

   Option B is correct. RoPE applies a position-dependent rotation to the query and key vectors before the dot product. Because rotating both by angles proportional to their absolute positions leaves the result depending only on their difference, relative position is baked into the attention score itself — without adding a separate positional vector. This relative formulation extrapolates more gracefully than additive sinusoidal encodings (information-limited at high frequency) or learned absolute embeddings (which simply have no entries past the trained max length). Long-context extension exploits RoPE's frequency structure: position interpolation squeezes positions into the trained range, NTK-aware scaling adjusts the rotation base so high-frequency components are not over-compressed, and YaRN combines interpolation with attention-temperature correction to push context length up by large factors with minimal fine-tuning. Option A is backwards — permutation-invariance is the problem positional encodings exist to fix. Option C describes learned absolute embeddings. Option D is wrong — RoPE is the de-facto standard in decoder-only LLMs (Llama, Mistral, Qwen).

10. A teammate says "FlashAttention makes attention O(N) instead of O(N^2)." What is the precise correction?
    • A. They are right — FlashAttention changes the algorithmic complexity of attention to linear.
    • B. FlashAttention keeps the same O(N^2) arithmetic complexity but is IO-aware: it tiles the computation and computes the softmax online in fast SRAM, never materializing the full N×N attention matrix in slow HBM — so the win is wall-clock time and memory (O(N) memory, far fewer HBM reads/writes), not fewer FLOPs. Sub-quadratic compute requires sparse or linear-attention variants instead. ✓
    • C. FlashAttention reduces complexity by pruning low-weight attention heads at runtime.
    • D. FlashAttention is a CPU-only kernel and does not affect GPU attention.

    Option B is correct. The distinction is between arithmetic complexity (number of multiply-adds) and IO complexity (bytes moved across the memory hierarchy). FlashAttention does not change the former — it still computes all N^2 query-key interactions. What it changes is the latter: it processes attention in tiles that fit in on-chip SRAM and uses an online (streaming) softmax so it never writes the full N×N score matrix to HBM. That drops attention memory from O(N^2) to O(N) and slashes HBM traffic, which is the actual bottleneck on modern GPUs — yielding large wall-clock speedups (and enabling longer contexts) with bit-for-bit exact results. If you genuinely need sub-quadratic compute, you must change the algorithm: sparse attention (Longformer, BigBird), low-rank/kernelized linear attention (Performer, Linformer), or state-space models (Mamba). Options C and D are simply false descriptions of what FlashAttention does.

11. You need (a) high-quality sentence embeddings for semantic search and (b) open-ended text generation. Which transformer family fits each and why?
    • A. Use a decoder-only model for both, since causal attention is strictly superior.
    • B. Use an encoder-only model (bidirectional attention, BERT-style) for embeddings/classification because every token sees full left-and-right context; use a decoder-only model (causal attention, GPT/Llama-style) for generation because next-token training must mask future tokens. Encoder-decoder (T5/BART) fits seq2seq tasks like translation, where an autoregressive decoder cross-attends to a bidirectional encoding of the input. ✓
    • C. Encoder-only models are required for generation and decoder-only models for embeddings.
    • D. Only encoder-decoder models can perform either task.

    Option B is correct. Pure representation tasks (embeddings, classification, span extraction) benefit from bidirectional attention: a token's encoding can incorporate both left and right context, which is exactly what masked-language-model pretraining (BERT/RoBERTa) optimizes. Generation requires causal masking — at training time a token must not see its own future, or next-token prediction degenerates — which is the decoder-only design (GPT, Llama, Mistral). Encoder-decoder models keep both: a bidirectional encoder builds a rich input representation and an autoregressive decoder cross-attends to it, which is the natural fit for sequence-to-sequence problems such as translation and summarization (T5, BART). Nuance worth stating in an interview: decoder-only models can be adapted to produce embeddings (last-token or mean pooling, contrastive fine-tuning), and several strong modern embedding models are decoder-based — but classically, and for parameter-matched quality on pure representation, bidirectional encoders remain the cleaner choice. Option A overgeneralizes; options C and D invert or overclaim the families' roles.

Frequently asked questions

Why does the transformer use scaled dot-product attention instead of plain dot-product attention?

Without the 1/sqrt(d_k) scaling, the dot products grow in magnitude with the embedding dimension d_k because they sum d_k independent products. Large pre-softmax values push the softmax into a saturated region where one logit dominates and the gradient w.r.t. the others vanishes — training stalls. Dividing by sqrt(d_k) keeps the variance of the dot product approximately unit, the softmax stays in its informative range, and gradients flow. Empirically, removing the scaling on a 512-dim model degrades both convergence speed and final perplexity.

What is the difference between encoder-only, decoder-only, and encoder-decoder transformers?

Encoder-only models (BERT, RoBERTa) use bidirectional self-attention — every token sees every other token — and are trained with masked-language-modeling objectives. They excel at classification and span-extraction tasks. Decoder-only models (GPT family, Llama, Mistral) use causal self-attention where each token attends only to prior tokens, are trained with next-token prediction, and are the dominant architecture for generative tasks. Encoder-decoder models (T5, BART, original transformer) combine both — the encoder builds a bidirectional representation of the input and the decoder cross-attends to it autoregressively. Cross-attention is what distinguishes the encoder-decoder from a stacked monolithic decoder.

Why are positional encodings necessary, and what are the tradeoffs between sinusoidal, learned, and rotary positional encodings (RoPE)?

Self-attention is permutation-equivariant (permuting the input tokens permutes the outputs identically, so the layer carries no inherent notion of order) — without positional information the model cannot distinguish "dog bites man" from "man bites dog". Sinusoidal encodings are deterministic functions of position and generalize trivially to longer sequences but are additive and information-limited at high frequencies. Learned absolute positional embeddings are flexible but cap out at the training-time max length. RoPE rotates query and key vectors by an angle proportional to position, which encodes relative position multiplicatively inside the dot product, generalizes better to longer contexts, and is the de-facto standard in modern open-weights LLMs. ALiBi is the other common choice for long-context efficiency.

What is the computational complexity of self-attention and why is this a problem for long sequences?

Self-attention is O(N^2) in both time and memory in the sequence length N because every token attends to every other token, producing an N×N attention matrix. At N=128k context, the attention matrix alone is 16 billion entries per head — infeasible to materialize naively. The fix is one of three families: sparse attention (Longformer, BigBird) restricts each token to attend to a local window plus a few global tokens; linear-attention variants (Performer, Linformer, Mamba/SSM) reformulate attention to scale linearly; and IO-aware exact methods (FlashAttention, FlashAttention-2/3) keep complexity O(N^2) but tile the computation to avoid materializing the full matrix in HBM, dramatically reducing wall-clock time and memory.

How does multi-head attention differ from single-head, and what does each head learn?

Multi-head attention splits the model dimension into h heads, runs scaled dot-product attention independently on each, and concatenates the results — at the same parameter count and FLOP budget as a single full-dimension head. The benefit is representational: different heads can specialize in different relations (one head tracks subject-verb agreement, another tracks coreference, another tracks positional offsets). Probing studies on BERT and GPT-style models show some heads are clearly interpretable (induction heads, copy heads, syntactic-role heads) while others are noisy. In production, head pruning often removes 20-40% of heads with minimal quality loss, suggesting redundancy.

Why do transformers use layer normalization, and what is the difference between pre-LN and post-LN?

Layer normalization stabilizes activations by normalizing across the feature dimension at each token position, which keeps gradients well-conditioned through deep stacks. The original transformer used post-LN (LayerNorm applied after the residual connection), which trains well at moderate depth but requires careful learning-rate warmup at large scale because the residual path can dominate early. Pre-LN (LayerNorm before each sublayer, residual added unmodified) is the modern default — it is easier to train deep models with, removes the need for warmup in many cases, and is what GPT-style and Llama-style models use. RMSNorm replaces LayerNorm with a cheaper variant that omits the mean-centering and is now standard in open-weights LLMs.

How does KV caching work, and why is prefill so much faster than decode?

During autoregressive decoding, the keys and values of all prior tokens are reused at every step — recomputing them is wasted work. KV caching stores K and V tensors per layer per head and grows them by one token per decode step, so each new token only needs to compute its own Q and attend over the cached K/V. Prefill processes the entire prompt in parallel (compute-bound, near-peak FLOPs); decode processes one token at a time (memory-bandwidth-bound, because the entire KV cache must be read every step). This asymmetry is why batched serving, paged KV caches (vLLM PagedAttention), and speculative decoding are the dominant optimizations in inference systems.

References & standards

• Vaswani et al. — Attention Is All You Need (the original paper that introduced the transformer) — NeurIPS 2017 (Google)
• Su et al. — RoFormer: Enhanced Transformer with Rotary Position Embedding (RoPE) — arXiv 2021
• Dao et al. — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — NeurIPS 2022

Related topics