9 ML Engineering From Scratch Interview Questions and Answers

Short answer: A two-step loop: assign every point to its nearest centroid, then recompute centroids as the mean of their clusters — until the centroids stop moving. Distances come from a single broadcasting expression, no loops over points.

In depth:

1. Distances in one line — ((X[:, None, :] - C[None, :, :])**2).sum(-1) yields an (n, k) matrix. At scale, use the ‖x‖² − 2xᵀc + ‖c‖² expansion to avoid materializing the (n, k, d) tensor.
2. Assign and update — argmin over the cluster axis, then the mean of each cluster's points.
3. Empty cluster is THE trap — if no point lands on a centroid, mean returns NaN; re-seed it from a random data point.
4. Stopping — on centroid movement, not just an iteration cap.

import numpy as np

def kmeans(X, k, iters=100, tol=1e-6, seed=0):
    rng = np.random.default_rng(seed)
    C = X[rng.choice(len(X), k, replace=False)]
    for _ in range(iters):
        d2 = ((X[:, None, :] - C[None, :, :])**2).sum(-1)  # (n, k)
        labels = d2.argmin(1)
        newC = C.copy()
        for j in range(k):  # looping over k only is fine
            pts = X[labels == j]
            # empty cluster -> re-seed from a random point
            newC[j] = pts.mean(0) if len(pts) else X[rng.integers(len(X))]
        if np.linalg.norm(newC - C) < tol:
            break
        C = newC
    return C, labels

⚠️ Common mistake: a double loop over points and centroids instead of broadcasting — and the silent NaN from an empty cluster: interviewers almost always probe exactly that.

Short answer: The model is p = σ(Xw), the loss is binary cross-entropy, and its gradient collapses to Xᵀ(p − y)/n. What remains: a stable sigmoid, probability clipping in the loss, and the descent loop; bias comes from a column of ones in X.

In depth:

1. Stable sigmoid — the naive 1/(1+exp(-z)) overflows for large |z|; branch on the sign of z.
2. Loss without log(0) — a saturated sigmoid pushes p to 0 or 1; clip with np.clip(p, eps, 1-eps) or compute via logaddexp.
3. Why the gradient is that clean — chain rule: dL/dp × dp/dz produces σ′, which cancels against the log's derivative, leaving (p − y). Hence Xᵀ(p − y)/n — a favorite follow-up.

import numpy as np

def sigmoid(z):
    out = np.empty_like(z)          # stable for large |z|
    pos = z >= 0
    out[pos] = 1 / (1 + np.exp(-z[pos]))
    e = np.exp(z[~pos])
    out[~pos] = e / (1 + e)
    return out

def fit(X, y, lr=0.1, epochs=1000, eps=1e-12):
    Xb = np.hstack([X, np.ones((len(X), 1))])  # bias — column of ones
    w = np.zeros(Xb.shape[1])
    for _ in range(epochs):
        p = sigmoid(Xb @ w)
        pc = np.clip(p, eps, 1 - eps)          # log(0) cannot happen
        loss = -np.mean(y*np.log(pc) + (1-y)*np.log(1-pc))
        w -= lr * Xb.T @ (p - y) / len(y)
    return w

⚠️ Common mistake: log(p) without clipping — a saturated sigmoid yields log(0) = −inf and NaN; and a forgotten bias — the decision boundary is forced through the origin.

Short answer: scores = QKᵀ/√d_k; fill the upper triangle (the future) with −inf BEFORE softmax; softmax over the last axis with the row max subtracted; output is weights @ V.

In depth:

1. Why √d_k — the variance of a dot product grows with dimensionality; without scaling, softmax saturates and gradients vanish. The standard follow-up: "why the square root?"
2. Mask strictly before softmax — −inf becomes exactly 0 weight after the exponent, while each row still sums to 1. Zeroing weights after softmax breaks the normalization.
3. Stable softmax — subtract the row max before exp.
4. Axis — softmax over keys (the last axis of scores), not over queries.

import numpy as np

def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)  # stability
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V, causal=True):
    d_k = Q.shape[-1]
    scores = Q @ K.swapaxes(-1, -2) / np.sqrt(d_k)
    if causal:
        T = scores.shape[-1]
        mask = np.triu(np.ones((T, T), dtype=bool), k=1)
        scores = np.where(mask, -np.inf, scores)  # BEFORE softmax!
    w = softmax(scores, axis=-1)  # over the key axis
    return w @ V

⚠️ Common mistake: softmax over the wrong axis — and masking after softmax: weights no longer sum to one, and the model quietly peeks into the future.

Short answer: In train mode, normalize with the current batch statistics and update running mean/var with momentum; in eval mode, use ONLY the running statistics. Then scale-shift: y = γ·x̂ + β.

In depth:

1. Train — μ, σ² over the batch, x̂ = (x − μ)/√(σ² + ε), running stats updated as an exponential moving average.
2. Eval — batch statistics are off-limits: predictions would depend on batch composition, and with batch=1 the variance is zero.
3. γ and β — learnable; they let the layer undo the normalization if needed.
4. Why LayerNorm differs — it normalizes over the features of a single example, not the batch, so it has no modes and lives in transformers and RNNs.

import numpy as np

class BatchNorm:
    def __init__(self, d, momentum=0.1, eps=1e-5):
        self.gamma, self.beta = np.ones(d), np.zeros(d)
        self.run_mean, self.run_var = np.zeros(d), np.ones(d)
        self.momentum, self.eps = momentum, eps

    def __call__(self, x, training):
        if training:
            mu, var = x.mean(0), x.var(0)
            m = self.momentum
            self.run_mean = (1 - m) * self.run_mean + m * mu
            self.run_var = (1 - m) * self.run_var + m * var
        else:
            mu, var = self.run_mean, self.run_var  # NOT batch stats!
        x_hat = (x - mu) / np.sqrt(var + self.eps)
        return self.gamma * x_hat + self.beta

⚠️ Common mistake: one mode for everything — batch statistics at inference. Metrics jump with batch composition, batch=1 breaks. This is exactly what model.eval() toggles in PyTorch.

Short answer: Subtract the row max before exp — exp(89) already overflows float32. For the loss, don't compute softmax and log separately; combine them as log-sum-exp: CE = logsumexp(z) − z[y].

In depth:

1. Shift by the max — softmax is shift-invariant: exp(z − m)/Σexp(z − m) is the same value, minus the overflow.
2. log-softmax in one piece — log(softmax(z)) = z − m − log Σ exp(z − m); we never take the log of a near-zero.
3. CE as indexing — minus log-softmax at the correct class position, averaged over the batch.

import numpy as np

def log_softmax(z):
    z = z - z.max(axis=1, keepdims=True)  # exp(89) overflows fp32
    return z - np.log(np.exp(z).sum(axis=1, keepdims=True))

def cross_entropy(z, y):
    # CE = logsumexp(z) - z[y]; no separate softmax -> log
    n = len(y)
    return -log_softmax(z)[np.arange(n), y].mean()

⚠️ Common mistake: computing softmax, then taking the log — under saturation the probability rounds to 0 and log yields −inf. This is exactly why PyTorch's F.cross_entropy takes logits, not probabilities.

Short answer: Train: model.train(), then loop zero_grad → forward → loss → backward → step. Eval: model.eval() plus torch.no_grad(), accumulating metrics weighted by batch size.

In depth:

1. zero_grad is mandatory — PyTorch sums gradients into .grad; forget it and you train on gradients from several batches at once.
2. eval() and no_grad() are different things — the first switches dropout and BatchNorm to inference mode, the second disables graph construction. You need both.
3. Metric weighting — the last batch is usually shorter; mean-over-batches ≠ mean-over-dataset, so multiply by len(x).
4. .item() — otherwise the tensor drags its computation graph along and memory leaks.

def train_epoch(model, loader, opt, crit, device):
    model.train()
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        opt.zero_grad()            # otherwise gradients accumulate
        loss = crit(model(x), y)
        loss.backward()
        opt.step()

@torch.no_grad()
def evaluate(model, loader, crit, device):
    model.eval()                   # dropout off, BN -> running stats
    total, n = 0.0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        loss = crit(model(x), y)
        total += loss.item() * len(x)  # last batch is shorter
        n += len(x)
    return total / n

⚠️ Common mistake: a forgotten zero_grad — gradients accumulate and the loss "jumps strangely"; validating without model.eval() — dropout keeps dropping neurons; loss without .item() — graphs pile up in memory.

Short answer: Divide the logits by the temperature FIRST, then truncate: top-k keeps the k largest logits; top-p keeps the smallest sorted prefix whose probability sum reaches p (at least one token). After truncation — renormalize and sample.

In depth:

1. Order of operations — temperature reshapes the whole distribution, so it goes before truncation; T after top-k reweights an already-truncated set.
2. top-k — threshold on the k-th logit (np.sort(z)[-k] or argpartition), everything else to −inf.
3. top-p — sort descending, cumsum, cut where the sum reaches p; searchsorted + 1 includes the crossing token.
4. Renormalize — after zeroing the tail, divide probabilities by the new sum.

import numpy as np

def sample(logits, T=1.0, k=None, p=None, rng=np.random.default_rng()):
    z = logits / T                        # temperature BEFORE truncation
    if k is not None:
        kth = np.sort(z)[-k]              # k-th largest logit
        z = np.where(z < kth, -np.inf, z)
    probs = np.exp(z - z.max())
    probs /= probs.sum()
    if p is not None:
        order = np.argsort(-probs)
        cum = np.cumsum(probs[order])
        cut = np.searchsorted(cum, p) + 1  # keep at least one token
        probs[order[cut:]] = 0.0           # zero the tail
        probs /= probs.sum()               # renormalize
    return rng.choice(len(probs), p=probs)

⚠️ Common mistake: temperature after truncation — and cutting top-p strictly below p without the crossing token: on a peaked distribution you can end up with no candidates at all.

Short answer: Normalize the base rows once — cosine becomes a matmul of unit vectors; take top-k with np.argpartition in O(n) instead of a full sort. For 1M×d that's a single BLAS operation.

In depth:

1. Normalize up front — cos(a, b) = â·b̂; normalize the base once offline, so each query costs only a matrix multiply.
2. argpartition — selects the k largest in O(n) without sorting the whole array; you only sort the k survivors.
3. Zero-norm vectors — dividing by a zero norm yields NaN; add ε to the norm.
4. Bridge to the production answer — brute force at 1M vectors is still alive (milliseconds on BLAS); beyond that you reach for ANN: HNSW or IVF (faiss), trading recall for speed.

import numpy as np

def build_index(X, eps=1e-12):
    # once, offline: normalize the rows
    return X / (np.linalg.norm(X, axis=1, keepdims=True) + eps)

def top_k(Xn, q, k=10, eps=1e-12):
    qn = q / (np.linalg.norm(q) + eps)
    sims = Xn @ qn                        # cosine = matmul of unit vectors
    idx = np.argpartition(-sims, k)[:k]   # O(n), not a full sort
    return idx[np.argsort(-sims[idx])]    # sort only the k survivors

⚠️ Common mistake: np.sort on the entire similarity array for a top-10 — O(n log n) instead of O(n); and "naive" cosine that renormalizes the whole base on every query.

Short answer: Divide the loss by the number of accumulation steps, call backward on every micro-batch (gradients sum into .grad), and run optimizer.step() + zero_grad() once per accum_steps. This emulates a batch K times larger with no extra memory.

In depth:

1. Dividing the loss is mandatory — backward sums gradients; without the division the gradient comes out K times larger, i.e. the effective learning rate silently multiplies by K.
2. step/zero_grad on schedule — only every accum_steps micro-batches; zero_grad right after step.
3. Tail of the epoch — if the batch count isn't divisible by accum_steps, the last accumulated gradients must still be stepped, or they vanish — or leak into the next epoch.
4. BatchNorm caveat — BN statistics are still computed per physical micro-batch; accumulation does not "merge" them. For small batches, use GroupNorm/LayerNorm.

accum = 4
opt.zero_grad()
for i, (x, y) in enumerate(loader):
    loss = crit(model(x), y) / accum  # else gradients x K -> LR x K
    loss.backward()                   # gradients accumulate in .grad
    if (i + 1) % accum == 0:
        opt.step()
        opt.zero_grad()
# tail: batch count not divisible by accum — flush the remainder
if (i + 1) % accum != 0:
    opt.step()
    opt.zero_grad()

⚠️ Common mistake: not dividing the loss by accum_steps — training "suddenly" diverges because the effective LR grew K-fold. And don't rescale the LR "just in case": with the division done right, it needs no scaling.