Research Scientist
Mathematical Foundations
Research Scientist roles at top AI labs require graduate-level mathematics. You must derive estimators from scratch, compute gradients of matrix functions, prove convergence bounds, and work with information theory. These skills are tested in coding questions ('implement PCA from scratch'), derivation questions ('derive the ELBO'), and discussion questions ('what is the condition number?').
Core Theory
- 1Matrix calculus: key identities for ML derivations. ∂(Ax)/∂x = A^T. ∂(x^T Ax)/∂x = (A + A^T)x = 2Ax if A symmetric. ∂(log det A)/∂A = A^{-T}. ∂tr(BA)/∂A = B^T. For neural networks: if y = Wx, then ∂L/∂W = (∂L/∂y)^T x^T. Memorize these — they come up in deriving gradient updates for any model.
- 2Eigendecomposition and SVD: symmetric matrix A = QΛQ^T where Q is orthogonal (eigenvectors) and Λ diagonal (eigenvalues). PCA: eigenvectors of covariance matrix = principal components. SVD: A = UΣV^T for any matrix. Connection: left singular vectors of A are eigenvectors of AA^T; right singular vectors are eigenvectors of A^T A. SVD enables low-rank approximation: best rank-k approximation is A_k = Σᵢ₌₁ᵏ σᵢuᵢvᵢ^T (Eckart-Young theorem).
- 3Probability: joint, marginal, conditional distributions. Chain rule: P(A,B,C) = P(A)P(B|A)P(C|A,B). Bayes' theorem: P(θ|x) ∝ P(x|θ)P(θ). Law of total expectation: E[X] = E[E[X|Y]]. Law of total variance: Var(X) = E[Var(X|Y)] + Var(E[X|Y]). Moment generating functions: M_X(t) = E[exp(tX)], useful for proving CLT and sums of distributions.
- 4Information theory: Shannon entropy H(X) = -Σ P(x) log P(x) — average bits needed to encode X. Cross-entropy H(p,q) = -Σ p(x) log q(x) — expected bits when using code optimized for q to encode p. KL divergence: KL(p||q) = Σ p(x) log(p(x)/q(x)) = H(p,q) - H(p). KL ≥ 0, KL = 0 iff p = q. KL is not symmetric (KL(p||q) ≠ KL(q||p)).
- 5MLE and MAP: MLE finds θ that maximizes P(data|θ). MAP finds θ that maximizes P(θ|data) ∝ P(data|θ)P(θ) — equivalent to MLE with regularization. Gaussian prior on θ → L2 regularization (Ridge). Laplace prior → L1 regularization (Lasso). Bayesian inference: P(θ|data) = P(data|θ)P(θ)/P(data). The posterior is the full distribution, not a point estimate.
- 6Convex optimization: a function f is convex if f(λx + (1-λ)y) ≤ λf(x) + (1-λ)f(y). L-smooth: ||∇f(x) - ∇f(y)|| ≤ L||x-y||. μ-strongly convex: f(y) ≥ f(x) + ∇f(x)^T(y-x) + μ/2||y-x||². Gradient descent with step size 1/L converges at rate O(1/T) for L-smooth convex. Exponential rate O((1-μ/L)^T) for strongly convex. Condition number κ = L/μ determines convergence speed.
- 7ELBO (Evidence Lower Bound) for VAE: we want to maximize log P(x) = log ∫ P(x|z)P(z)dz. This is intractable. By Jensen's inequality: log P(x) = log E_q[P(x,z)/q(z)] ≥ E_q[log P(x|z)] - KL(q(z)||P(z)) = ELBO. Maximizing the ELBO wrt q(z) minimizes KL(q(z)||P(z|x)) — makes the variational posterior approximate the true posterior. Tight when q = P(z|x).
- 8Conjugate priors: if prior P(θ) and likelihood P(x|θ) have conjugate relationship, the posterior P(θ|x) has the same family as the prior. Examples: Beta-Binomial (beta prior, binomial likelihood → beta posterior), Normal-Normal (Gaussian prior, Gaussian likelihood → Gaussian posterior), Dirichlet-Multinomial. Enables closed-form Bayesian updates — crucial for bandit algorithms and Bayesian optimization.
Key Formulas
| 1 | import numpy as np |
| 2 | from scipy import stats |
| 3 | |
| 4 | np.random.seed(42) |
| 5 | |
| 6 | # ─── 1. PCA from Scratch via SVD ───────────────────────────────────────────── |
| 7 | n_samples, n_features = 500, 20 |
| 8 | k = 3 # number of principal components to keep |
| 9 | |
| 10 | # Generate correlated data |
| 11 | A = np.random.randn(n_features, n_features) |
| 12 | cov = A @ A.T / n_features + np.eye(n_features) * 0.1 |
| 13 | X = np.random.multivariate_normal(np.zeros(n_features), cov, n_samples) |
| 14 | |
| 15 | # PCA step 1: center the data |
| 16 | X_centered = X - X.mean(axis=0) |
| 17 | |
| 18 | # PCA step 2: SVD (NOT eigendecomposition — SVD is numerically more stable) |
| 19 | U, S, Vt = np.linalg.svd(X_centered, full_matrices=False) |
| 20 | |
| 21 | # Principal components = right singular vectors (rows of Vt) |
| 22 | # Eigenvalues of covariance matrix = S² / (n-1) |
| 23 | eigenvalues = S**2 / (n_samples - 1) |
| 24 | explained_var_ratio = eigenvalues / eigenvalues.sum() |
| 25 | |
| 26 | print("PCA via SVD:") |
| 27 | print(f"Top-{k} explained variance: {explained_var_ratio[:k].sum()*100:.1f}%") |
| 28 | |
| 29 | # Project to k-dimensional space |
| 30 | X_projected = X_centered @ Vt[:k].T # (500, 3) |
| 31 | print(f"Projected shape: {X_projected.shape}") |
| 32 | |
| 33 | # Verify: same as sklearn PCA |
| 34 | from sklearn.decomposition import PCA |
| 35 | pca_sklearn = PCA(n_components=k).fit(X) |
| 36 | print(f"sklearn PCA match: {np.allclose(np.abs(X_projected), np.abs(pca_sklearn.transform(X)), atol=1e-6)}") |
| 37 | |
| 38 | # ─── 2. KL Divergence between two Gaussians ─────────────────────────────────── |
| 39 | def kl_gaussian(mu1, sigma1, mu2, sigma2): |
| 40 | """ |
| 41 | KL(N(mu1, sigma1²) || N(mu2, sigma2²)) — closed form. |
| 42 | """ |
| 43 | return (np.log(sigma2/sigma1) |
| 44 | + (sigma1**2 + (mu1-mu2)**2) / (2*sigma2**2) |
| 45 | - 0.5) |
| 46 | |
| 47 | mu1, sigma1 = 0.0, 1.0 # standard normal |
| 48 | mu2, sigma2 = 1.0, 2.0 # shifted, wider normal |
| 49 | |
| 50 | kl_12 = kl_gaussian(mu1, sigma1, mu2, sigma2) |
| 51 | kl_21 = kl_gaussian(mu2, sigma2, mu1, sigma1) |
| 52 | |
| 53 | print(f"\nKL(N(0,1) || N(1,2)) = {kl_12:.4f}") |
| 54 | print(f"KL(N(1,2) || N(0,1)) = {kl_21:.4f}") |
| 55 | print(f"KL is asymmetric: {kl_12:.4f} ≠ {kl_21:.4f}") |
| 56 | |
| 57 | # Verify with Monte Carlo |
| 58 | samples = np.random.normal(mu1, sigma1, 100_000) |
| 59 | log_ratio = stats.norm.logpdf(samples, mu1, sigma1) - stats.norm.logpdf(samples, mu2, sigma2) |
| 60 | kl_mc = log_ratio.mean() |
| 61 | print(f"KL (Monte Carlo): {kl_mc:.4f} (analytical: {kl_12:.4f})") |
| 62 | |
| 63 | # ─── 3. ELBO for a simple Gaussian VAE (1D) ─────────────────────────────────── |
| 64 | def elbo(x: np.ndarray, mu_q: np.ndarray, log_var_q: np.ndarray) -> float: |
| 65 | """ |
| 66 | ELBO = E_q[log P(x|z)] - KL(q(z|x) || P(z)) |
| 67 | where: q(z|x) = N(mu_q, exp(log_var_q)) |
| 68 | P(z) = N(0, 1) |
| 69 | P(x|z) = N(z, 1) (simplified decoder) |
| 70 | """ |
| 71 | sigma_q = np.exp(0.5 * log_var_q) |
| 72 | z = mu_q + sigma_q * np.random.randn(*mu_q.shape) # reparameterization |
| 73 | |
| 74 | # Reconstruction term: E_q[log P(x|z)] |
| 75 | recon = -0.5 * ((x - z)**2 + np.log(2 * np.pi)) |
| 76 | |
| 77 | # KL term: KL(N(mu_q, sigma_q²) || N(0,1)) — closed form for Gaussians |
| 78 | kl = -0.5 * (1 + log_var_q - mu_q**2 - np.exp(log_var_q)) |
| 79 | |
| 80 | return float((recon - kl).mean()) |
| 81 | |
| 82 | x = np.array([1.5, -0.5, 2.0, 0.3, -1.2]) |
| 83 | mu_q = np.array([1.2, -0.3, 1.8, 0.2, -1.0]) # posterior mean |
| 84 | log_var_q = np.full(5, -1.0) # log variance ≈ -1 → σ ≈ 0.6 |
| 85 | |
| 86 | elbo_val = elbo(x, mu_q, log_var_q) |
| 87 | print(f"\nVAE ELBO: {elbo_val:.4f} (higher is better)") |
Three fundamental math computations: (1) PCA via SVD — center data, compute SVD, project using top-k right singular vectors. SVD is more numerically stable than eigendecomposition of the covariance matrix. (2) KL divergence between two Gaussians in closed form, showing KL is asymmetric. Monte Carlo verification. (3) VAE ELBO: reconstruction term E[log P(x|z)] minus KL(q(z|x)||P(z)) in closed form for Gaussian variational posterior.
Worked Interview Problems
3 problemsProblem
Observations X₁,...,Xₙ ~ N(μ, σ²) with known σ². Prior: μ ~ N(μ₀, τ²). Derive the MLE for μ and the MAP estimate. Show MAP = weighted average of MLE and prior mean.
Solution
MLE: maximize log P(data|μ) = -n/2 × log(2πσ²) - (1/2σ²)Σ(xᵢ-μ)². Take derivative: dL/dμ = (1/σ²)Σ(xᵢ-μ) = 0. Solve: μ_MLE = X̄.
MAP: maximize log P(μ|data) ∝ log P(data|μ) + log P(μ). Log prior: log P(μ) = -(μ-μ₀)²/(2τ²) + const.
MAP objective: -(1/2σ²)Σ(xᵢ-μ)² - (μ-μ₀)²/(2τ²). Take derivative: (1/σ²)Σ(xᵢ-μ) - (μ-μ₀)/τ² = 0.
Solve: n(X̄-μ)/σ² = (μ-μ₀)/τ². Rearranging: μ × (n/σ² + 1/τ²) = nX̄/σ² + μ₀/τ².
MAP estimate: μ_MAP = (n/σ²)/(n/σ² + 1/τ²) × X̄ + (1/τ²)/(n/σ² + 1/τ²) × μ₀. This is a weighted average of X̄ and μ₀, with weights proportional to data precision n/σ² and prior precision 1/τ².
Answer
. As n→∞, MAP→MLE. As n→0, MAP→μ₀ (prior mean). Gaussian prior on μ is equivalent to L2 regularization on the MLE objective.
Common Mistakes
Not knowing that KL divergence is asymmetric. KL(p||q) ≠ KL(q||p) in general. KL(p||q) (forward KL) is 'mode-seeking' — when minimized wrt q, the learned distribution covers the full support of p. KL(q||p) (reverse KL) is 'mode-averaging' — when minimized wrt q, the learned distribution concentrates on one mode of p.
Confusing covariance matrix eigendecomposition and SVD. For the covariance matrix (symmetric PSD), eigenvectors ARE the principal components. For the data matrix X, SVD gives U, S, V^T — the principal components are the right singular vectors V^T, and eigenvalues = S²/(n-1). Always center the data before PCA.
Not knowing the closed-form KL divergence between two Gaussians: KL(N(μ₁,σ₁²) || N(μ₂,σ₂²)) = log(σ₂/σ₁) + (σ₁² + (μ₁-μ₂)²)/(2σ₂²) - 1/2. This appears in VAE training, RLHF, and many other settings.
Confusing MLE (maximize likelihood) with MAP (maximize posterior). MAP = MLE + regularization. Gaussian prior → L2 regularization, Laplace prior → L1 regularization. When someone says 'add L2 regularization,' they are implicitly using a Bayesian MAP framework with a Gaussian prior.
Claiming the ELBO derivation uses strict equality. Jensen's inequality gives a lower bound: log E[X] ≥ E[log X] (log is concave). The ELBO is a lower bound on log P(x), not an equality. The gap is KL(q||P(z|x)) — this is what VAE training minimizes.