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

47 practiced

You're training a 10B-parameter Transformer. Select an optimizer (AdamW, LAMB, Adafactor), a learning rate schedule (warmup, decay), and core hyperparameters (weight decay, betas). Justify your choices and explain how you'd tune them to stabilize large-batch distributed training.

Sample Answer

Choice and justification:- Optimizer: Adafactor. For a 10B-parameter Transformer, memory and optimizer state overhead are major constraints. Adafactor gives near-AdamW performance while storing only factored second-moment estimates (much lower memory), enabling larger micro-batches or fewer gradient-accumulation steps. If memory is abundant and you need aggressive large-batch scaling, consider LAMB as an alternative (better at scaling to very large global batch sizes).

Learning-rate schedule:- Linear warmup for 2–10k steps (depending on dataset size), then cosine decay or inverse-sqrt decay. Warmup avoids early instabilities; cosine gives smooth final convergence, inverse-sqrt works well for very long training.

Core hyperparameters:- Weight decay (decoupled): 1e-2 (tune 1e-3–1e-1). Apply to non-bias, non-layernorm weights.- For Adafactor: set multiply_by_parameter_scale=True, clipping_threshold=1.0, decay_rate=0.8–0.999 (tune), epsilon=1e-30 default. If using AdamW: betas=(0.9,0.98) and eps=1e-6 (more stable than 1e-8).- Learning rate: base LR per-parameter such that LR_effective = LR * sqrt(batch_size) heuristic is a starting point; but do an LR sweep.

Stabilizing large-batch distributed training (practical tuning steps):1. Start conservative: small LR, long warmup (warmup steps proportional to total steps, e.g., 0.5–2% of training steps).2. Gradient clipping: global-norm clip at 1.0–2.0 to avoid spikes.3. Mixed precision + dynamic loss scaling to save memory and maintain numeric stability.4. Scale LR with batch size carefully: run short LR sweep (log-uniform) at target global batch size; monitor loss scale and train/val gap.5. Use gradient accumulation to emulate desired batch size while testing stability.6. Monitor gradient variance / gradient noise scale; if noise low (too deterministic), reduce LR or add small weight decay/dropout.7. If divergence occurs, increase warmup length, reduce LR or increase eps in optimizer.8. Consider layer-wise LR decay (0.8–0.95) for transfer/fine-tuning.

Validation strategy:- Do small-scale reproducibility runs (few steps) across different seeds and hardware, verify step-level loss curves before full run. Use checkpointing and early rollback when divergence detected.

This setup prioritizes memory efficiency and robust scaling; tune LR and warmup first, then weight decay and clipping.

MediumTechnical

29 practiced

You need to implement rotary positional embeddings (RoPE) in a Transformer training pipeline using PyTorch. Describe the mathematical operation, exactly how you apply RoPE to Q and K tensors (shapes and index mapping), how to implement it efficiently for batched tensors, and how to handle inference when sequences are longer than training sequences.

Sample Answer

RoPE idea (math): for each head vector element pair (x_{2i}, x_{2i+1}) at position p, apply a rotation by angle θ_{p,i}: [x'_{2i}, x'_{2i+1}] = R(θ_{p,i}) · [x_{2i}, x_{2i+1}] where R = [[cos, -sin],[sin, cos]] and θ_{p,i} = p / (10000^{2i/d_head}). This multiplies Q and K by the same position-dependent rotation so their dot-product encodes relative positions.

Shapes & index mapping:- Q, K: (B, H, S, D) where D = d_head. RoPE pairs dims: index pairs (0,1), (2,3), ..., (D-2,D-1). For pair i use frequency ω_i = 1 / (10000^{2i/D}).

Efficient batched PyTorch implementation (vectorized):

python

import torch
def apply_rope(q, k, max_seq_len=2048, base=10000.0):
    # q,k: (B, H, S, D)
    B,H,S,D = q.shape
    device = q.device
    assert D % 2 == 0
    dim_half = D // 2
    # prepare freqs: (D/2,)
    inv_freq = 1.0 / (base ** (torch.arange(0, dim_half, device=device) * 2.0 / D))
    # positions: (S,)
    positions = torch.arange(S, device=device).type_as(inv_freq)
    # angles: (S, D/2) -> broadcastable
    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)  # (S, D/2)
    cos = torch.cos(angles).unsqueeze(0).unsqueeze(0)  # (1,1,S,D/2)
    sin = torch.sin(angles).unsqueeze(0).unsqueeze(0)
    # reshape q/k to pair format (B,H,S,D/2,2)
    q_ = q.view(B,H,S,dim_half,2)
    k_ = k.view(B,H,S,dim_half,2)
    # apply rotation: [x,y] -> [x*cos - y*sin, x*sin + y*cos]
    q_rot = torch.empty_like(q_)
    q_rot[...,0] = q_[...,0]*cos - q_[...,1]*sin
    q_rot[...,1] = q_[...,0]*sin + q_[...,1]*cos
    k_rot = torch.empty_like(k_)
    k_rot[...,0] = k_[...,0]*cos - k_[...,1]*sin
    k_rot[...,1] = k_[...,0]*sin + k_[...,1]*cos
    return q_rot.view(B,H,S,D), k_rot.view(B,H,S,D)

Key points:- All broadcasting; no loops; memory O(B*H*S*D).- Use half-precision carefully; precompute cos/sin for training seq length.

Handling longer inference sequences:- Precompute cos/sin for maximum expected inference length (or compute on the fly) and apply to cached key/value positions when doing autoregressive caching. RoPE extrapolates naturally since θ_p = p * ω_i; as long as you can compute cos/sin for new p you can rotate Q/K consistently. If training used rotary with limited S_train, you must ensure ω_i uses base (not learned) so angles remain valid out-of-distribution. Optionally, scale frequencies or use ALiBi for explicit longer-range bias if needed.

Edge cases & tips:- D must be even; pairwise mapping is crucial.- For perf, precompute cos/sin once per forward (or cache per device) and reuse across heads.

EasyTechnical

30 practiced

Describe how causal (autoregressive) masking is implemented in attention for decoder-only models. Provide the mathematical effect on attention logits and explain how it enables autoregressive generation during inference and training.

MediumTechnical

29 practiced

You notice many attention heads appear to contribute little to the model output. Describe methods to prune attention heads during training or after training to reduce inference cost. Explain how to measure head importance, step-by-step pruning strategies, required retraining or distillation, and potential pitfalls (for example, information transfer to remaining heads).

Sample Answer

Approach summary: treat head pruning as structured model compression—measure per-head importance, remove low-importance heads with an iterative strategy, then recover quality via fine-tuning or distillation while watching for information re-allocation and distributional effects.

How to measure head importance- Magnitude / L1/L2 of head projection matrices (Wq, Wk, Wv, Wo) — cheap proxy.- Attention output norm: average ||head_output|| across dataset — directly reflects contribution.- Gradient-based saliency (Taylor expansion): importance ≈ ∑ parameters * gradient — estimates loss change if head removed.- Leave-one-out ablation: zero out head at inference on validation set and measure Δloss/Δmetric — most reliable but costly.- Learned gates: add a scalar gate g_h per head and train with L1 or L0 regularizer; gates converge to indicate importance.

Step-by-step pruning strategies1. Baseline: evaluate full model and compute head importance on representative data.2. One-shot pruning: remove all heads below threshold (fast, higher risk).3. Iterative pruning (recommended): remove small fraction (e.g., 5–20%) of lowest-importance heads, then: - Fine-tune on original objective for a few epochs. - Recompute importances; repeat until target inference cost or metric drop reached.4. Structured retraining: after pruning, reinitialize or re-scale remaining head outputs (layernorm/gain) and fine-tune.5. Distillation: optionally distill from full model to pruned student — use logits and intermediate attention distributions to transfer behavior.

Required retraining / distillation- Minimal: fine-tune pruned model on original training data; often recovers most performance.- Stronger: knowledge distillation (teacher-student) speeds recovery and preserves accuracy, helpful if many heads removed.- If using learned gates during training, you can train with sparsity prior from scratch (global objective) to get pruned model without expensive post-hoc steps.

Potential pitfalls and mitigations- Information transfer: remaining heads can re-learn functionality, hiding true importance. Mitigate by re-evaluating importance after small pruning steps and using ablation tests.- Layer interactions: some layers are more sensitive; prune per-layer quotas or use per-layer thresholds.- Distribution shift: importance computed on one dataset may not generalize — use representative validation sets.- Calibration and latency trade-offs: smaller models may become overconfident; recalibrate post-pruning.- Implementation overhead: ensure pruning yields actual speedup (fused kernels, reduced dims). Simply zeroing weights may not reduce latency.- Stability: pruning many heads at once risks catastrophic degradation; iterative + distillation reduces risk.

Practical tips- Prefer iterative pruning with validation-aware stopping.- Use learned gates + L0 regularization if you can afford extra training.- Measure real inference latency and memory—not just FLOPs.- Log per-head behavior (attention maps) before/after pruning to detect semantic shifts.

MediumTechnical

32 practiced

Explain how to implement relative positional encodings in attention (for example, Shaw et al. or T5-style relative biases). Provide pseudocode showing how relative distances are converted to bias terms and added to attention logits, and discuss memory implications for long sequences and clipping strategies.

Sample Answer

To implement relative positional encodings as a bias added to attention logits, map pairwise relative distances (j - i) to a small learnable table of biases and add them to the Q·K^T logits before softmax. Two common styles: Shaw et al. (full relative positional embeddings added to keys/values) and T5-style relative biases (a small lookup table indexed by clipped distance added to logits). Below is concise pseudocode for T5-style:

python

# Pseudocode (Python-like)
# seq_len: length of sequence
# num_heads: attention heads
# max_rel: maximum relative distance to store (e.g., 128)
# rpb: learnable tensor shape (num_heads, 2*max_rel+1)

def make_relative_index_matrix(seq_len, max_rel):
    # produce matrix of relative distances clipped to [-max_rel, max_rel]
    idx = [[j - i for j in range(seq_len)] for i in range(seq_len)]
    # clip and shift to 0..2*max_rel
    clipped = [[min(max(d, -max_rel), max_rel) + max_rel for d in row] for row in idx]
    return clipped  # shape (seq_len, seq_len)

# During init
rel_index = make_relative_index_matrix(seq_len, max_rel)  # or compute on the fly
# rpb initialized randomly, shape (num_heads, 2*max_rel+1)

# During forward
# logits = Q @ K.T  # shape (batch, num_heads, seq_len, seq_len)
bias = rpb[:, rel_index]  # broadcasting: (num_heads, seq_len, seq_len)
logits = logits + bias.unsqueeze(0)  # add to logits per head
attn = softmax(logits)

Key points and reasoning:- T5-style biases are tiny (num_heads × (2*max_rel+1)) and applied per-head, avoiding increasing Q/K dimensionality.- Clipping distances bounds memory and allows parameter sharing for far-away tokens.Memory implications and strategies:- Precomputing an (seq_len, seq_len) integer index matrix costs O(seq_len^2) memory; for long sequences, compute indices on-the-fly in streaming or use causal patterns to generate triangular indices cheaply.- Storing dense bias per pair would be O(num_heads × seq_len^2) if expanded — avoid by indexing into the small table and broadcasting.- Choose max_rel << seq_len (e.g., 64–512) so table size is small; clip distances beyond max_rel to the same bucket ("clipped" or "bucketing").- Alternative bucketing: log-scale buckets for very long contexts (e.g., Transformer-XL / T5x use bucketing) to reduce number of buckets while preserving resolution for nearby tokens.Edge cases:- Causal masks: ensure bias respects masking (e.g., don't add future biases for causal attention).- Variable sequence lengths: recompute or slice index matrix; keep rpb independent of seq_len.Trade-offs:- Shaw-style gives richer interactions (adds embeddings into keys/values) but increases parameter count and compute; T5-style is simpler, memory-efficient, and works well in practice.

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.