Implementing machine learning from scratch in an interview: the details that break

Machine learning engineer and research loops at technology companies, AI labs and quant funds now often include a round that is neither algorithms nor theory. You implement a piece of machine learning yourself, usually in NumPy, while the interviewer watches. Nobody expects a new idea. They expect the standard thing, written correctly, with the shapes right and the numerics safe.

Most failures in this round are not about understanding the model. They come from a single detail: an overflow, a transposed gradient, a mask applied in the wrong place. Here are eight implementations that come up again and again, and the detail in each one. Each links to a problem in the machine learning coding course, where it runs in a Python and NumPy editor in the page against the tests an interviewer would use.

1. Softmax that survives large logits

np.exp(z) / np.exp(z).sum() overflows as soon as a logit passes about 710, and the result is nan. Subtract the maximum first: softmax does not change when every logit moves by the same amount, and now the largest exponent is 1. The numerical stability lesson steps through it on logits near 1000; practise it in Softmax that never overflows.

2. Cross-entropy from logits, not from probabilities

Taking the log of a softmax that underflowed to zero gives minus infinity. Work in log space throughout with log-sum-exp, and the gradient with respect to the logits comes out as the famously simple softmax − one-hot, divided by the batch size if the loss is a mean. Interviewers ask for that gradient. See Cross-entropy loss and its gradient from logits.

3. Backpropagation through a two-layer network

The chain rule is not the hard part; the shapes are. Keep the batch in the first dimension, cache what the backward pass needs during the forward pass, and check every gradient has the shape of the thing it differentiates: dW1 the shape of W1. Then check the result numerically with finite differences, which interviewers love to see. The backpropagation lesson walks a forward and backward pass one step at a time; the problem is backpropagation through a two-layer network.

4. Attention with a causal mask

Scores are Q Kᵀ / √d; the mask must be applied before the softmax, by setting future positions to minus infinity (or a large negative number), not by zeroing weights afterwards, which leaves rows that no longer sum to one. Forgetting the √d pushes softmax into saturation as the dimension grows. See scaled dot-product attention.

5. Batch norm's two variances

During training, batch norm normalises with the batch's own mean and variance and updates running averages; at inference it uses the running averages. Mixing the two, or using the unbiased variance where the biased one belongs, is the usual bug. Layer norm normalises across features instead and needs neither. See batch normalisation and layer normalisation.

6. k-means and the empty cluster

Lloyd's algorithm alternates two steps: assign each point to its nearest centroid, then move each centroid to the mean of its points. If a centroid ends up with no points, its mean is undefined, and a careless implementation returns nan. Decide a rule (keep it, or reseed it) and say it out loud. Compute distances without a Python loop, using ‖x‖² + ‖c‖² − 2x·c. See k-means clustering.

7. PCA and the sign of an eigenvector

Centre the data first, then take the eigenvectors of the covariance matrix or, better, the singular vectors of the centred data. Each component is defined only up to sign, so two correct implementations can disagree; a deterministic sign convention makes results comparable. See PCA.

8. ROC AUC with tied scores

AUC is the probability that a random positive scores above a random negative, with ties counted as a half. Sorting and counting is O(n log n); the trap is ties, which a naive rank-based formula mishandles unless it uses average ranks. See ROC AUC.

How to practise

Write each one from a blank file, without looking, against tests that include the awkward cases: a logit of 1000, a tied score, an empty cluster. Then explain it out loud as you would to an interviewer, including the detail above. Twenty short implementations done that way prepare you better than reading a hundred.

The machine learning coding course has 36 such problems in 11 lessons, each with an editor that runs NumPy in your browser, the tests it is judged on, and a worked solution checked against an independent implementation. Three lessons are free to read in full, including numerical stability and backpropagation.