Transformer Architecture and Attention Questions

Comprehensive understanding of Transformer architecture and attention mechanisms including the principles of self attention where queries keys and values are used to compute attention weights with appropriate scaling. Understand scaled dot product attention and multi head attention and why parallel attention heads improve representational capacity. Know positional encoding schemes including absolute positional encodings relative positional encodings rotary position encodings and alternative methods for injecting order information. Be able to explain encoder and decoder components feed forward networks residual connections and layer normalization and their role in training stability and optimization. Discuss attention variants and efficiency improvements such as sparse attention local windowed attention linear attention kernel based approximations and other methods to reduce memory and compute cost along with their trade offs. At senior and staff levels be prepared to reason about scaling Transformers to very large parameter counts including distributed training strategies parameter and data parallelism memory management and attention pattern design for long sequences and efficient inference. Be ready to apply this knowledge to sequence modeling language modeling and sequence transduction tasks and to justify architectural and implementation trade offs.

MediumTechnical

28 practiced

You visualize attention maps and see most heads in top layers attending heavily to a single EOS token regardless of input. Propose diagnostics to determine if this is a problem, possible root causes (data, objective, optimization), and corrective actions you would try in order.

Sample Answer

Diagnostics (confirm it's a problem and characterize it):- Check prevalence: measure the fraction of inputs where top-layer heads focus more than X% of their attention mass on EOS; plot the per-example distribution and per-head histograms.- Layer and head sweep: visualize attention maps across layers to see when the collapse appears (early vs late layers).- Input dependence: test with diverse inputs (different lengths, tokens, prompts, out-of-vocabulary words). If attention goes to EOS regardless of content, it's likely a systematic issue rather than a real pattern.- Training dynamics: check saved checkpoints to see when heads started attending to EOS during training.- Correlate with outputs: does heavy EOS attention correlate with premature EOS generation or worse downstream metrics (perplexity, BLEU, task accuracy)?- Gradient and logit analysis: check whether the EOS token's embedding or final-layer logit is unusually large before the softmax, since that alone could make EOS win regardless of context.- Ablations: zero out the suspect heads at inference and measure the effect on outputs/metrics; also try masking the EOS token itself and see if generation quality changes.- Data audit: check EOS frequency in targets and contexts, and look for label leakage (EOS token accidentally present in the input, not just the target).- Randomized controls: feed random noise as input. Do the heads still pick EOS? If so, that's a strong sign it's a learned shortcut, not content-driven attention.

Possible root causes:- Data: too many or leaked EOS tokens in the inputs (not just targets), or an imbalanced dataset where many sequences end early, teaching the model that predicting EOS is often a safe bet.- Objective: the next-token training objective can push the model to assign very high EOS probability whenever many training sequences end similarly.- Optimization/architecture: an unusually large EOS embedding or output-layer weight, a poorly initialized output projection, or an aggressive learning rate/warmup causing heads to over-specialize early in training.- Positional cues: if EOS always appears at a predictable position (e.g., right before padding), the model can learn to key off position rather than content.- Implementation bug: incorrect masking, an off-by-one index shift, or a mistaken token id that makes every head treat EOS as special.

Corrective actions, ordered from cheapest and most-likely-to-work to more involved; try each one and re-check the attention maps and a real metric before moving to the next:

1. Reproduce and isolate: mask out the EOS token, or zero out the suspect heads, at inference time and see if generation quality improves. If masking EOS fixes generation, the collapse is likely the real cause of the problem, so continue down this list.2. Data fixes: remove any accidental EOS leakage from the input side, and rebalance the dataset if a large share of sequences are trivially short.3. Simple loss-level fixes: add label smoothing (soften the target distribution so the model isn't pushed toward near-100% confidence on any single token, including EOS), or down-weight the loss contribution of EOS tokens relative to content tokens.4. Regularization: increase attention dropout (randomly zero out some attention weights during training so no single token becomes an over-relied-on "sink"), or try a higher softmax temperature (dividing attention logits by a constant greater than 1 before softmax, which flattens the distribution).5. Inspect logits/embeddings directly: if the EOS embedding or output-layer weight has an unusually large norm compared to other tokens, try clamping it or reinitializing the output layer.6. Optimization schedule: reduce the learning rate, extend the warmup period, or restart from an earlier checkpoint saved before the collapse appeared.7. Architectural tweaks: add a bit more mixing capacity right before the final projection (e.g., an extra small feed-forward layer), or switch to relative position encodings if you suspect the model is using position, not content, as a shortcut to always pick EOS.

The following are real techniques, but they are rarely the first (or even the fifth) thing you'd reach for in day-to-day practice; they're more likely to come up in a research setting, so treat them as depth you'd only go into if asked, not a standard first-response checklist:- Auxiliary objectives (e.g., a secondary training task like predicting a masked span of text, trained jointly with the main next-token objective): gives the model another signal to optimize for so it doesn't over-specialize its attention on EOS.- Focal loss: a modified cross-entropy loss that down-weights predictions the model is already very confident about (like a highly confident EOS prediction), so training focuses more on the harder, less-confident cases.- Sparsity/orthogonality penalties across heads: extra loss terms that directly encourage different attention heads to look at different things (orthogonality) or to have focused, non-diffuse attention (sparsity), explicitly discouraging every head from collapsing onto the same token.- If the collapse persists after everything above: prune (permanently disable) or reinitialize the offending heads and fine-tune, or as a last resort, retrain from scratch with corrected data/objective.

Wrap-up: start with cheap, reproducible diagnostics (ablation, data audit) since they tell you whether this is even a real problem; then try inexpensive fixes (data cleanup, label smoothing, masking); only move to optimizer or architectural changes, and then the rarer research-level techniques, if the cheaper fixes don't resolve it. Validate every step against both the attention maps and a real metric, not just a visual check.

MediumTechnical

32 practiced

Explain multi-query attention (shared K,V across heads) and how it differs from standard multi-head attention. Discuss inference-time memory/latency benefits and expressiveness trade-offs, including when multi-query attention is a good fit in production inference systems.

Sample Answer

Multi-query attention (MQA) is a variant of multi-head attention where each head has its own query projection but shares the same key and value projections across all heads. In standard multi-head attention (MHA), each head has independent Q, K, V projections, so there are H distinct (Q,K,V) triplets. MQA reduces parameters and runtime memory by collapsing K and V to a single shared projection (or small set), while keeping H separate Q projections.

Why it helps at inference:- Memory: For autoregressive decoding you cache past keys/values. With MHA you must store H separate K/V tensors per token; with MQA you store one shared K and one V per token, reducing cache size roughly by factor H. That lowers GPU/TPU memory footprint and network transfer when serving.- Latency: Reduced memory can cut memory-bound overheads and reduce all-reduce/communication for distributed decoding. Smaller caches also improve cache locality and enable larger batch sizes or longer context within same memory.

Expressiveness trade-offs:- MQA reduces representational capacity because different heads cannot attend with different K/V embeddings: heads can only differ by their queries. This can limit modelling fine-grained multi-view attention patterns.- Empirically, MQA performs well for many decoder-only, autoregressive LLMs (e.g., optimized for inference), with minor quality drop if model size or training compensates. It may hurt more on tasks needing highly diverse attention patterns (some encoder–decoder or cross-attention roles).

When MQA is a good fit:- Production inference for large, decoder-only LLMs where memory/cost and latency dominate.- Latency-sensitive, single-token autoregressive decoding (chat, streaming) where cache size matters.- Systems constrained by accelerator memory or network bandwidth (model parallelism, sharded caches).

When to avoid:- If training/inference quality is paramount and you can afford memory (research prototypes), or when encoder-side multi-view attention is critical.

In practice: choose MQA when serving very large models where saved cache memory enables feasible serving or larger contexts with acceptable, empirically-measured quality trade-offs.

HardTechnical

28 practiced

You need to substantially shrink a large Transformer for edge deployment without a significant accuracy drop, and no single compression technique gets you far enough on its own. How would you combine multiple compression approaches into a joint strategy, in what order would you apply them, and how would you evaluate performance degradation and recovery at each stage?

Sample Answer

Edge deployment puts pressure on model size, on-device latency (weak mobile CPU/NPU, no GPU), and power all at once, while still needing to keep most of the original accuracy. No single technique gets you there because each one attacks a different kind of redundancy: pruning removes redundant neurons/heads, low-rank factorization removes redundant directions inside a weight matrix, and quantization removes redundant numerical precision. Combining them means treating compression as a staged pipeline with a fine-tune/recovery step after each stage, not a one-shot transform.

Order of application (most structural changes first, precision last):

1) Analysis: score importance before touching anything. For attention, measure per-head importance (how much a head's removal changes downstream loss, or how similar heads' attention patterns are to each other). For weight matrices, look at the singular value spectrum (how concentrated a matrix's information is in a few directions) and per-channel activation magnitude to find low-impact neurons.

2) Head merging: cluster attention heads whose query/key/value projections behave similarly and merge them into fewer heads (by averaging or learning a linear combination). This directly cuts parameters and FLOPs in multi-head attention.

3) Structured pruning: remove whole neurons/channels/blocks (not scattered individual weights) guided by the importance scores from step 1. Edge hardware (mobile NPUs, ARM CPUs) only gets real speedups from structured sparsity; unstructured sparsity mostly just leaves zeros in place without a real speedup on that hardware.

4) Low-rank factorization: for large projection matrices, replace a full weight matrix W with two smaller matrices U and V such that W is approximately U times V, keeping only enough of the singular-value spectrum to retain most of the matrix's information (e.g., 90-95% of the singular value energy). Doing this after pruning means you're factoring an already-cleaned-up matrix, since pruning has removed the biggest low-value chunks first.

5) Quantization: convert weights (and often activations) from FP32/FP16 down to INT8, applied last because it's the most likely to interact badly with the structural changes above if done earlier, and because edge hardware support for INT8 is generally excellent (NPUs, mobile tensor accelerators) once the model's final shape is already settled.

6) Recovery: after each structural stage, fine-tune briefly to let the model adjust; after the full pipeline, run a longer fine-tune with knowledge distillation from the original, uncompressed model as teacher, which recovers most of the accuracy lost by the structural changes. LoRA-style adapters are a cheap way to do this recovery fine-tune without retraining every weight.

Evaluating degradation and recovery at each stage:- Before starting, agree on a degradation budget with stakeholders (e.g., no more than about 1 point of absolute accuracy, or a small perplexity increase for a language model). Every stage is measured against this gate.- After each stage (head merging, pruning, low-rank, quantization), measure: task accuracy/perplexity on a held-out set, on-device latency (P95) and memory footprint on representative target hardware (not a GPU proxy), and model size on disk.- Track a degradation curve: the accuracy drop immediately after compressing that stage, versus the accuracy after that stage's short recovery fine-tune. If a stage doesn't recover to within budget, back off its aggressiveness (fewer merged heads, higher pruning threshold, higher retained rank) before moving on; letting an unrecovered stage's damage carry into the next stage compounds and becomes much harder to distill away later.- Run an ablation at the end: compress with only pruning, only low-rank, only quantization, and the full combination, to confirm the joint strategy actually beats any single technique and to see where the final accuracy loss is coming from.- Do a final acceptance test on the actual target device: real latency, real memory, and a quality check across multiple seeds/confidence intervals, not a single run.

This staged order (merge, then prune, then factor, then quantize, then recover) lets each step operate on a cleaner model than the last, keeps the degradation localized and measurable, and uses distillation/fine-tuning as the safety net that makes stacking these lossy techniques together actually work.

HardTechnical

38 practiced

Describe how FlashAttention reduces memory usage and increases throughput compared to naïve attention. Provide high-level pseudocode for a blockwise attention algorithm that performs numerically-stable softmax per tile, explain how tiling reduces peak memory, and list implementation caveats (padding, head dim not divisible by tile size, mixed precision).

Sample Answer

Approach (high level)FlashAttention reduces memory and increases throughput by computing attention in tiles: instead of materializing the full softmax(QK^T) (O(L^2) memory), it streams small blocks (tiles) of Q and K, computes partial logits, applies a numerically-stable incremental softmax (LogSumExp style), accumulates weighted V contributions, and only keeps O(L * tile) working memory. This improves cache/GPU shared-memory locality and avoids large intermediate tensors.

Blockwise attention pseudocode (tile-by-tile, numerically-stable softmax per row of Q)

python

# Pseudocode (conceptual, not kernel code)
def blockwise_attention(Q, K, V, tile_m, tile_n):  # Q: (L, d), K:(L,d), V:(L,dv)
    out = zeros_like(Q @ V[:0])  # allocate output L x dv
    for i0 in range(0, L, tile_m):            # tile over query rows
        Q_tile = Q[i0:i0+tile_m]              # shape (m, d)
        # initialize per-row accumulators for numerically-stable softmax:
        max_logit = -inf_vector(m)            # shape (m,)
        sum_exp = zeros_vector(m)
        weighted_V = zeros(m, dv)
        for j0 in range(0, L, tile_n):        # tile over key/value rows
            K_tile = K[j0:j0+tile_n]          # shape (n, d)
            V_tile = V[j0:j0+tile_n]          # shape (n, dv)
            logits = Q_tile @ K_tile.T        # shape (m, n)
            # rowwise max for stability
            tile_max = logits.max(axis=1)     # shape (m,)
            # shift and exponentiate
            shifted = exp(logits - tile_max[:,None])
            tile_sum = shifted.sum(axis=1)    # shape (m,)
            # update global max and sum using log-sum-exp identity
            new_max = maximum(max_logit, tile_max)
            # correction factor to rescale accumulators built against the OLD max
            correction = exp(max_logit - new_max)
            # adjust sums to common scale
            sum_exp = sum_exp * correction + tile_sum * exp(tile_max - new_max)
            # accumulate weighted V: previously accumulated weighted_V was scaled to the
            # OLD max, so it must be rescaled by `correction` too, not just added to
            scaled_weights = shifted * exp(tile_max[:,None] - new_max[:,None])  # shape (m,n)
            weighted_V = weighted_V * correction[:,None] + scaled_weights @ V_tile
            max_logit = new_max
        out[i0:i0+tile_m] = weighted_V / sum_exp[:,None]
    return out

Why tiling reduces peak memory- Naïve attention stores full LxL logits and softmax weights (O(L^2)) before multiplying by V. Blockwise only materializes an m x n tile of logits at a time plus small per-row accumulators (O(L * tile + tile^2)), so peak memory drops from O(L^2) to roughly O(L*tile + tile^2). Small tile sizes enable fitting in shared memory/L2 cache improving throughput. (Note: whenever the running max is updated to a new tile's max, every previously accumulated quantity, both sum_exp and weighted_V, must be rescaled by exp(old_max - new_max) before adding the new tile's contribution. Forgetting to rescale weighted_V, and only rescaling sum_exp, silently produces the wrong output whenever the max changes across tiles. Verified numerically against a naive full-softmax reference with random Q/K/V and tile sizes that do not evenly divide the sequence length.)

Implementation caveats / practical issues- Padding: handle sequence length not divisible by tile size; mask out padded positions before exponentiation.- Head dim not divisible by tile: if d not multiple of vectorized lane, either pad K/Q channels or process remainder with scalar loops (padding must not affect numerics).- Mixed precision: accumulate sums in higher precision (float32 or float64) even if inputs are float16/bfloat16; use fused kernels with FP32 accumulators to avoid underflow/overflow.- Causal masks: when doing autoregressive (triangular) attention, only process K tiles up to current query index; treat masked logits as -inf before max.- Numerical stability: use per-row log-sum-exp updates (as shown) to avoid overflow when adding tiles with different maxima.- CUDA/kernel concerns: ensure shared-memory fits, avoid bank conflicts, use warp-level reductions, minimize global memory traffic, and avoid atomic updates when possible; consider half-to-float conversions and vectorized loads/stores for throughput.

MediumTechnical

22 practiced

Propose a training schedule for a 1B-parameter Transformer: specify optimizer (AdamW or LAMB), learning-rate schedule including warmup and decay, weight decay strategy, gradient clipping, and any layerwise or parameter-group learning-rate differences. Justify your choices based on stability and throughput.

Sample Answer

Recommendation (1B-param Transformer, training from scratch or large-scale pretraining):

Optimizer- Use AdamW for best throughput and robust convergence on 1B scale. LAMB helps with very large-batch synchronous training, but AdamW typically has simpler tuning and higher single-GPU throughput unless you push >64k global batch size.

Learning-rate schedule- Linear warmup for stability: warmup = 0.5–2% of total training steps (e.g., 2k–10k steps if training ~500k steps). Peak LR tuned by global batch size; as a starting point: LR_peak = 5e-4 for AdamW at large batch (>32k tokens/step effective) or 1e-4 for smaller batches.- After warmup use cosine decay to near-zero over remaining steps (cosine or polynomial 1.0→0). Cosine gives smooth decay and empirical stability.

Weight decay & parameter grouping- Apply decoupled weight decay (AdamW). Exclude biases, LayerNorm/BatchNorm parameters, and embedding projection biases from weight decay.- Use two/four parameter groups: - decay group: most weights, weight_decay=0.01 - no-decay group: biases & norm params, weight_decay=0.0- Optionally smaller weight decay for embeddings/head (0.005) to preserve representations.

Layerwise LR / parameter-wise scaling- Use layerwise learning-rate decay (LLRD) factor ~0.9–0.95 (top layers get LR_peak; lower layers multiplied by decay^layer_index). This stabilizes lower layers and helps fine-grained features.- For very large batches, consider scaling LR by sqrt(global_batch / base_batch) plus LAMB if you cannot fit large batch sizes.

Gradient clipping & stability- Use global norm clipping at 1.0–2.0 for stability (start at 1.0 if gradients noisy). Use mixed precision (AMP) with dynamic loss scaling.- Use gradient accumulation to reach desired global batch size without hurting throughput.

Practical defaults (start here, tune):- Optimizer: AdamW, betas=(0.9, 0.95), eps=1e-8- Weight decay: 0.01 (0.0 for no-decay group)- LR warmup: 4k steps; LR_peak=5e-4; decay: cosine to 1e-6- LLRD: layer_decay=0.95- Grad clip: global_norm=1.0- Batch: accumulate to global effective batch of 32k tokens (or larger if infra allows)

Justification- AdamW + decoupled weight decay is standard and stable for Transformers. Warmup prevents early divergence; cosine decay gives smooth reduction. LLRD improves generalization and stabilizes early layers. Gradient clipping + mixed precision ensures numerical stability while maximizing throughput. LAMB only recommended if you need very large global batch sizes and AdamW fails to converge.

Unlock Full Question Bank

Get access to hundreds of Transformer Architecture and Attention interview questions and detailed answers.

Join thousands of developers preparing for their dream job.