Katkat
Back

Explain Transformer Attention Fundamentals

AdobeMachine Learning Engineer · Onsite
HardMachine Learning EngineerOnsiteMachine Learning

Last updated

This question tests the candidate's deep conceptual knowledge of the Transformer architecture, focusing on the mathematical mechanics of scaled dot-product attention, memory optimization, and hardware-aware training techniques.

Transformer Architecture Fundamentals

Explain the mechanics of the Transformer architecture by addressing the following components and their implications for model performance and scalability.

1. The Attention Mechanism

Define scaled dot-product attention mathematically: Attention(Q, K, V) = softmax((QK^T) / sqrt(d_k))V. Explain why the scaling factor 1 / sqrt(d_k) is critical for the stability of the softmax gradient. What happens to the attention distribution if this scaling is omitted as the dimensionality d_k increases?

2. KV-Caching

During inference, we compute the Key and Value states for each token. Describe the memory footprint of the Key-Value (KV) cache for a sequence of length N in a model with L layers, H attention heads, and hidden dimension d_model. Explain how KV-caching changes the computational complexity of the forward pass from O(N^2) to O(N) per generated token.

3. Training Stability and Precision

Explain the rationale behind using mixed-precision training (FP16/BF16) in modern LLM pre-training. Specifically, discuss the role of the 'loss scaler' in FP16 training and why BF16 has largely superseded it in recent hardware architectures.

4. Architectural Trade-offs

Contrast 'Masked Self-Attention' in a decoder-only architecture with 'Cross-Attention' in an encoder-decoder architecture. In what scenarios would you choose one over the other?

Hints

Hint 1

Consider the variance of the dot product QK^T when the values in Q and K are independent random variables with mean 0 and variance 1.

Hint 2

Think about the memory complexity of storing KV pairs per token versus recomputing them, and how it impacts the latency of the 'prefill' versus 'decoding' phases.

Hint 3

Compare the dynamic range of the exponent representable in IEEE 754 half-precision versus the Bfloat16 format.

Follow-up questions

  • How does FlashAttention optimize the memory hierarchy (SRAM vs HBM) to improve attention performance?
  • What is the impact of positional encoding choices (e.g., RoPE vs Absolute embeddings) on the model's ability to handle inference sequences longer than those seen during training?
  • If you observe 'activation spikes' during training, which architectural modifications (e.g., LayerNorm placement or QKV initialization) would you investigate?