Research Scientist
Competitive Coding (Hard Level)
Research Scientist roles at top AI labs (OpenAI, DeepMind, Google Brain, Meta AI) still require coding rounds โ often harder than MLE roles. Expect LeetCode Hard, algorithm design problems, and 'implement this ML component from scratch in 30 minutes' questions. Building transformers, attention, and backprop from memory is tested.
Core Theory
- 1Advanced DP patterns: bitmask DP โ state is a subset of elements, encode as an integer. Used for: TSP, assignment problems, shortest Hamiltonian path. Interval DP โ state is a range [i,j], transitions merge intervals. Used for: matrix chain multiplication, burst balloons, palindrome problems. DP on trees โ combine results from subtrees. Used for: maximum independent set on trees, tree rerooting technique.
- 2Segment trees: support range queries and point/range updates in O(log n). Build: recursively divide [0,n-1] into halves. Update: update leaf, propagate up. Query: combine results from covering segments. Lazy propagation: defer updates to avoid O(n log n) per operation. Used for: range minimum query (RMQ), range sum, range maximum with lazy updates.
- 3Fenwick trees (Binary Indexed Trees): O(log n) prefix sum query and point update. Simpler than segment tree for prefix sum queries. update(i, delta): add delta to all nodes responsible for position i. query(i): sum of all elements [1, i]. Built on the observation that i & (-i) gives the lowest set bit of i.
- 4Advanced graph algorithms: Tarjan's SCC โ DFS-based algorithm for strongly connected components in O(V+E). Critical for dependency resolution and circuit analysis. Dinic's max-flow โ BFS to build level graph, DFS to find blocking flow, repeat. O(VยฒE) general, O(EโV) for unit capacity graphs. Max-flow min-cut theorem.
- 5String algorithms: KMP (Knuth-Morris-Pratt) โ build failure function in O(m), then match in O(n). Z-algorithm โ for each position, compute longest substring starting there that is also a prefix. Suffix arrays + LCP: sort all suffixes, build LCP array. Applications: longest repeated substring, number of distinct substrings.
- 6Implement ML components from scratch: multi-head attention in PyTorch (without nn.MultiheadAttention), backpropagation through a custom layer in NumPy, LSTM cell forward pass, batch normalization backward pass. These are tested at every AI research lab interview.
- 7Randomized algorithms: reservoir sampling โ uniformly sample k elements from a stream of unknown length without storing all elements. Miller-Rabin primality test โ probabilistic primality testing in O(k logยฒ n). Treap โ BST + heap, randomized BST with expected O(log n) all operations.
- 8Competitive programming time complexity: for n=10^6: O(n) or O(n log n) required. For n=10^5: O(n logยฒn) acceptable. For n=10^4: O(nยฒ) acceptable. For n=10^3: O(nยณ) acceptable. Always state the time complexity of your solution and verify it fits within these bounds.
Key Formulas
| 1 | import numpy as np |
| 2 | |
| 3 | def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray: |
| 4 | """Numerically stable softmax.""" |
| 5 | e = np.exp(x - x.max(axis=axis, keepdims=True)) |
| 6 | return e / e.sum(axis=axis, keepdims=True) |
| 7 | |
| 8 | def scaled_dot_product_attention(Q: np.ndarray, K: np.ndarray, |
| 9 | V: np.ndarray, |
| 10 | mask: np.ndarray = None) -> tuple: |
| 11 | """ |
| 12 | Compute scaled dot-product attention. |
| 13 | Q: (batch, heads, seq, d_k) |
| 14 | K: (batch, heads, seq, d_k) |
| 15 | V: (batch, heads, seq, d_v) |
| 16 | Returns: output (batch, heads, seq, d_v), weights (batch, heads, seq, seq) |
| 17 | """ |
| 18 | d_k = Q.shape[-1] |
| 19 | scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(d_k) # (B, H, seq, seq) |
| 20 | |
| 21 | if mask is not None: |
| 22 | scores = np.where(mask == 0, -1e9, scores) |
| 23 | |
| 24 | weights = softmax(scores, axis=-1) |
| 25 | output = weights @ V |
| 26 | return output, weights |
| 27 | |
| 28 | def multi_head_attention(x: np.ndarray, |
| 29 | W_q: np.ndarray, W_k: np.ndarray, |
| 30 | W_v: np.ndarray, W_o: np.ndarray, |
| 31 | n_heads: int, mask: np.ndarray = None) -> np.ndarray: |
| 32 | """ |
| 33 | Full multi-head attention. |
| 34 | x: (batch, seq, d_model) |
| 35 | W_q, W_k, W_v: (d_model, d_model) |
| 36 | W_o: (d_model, d_model) |
| 37 | """ |
| 38 | B, T, d_model = x.shape |
| 39 | d_k = d_model // n_heads |
| 40 | |
| 41 | # Project to Q, K, V |
| 42 | Q = x @ W_q # (B, T, d_model) |
| 43 | K = x @ W_k |
| 44 | V = x @ W_v |
| 45 | |
| 46 | # Reshape to (B, n_heads, T, d_k) |
| 47 | def split_heads(t): |
| 48 | return t.reshape(B, T, n_heads, d_k).transpose(0, 2, 1, 3) |
| 49 | |
| 50 | Q, K, V = split_heads(Q), split_heads(K), split_heads(V) |
| 51 | |
| 52 | # Attention |
| 53 | out, attn_weights = scaled_dot_product_attention(Q, K, V, mask) |
| 54 | |
| 55 | # Merge heads: (B, n_heads, T, d_k) โ (B, T, d_model) |
| 56 | out = out.transpose(0, 2, 1, 3).reshape(B, T, d_model) |
| 57 | return out @ W_o, attn_weights |
| 58 | |
| 59 | # โโโ Test โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ |
| 60 | np.random.seed(42) |
| 61 | B, T, d_model, n_heads = 2, 8, 64, 4 |
| 62 | d_k = d_model // n_heads # 16 |
| 63 | |
| 64 | x = np.random.randn(B, T, d_model) |
| 65 | W_q = np.random.randn(d_model, d_model) / np.sqrt(d_model) |
| 66 | W_k = np.random.randn(d_model, d_model) / np.sqrt(d_model) |
| 67 | W_v = np.random.randn(d_model, d_model) / np.sqrt(d_model) |
| 68 | W_o = np.random.randn(d_model, d_model) / np.sqrt(d_model) |
| 69 | |
| 70 | # Causal (autoregressive) mask: position i can only attend to positions 0..i |
| 71 | causal_mask = np.tril(np.ones((T, T)))[np.newaxis, np.newaxis, :, :] |
| 72 | |
| 73 | output, weights = multi_head_attention(x, W_q, W_k, W_v, W_o, n_heads, causal_mask) |
| 74 | print(f"Input: {x.shape}") |
| 75 | print(f"Output: {output.shape}") |
| 76 | print(f"Weights: {weights.shape} (B, heads, seq, seq)") |
| 77 | print(f"\nCausal mask check โ row 0 should attend only to position 0:") |
| 78 | print(f" weights[0, 0, 0, :] = {weights[0, 0, 0, :].round(3)}") |
| 79 | print(f" weights[0, 0, 1, :] = {weights[0, 0, 1, :].round(3)}") |
Multi-head attention implemented purely in NumPy (no PyTorch). Key implementation decisions: (1) numerically stable softmax via subtracting max; (2) Q,K,V projected with a single matmul then reshaped โ more efficient than per-head projections; (3) transpose pattern (0,2,1,3) to get (B, heads, T, d_k); (4) causal masking using -1e9 (not -inf to avoid NaN after softmax); (5) merge heads by transposing back and reshaping.
Worked Interview Problems
3 problemsProblem
You're asked to implement the backward pass for the scaled dot-product attention in NumPy. Given upstream gradient dL/dOut, compute dL/dQ, dL/dK, dL/dV.
Solution
Forward pass (recap): scores = QK^T/โd_k. weights = softmax(scores). out = weights @ V.
Backward through out = weights @ V: dL/dweights = dL/dOut @ V^T. dL/dV = weights^T @ dL/dOut.
Backward through softmax: for softmax output s with upstream gradient g, gradient of each element: ds_i/dscores_j = s_i(ฮด_{ij} - s_j). dL/dscores = s โ (g - (g โ s).sum(keepdims=True)). [This is the softmax Jacobian applied efficiently without materializing the full nรn Jacobian.]
Backward through QK^T/โd_k: dL/dQ = (dL/dscores @ K) / โd_k. dL/dK = (dL/dscores^T @ Q) / โd_k.
Stack these and the chain rule through the projection matrices W_q, W_k, W_v, W_o to complete the full backward pass.
Answer
dL/dV = weights^T @ dL/dOut. dL/dweights = dL/dOut @ V^T. dL/dscores = softmax_backward(weights, dL/dweights). dL/dQ = dL/dscores @ K / โd_k. dL/dK = dL/dscores^T @ Q / โd_k. The softmax backward is the most subtle step โ it uses s โ (g - sum(gโs)) rather than the full Jacobian.
Common Mistakes
Not knowing the reshape + transpose pattern for multi-head attention from memory. The sequence is: project โ reshape(B,T,H,d_k) โ transpose(0,2,1,3) to get (B,H,T,d_k). And the merge: transpose(0,2,1,3) โ reshape(B,T,d_model). Practice this until it comes automatically.
Integer overflow in competitive programming: in Python this is not an issue, but when computing products of large numbers (2^31 ร 2^31 in bitmask DP), be aware that other languages (C++, Java) require long long. Python integers are arbitrary precision.
Forgetting to handle the edge cases in segment trees: empty ranges (l > r โ return identity element), single elements (l == r โ base case), and the recursive call splitting vs. complete overlap cases.
Using recursion for complex DP without memoization. LeetCode Hard DP problems with O(nยณ) states will TLE if you use pure recursion. Use either top-down with @functools.lru_cache or bottom-up tabulation.
Not knowing the softmax backward formula: the naive approach materializes the nรn Jacobian (O(nยฒ) per example). The efficient formula is: dL/dscores = softmax_output โ (upstream_grad - sum(upstream_grad โ softmax_output)). This is the formula used in every production implementation.