ML Engineer Interview Questions

My first move is to add gradient norm logging immediately before the NaN appears, because NaN in the loss usually traces back to either an exploding gradient or a numerical issue in the forward pass, and gradient norms tell me which direction to look. If the gradient norm spikes to infinity in the backward pass, I know the issue is exploding gradients and I add gradient clipping as a short-term fix while investigating the root cause. If the loss is NaN from the first iteration, I look at the forward pass: check for log of zero in cross-entropy, division by zero in normalization layers, or large activations from poor initialization. I'll reduce batch size to 1 and inspect activations layer by layer to find where the NaN first appears. Common culprits I've seen in production: mixed-precision training with FP16 where activations overflow, a log function applied to a softmax that's exactly zero, or a learning rate that's two orders of magnitude too high. The systematic approach is to binary search: if batch 100 is NaN, check batch 50, then narrow down. I also check that the input data itself doesn't contain NaNs, which sounds obvious but has burned me before.

Self-attention computes a weighted combination of values, where the weights reflect how much each position should attend to every other position. We start by projecting each token into three vectors: a query (what this token is looking for), a key (what this token offers to others), and a value (the content it contributes if selected). We compute attention scores as the dot product of each query with all keys, then softmax-normalize them to get weights, and finally take a weighted sum of the values. The scaling by sqrt(d_k) is critical: without it, dot products in high dimensions grow large, pushing softmax into regions where gradients vanish, which kills learning. Scaling keeps the variance of the dot product roughly at 1 regardless of dimension.

My default is REST for most deployments because it is universally supported, easy to debug with curl, and works across every client stack. I switch to gRPC when I need low-latency, high-throughput inference, specifically because it uses HTTP/2 multiplexing and protobuf serialization which can cut payload size by 50-80% compared to JSON. The tradeoff is operational overhead: gRPC requires proto schema management and not all load balancers handle it well. In practice I have used gRPC internally between microservices, like from an API gateway to a TensorFlow Serving backend, while keeping a REST facade for external clients.

In data parallelism, every GPU holds a full model copy and processes a different shard of the batch, then gradients are averaged across workers. It scales well when the model fits on one GPU. Model parallelism splits layers across GPUs sequentially, which helps when a single layer sequence is too large but introduces pipeline bubbles. Tensor parallelism, as used in Megatron-LM, splits individual weight matrices across GPUs so that each device handles a slice of every matrix multiply, reducing activation memory and allowing wider models. Pipeline parallelism splits the model into stages and streams micro-batches through them to overlap compute, but you have to tune the number of micro-batches carefully to minimize the idle bubble. In practice I combine all four: data parallel across nodes, tensor parallel within a node for intra-GPU bandwidth, and pipeline parallel across nodes when the model is deep enough to fill the pipeline.

I would start by initializing weight matrices W1 and W2 and bias vectors b1 and b2. The forward pass is two sequential steps: first I compute the hidden layer pre-activation as the matrix product of the input X and W1 transposed, add b1, then apply ReLU element-wise by clamping negatives to zero. Then I compute the output as the matrix product of the hidden activations and W2 transposed, plus b2. The key thing I watch is shape alignment: if X is shape (batch, input_dim) and W1 is (hidden_dim, input_dim), then X at-sign W1.T gives (batch, hidden_dim), which is what we want. I also cache the intermediate values like the pre-activation and the ReLU mask, because backprop will need them.

I would break this into three stages: candidate generation, ranking, and re-ranking. Candidate generation uses approximate nearest neighbor search over user and item embeddings trained with two-tower models, narrowing billions of items down to a few thousand candidates in milliseconds. The ranking stage runs a deeper model, often a neural network with cross-features like watch history, session context, and item freshness, to score those candidates. Re-ranking applies business rules like diversity, novelty, and policy filters before serving. I store user embeddings in a low-latency key-value store like Redis and update them asynchronously using streaming pipelines so they reflect recent behavior without blocking the serving path. For training, I use logged feedback with careful treatment of position bias and use a mixture of implicit signals like watch time and explicit signals like ratings.