ML Engineer
Deep Learning Engineering
Deep learning engineering interviews test whether you can implement, debug, and optimize neural networks from first principles. You must be able to derive backpropagation, implement attention from scratch, debug training instabilities, and explain why architectural choices like LayerNorm over BatchNorm matter in transformers.
Core Theory
- 1Backpropagation: application of the chain rule through a computational graph. For each node, compute local gradient × upstream gradient. The key identity: dL/dW = dL/dy × dy/dW. For a linear layer y = Wx + b: dL/dW = δ^T x, dL/dx = W^T δ where δ = dL/dy. Vanishing gradients occur when these products shrink exponentially through many layers.
- 2Scaled dot-product attention: Attention(Q,K,V) = softmax(QK^T/√d_k)V. The 1/√d_k scaling prevents the dot products from growing large (which pushes softmax into regions with tiny gradients). Multi-head attention projects Q,K,V into h subspaces, computes attention in each, then concatenates: MHA(Q,K,V) = Concat(head_1,...,head_h)W_O.
- 3BatchNorm vs LayerNorm: BatchNorm normalizes over the batch dimension (across samples, per feature) — requires large batches to estimate mean/variance, has different train/eval behavior (running stats at eval). LayerNorm normalizes over the feature dimension (per sample, across features) — batch-size independent, same behavior at train/eval, preferred for transformers and variable-length sequences.
- 4Adam optimizer: m_t = β_1 m_{t-1} + (1-β_1)g_t (first moment, momentum). v_t = β_2 v_{t-1} + (1-β_2)g_t² (second moment, RMS). Bias correction: m̂ = m/(1-β_1^t), v̂ = v/(1-β_2^t). Update: θ -= lr × m̂/(√v̂ + ε). AdamW = Adam + decoupled weight decay (W -= lr×λW added separately, not via gradient). This prevents weight decay from being effectively reduced by the adaptive learning rate.
- 5Gradient clipping: clip gradients by global norm before the optimizer step. global_norm = √(Σ‖g_i‖²). If global_norm > max_norm: g_i *= max_norm/global_norm. Prevents gradient explosion in RNNs and transformers with deep stacks. Standard max_norm = 1.0 for language models.
- 6Mixed precision (fp16/bf16): compute forward + backward pass in float16/bfloat16, store master weights in float32. fp16 has range [6e-5, 65504] — can underflow/overflow. Loss scaling multiplies loss before backward, unscales before optimizer step to prevent underflow. bf16 has same range as float32 but less precision — preferred for training LLMs (fewer NaN issues).
- 7KV caching for autoregressive inference: at each generation step, instead of recomputing all K,V for previous tokens, cache them. Memory: O(n_layers × seq_len × d_model). Reduces inference from O(T²) to O(T). Essential for long-context generation. PagedAttention (vLLM) manages KV cache memory as virtual pages to improve GPU utilization.
- 8Common training bugs: NaN loss (LR too high, bad initialization, exploding gradients — check with gradient clipping + LR warmup), loss plateau (learning rate too low or batch size too large), loss spike at step N (usually data pipeline issue — a corrupted batch). Debug by: logging gradient norms, loss on a tiny overfit test, learning rate range test.
Key Formulas
| 1 | import torch |
| 2 | import torch.nn as nn |
| 3 | import torch.nn.functional as F |
| 4 | import math |
| 5 | |
| 6 | class MultiHeadAttention(nn.Module): |
| 7 | """ |
| 8 | Multi-head self-attention from scratch. |
| 9 | Same interface as nn.MultiheadAttention but fully transparent. |
| 10 | """ |
| 11 | def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): |
| 12 | super().__init__() |
| 13 | assert d_model % n_heads == 0, "d_model must be divisible by n_heads" |
| 14 | self.d_model = d_model |
| 15 | self.n_heads = n_heads |
| 16 | self.d_k = d_model // n_heads # head dimension |
| 17 | |
| 18 | # Single projection matrices (more efficient than separate per-head) |
| 19 | self.W_q = nn.Linear(d_model, d_model, bias=False) |
| 20 | self.W_k = nn.Linear(d_model, d_model, bias=False) |
| 21 | self.W_v = nn.Linear(d_model, d_model, bias=False) |
| 22 | self.W_o = nn.Linear(d_model, d_model, bias=False) |
| 23 | self.dropout = nn.Dropout(dropout) |
| 24 | |
| 25 | def forward(self, x: torch.Tensor, |
| 26 | mask: torch.Tensor = None) -> torch.Tensor: |
| 27 | B, T, C = x.shape # batch, sequence length, d_model |
| 28 | |
| 29 | # 1. Project to Q, K, V |
| 30 | Q = self.W_q(x) # (B, T, d_model) |
| 31 | K = self.W_k(x) |
| 32 | V = self.W_v(x) |
| 33 | |
| 34 | # 2. Reshape to (B, n_heads, T, d_k) |
| 35 | def split_heads(t): |
| 36 | return t.view(B, T, self.n_heads, self.d_k).transpose(1, 2) |
| 37 | |
| 38 | Q, K, V = split_heads(Q), split_heads(K), split_heads(V) |
| 39 | |
| 40 | # 3. Scaled dot-product attention |
| 41 | scale = math.sqrt(self.d_k) |
| 42 | scores = torch.matmul(Q, K.transpose(-2, -1)) / scale # (B, h, T, T) |
| 43 | |
| 44 | if mask is not None: |
| 45 | scores = scores.masked_fill(mask == 0, float('-inf')) |
| 46 | |
| 47 | attn_weights = F.softmax(scores, dim=-1) # (B, h, T, T) |
| 48 | attn_weights = self.dropout(attn_weights) |
| 49 | |
| 50 | # 4. Weighted sum of values |
| 51 | out = torch.matmul(attn_weights, V) # (B, h, T, d_k) |
| 52 | |
| 53 | # 5. Concatenate heads and project |
| 54 | out = out.transpose(1, 2).contiguous().view(B, T, C) # (B, T, d_model) |
| 55 | return self.W_o(out) |
| 56 | |
| 57 | # ─── Test ───────────────────────────────────────────────────────────────────── |
| 58 | B, T, d_model, n_heads = 2, 16, 256, 8 |
| 59 | x = torch.randn(B, T, d_model) |
| 60 | mha = MultiHeadAttention(d_model, n_heads) |
| 61 | out = mha(x) |
| 62 | print(f"Input: {x.shape}") |
| 63 | print(f"Output: {out.shape}") # should be (2, 16, 256) |
| 64 | |
| 65 | # Causal mask for autoregressive generation |
| 66 | causal_mask = torch.tril(torch.ones(T, T)).unsqueeze(0).unsqueeze(0) |
| 67 | out_causal = mha(x, mask=causal_mask) |
| 68 | print(f"Causal output: {out_causal.shape}") |
Full multi-head self-attention implementation from scratch. Key implementation choices: (1) single W_q/W_k/W_v matrices projected and then reshaped — more memory efficient than n_heads separate matrices; (2) the reshape+transpose dance to get (B, n_heads, T, d_k); (3) causal masking fills -inf before softmax so future tokens get attention weight 0; (4) .contiguous() before view because transpose makes memory layout non-contiguous.
Worked Interview Problems
3 problemsProblem
Given logits z and one-hot target y, show that the gradient of cross-entropy loss through the softmax is simply ŷ - y. This comes up in almost every DL theory interview.
Solution
Softmax: ŷ_i = exp(z_i) / Σ_j exp(z_j). Cross-entropy loss: L = -Σ_i y_i log(ŷ_i). For one-hot target, L = -log(ŷ_c) where c is the correct class.
∂L/∂ŷ_c = -1/ŷ_c. ∂ŷ_c/∂z_c = ŷ_c(1 - ŷ_c) (diagonal term of softmax Jacobian). ∂ŷ_j/∂z_c = -ŷ_j·ŷ_c for j≠c (off-diagonal term).
By chain rule: ∂L/∂z_c = ∂L/∂ŷ_c × ∂ŷ_c/∂z_c + Σ_{j≠c} ∂L/∂ŷ_j × ∂ŷ_j/∂z_c.
= (-1/ŷ_c) × ŷ_c(1-ŷ_c) + Σ_{j≠c} (-y_j/ŷ_j) × (-ŷ_j·ŷ_c) = -(1-ŷ_c) + Σ_{j≠c} y_j·ŷ_c.
For one-hot y: y_c = 1, y_j = 0 for j≠c. So ∂L/∂z_c = -(1-ŷ_c) + 0 = ŷ_c - 1 = ŷ_c - y_c. For other classes k≠c: ∂L/∂z_k = ŷ_k - 0 = ŷ_k - y_k. Therefore ∂L/∂z = ŷ - y.
Answer
. The gradient is simply the difference between predicted probabilities and true labels. This beautiful simplification is why softmax + cross-entropy is so widely used — the gradient computation is numerically stable and extremely simple.
Common Mistakes
Calling optimizer.zero_grad() after optimizer.step() instead of before backward(). The correct order is: zero_grad() → forward() → loss.backward() → (optional: clip_grad_norm_) → optimizer.step().
Forgetting model.eval() during validation — BatchNorm and Dropout behave differently at train vs eval. BatchNorm uses running statistics at eval, dropout is disabled. Always use with torch.no_grad() + model.eval() for evaluation.
Using the same model instance for both teacher and student in knowledge distillation without detaching the teacher's outputs. Teacher outputs must be .detach()ed to prevent gradients flowing through them.
Transposing K incorrectly in attention: should be K.transpose(-2, -1) to swap the last two dimensions (seq_len × d_k → d_k × seq_len), not K.T which would transpose all dimensions for batched tensors.
Not using model.train() after model.eval() — once you set eval mode for validation, you must explicitly return to train mode for the next training step.