AI research engineer domain knowledge interview questions

As asked

Vanishing and exploding gradients in deep transformers: causes and solutions specific to large-scale LLM training. Explain what mathematical conditions lead to each, describe how residual connections and pre-norm placement specifically address them in a 96-layer transformer, and explain why gradient clipping is still necessary even with these architectural fixes.

Sample answer outline

Vanishing gradients occur when the gradient magnitude is repeatedly multiplied by values less than 1 during backpropagation through many layers, causing early-layer updates to effectively zero out. Exploding gradients occur when those multiplications exceed 1, causing instability. Solutions include residual connections (adding the input to the output bypasses layers whose gradients would vanish), normalization layers (LayerNorm stabilizes activation scale), careful initialization (Xavier or Kaiming scales weights to maintain variance through layers), and gradient clipping for the exploding case.

Expect these follow-ups

• Why did LSTMs address vanishing gradients in RNNs better than plain RNNs, and is that still relevant given transformers?
• What is the role of the residual stream in a transformer specifically in preventing gradient vanishing?

As asked

Define perplexity mathematically and explain its relationship to the cross-entropy loss your language model minimizes. Why is perplexity reported instead of raw loss in NLP papers?

Sample answer outline

Perplexity is exp(H) where H is the per-token cross-entropy loss in nats or bits. If loss is reported in nats, perplexity = exp(loss); if in bits, perplexity = 2^loss. Perplexity has a human-interpretable meaning: a perplexity of k means the model is as uncertain as if it had to choose uniformly among k alternatives at each token. Reporting perplexity instead of raw loss makes results comparable across tokenizer vocabulary sizes and easier to sanity-check (perplexity of 10 means roughly 10-way average uncertainty per token).

Expect these follow-ups

• Why does perplexity depend on the tokenizer, making it non-comparable across models with different vocabularies?
• What is a typical perplexity on WikiText-103 for a state-of-the-art small language model and what would a perplexity of 1.0 mean?

As asked

Walk me through the original encoder-decoder transformer from Attention Is All You Need. What is the role of the encoder, the decoder, and the cross-attention layer, and which components are dropped in decoder-only LLMs?

Sample answer outline

The encoder processes the input sequence with bidirectional self-attention and produces a context representation. The decoder generates the output sequence autoregressively using causal self-attention plus cross-attention that attends to the encoder output to incorporate source information. Decoder-only LLMs drop the encoder and cross-attention entirely, using only causal self-attention, since they are trained to predict the next token in a single unified sequence. The masking in the decoder's self-attention prevents each position from attending to future positions.

Expect these follow-ups

• For a translation task, what would you lose by using a decoder-only model instead of the original encoder-decoder?
• Why do encoder-only models like BERT use masked language modeling instead of causal language modeling?

As asked

Explain what temperature scaling does to the logit distribution before sampling. What does temperature = 1.0, temperature = 0.0 (approaching zero), and temperature = 2.0 produce, and when would you use each?

Sample answer outline

Temperature T divides logits before softmax: p_i = softmax(logits / T). At T=1, the distribution is unchanged from training. As T approaches zero, the distribution becomes a one-hot on the argmax token (greedy decoding). At T=2, the distribution becomes flatter and more uniform, increasing diversity. Low temperature is used for factual or code generation tasks where accuracy matters. High temperature is used for creative or diverse generation. Very high temperature produces incoherent text because low-probability tokens are assigned similar weight to high-probability ones.

Expect these follow-ups

• If you apply top-p filtering after temperature scaling, does the order matter and why?
• Why does temperature not help with the repetition problem in beam search?

As asked

Explain gradient accumulation. You want to train with an effective batch size of 1024 but can only fit a micro-batch of 64 on your GPU. How do you implement this, and what is the one common mistake that makes the accumulated gradients incorrect?

Sample answer outline

Run 16 forward-backward passes with micro-batch size 64, accumulating gradients in each parameter's .grad buffer without calling optimizer.step() between them. After 16 steps call optimizer.step() and zero_grad(). The common mistake is not dividing the loss by the accumulation steps before calling backward: if you compute loss over 64 tokens and accumulate 16 times, the effective loss magnitude is 16x larger than a single 1024-token batch, distorting gradient magnitudes. Dividing loss by accumulation_steps each step corrects this.

Expect these follow-ups

• Does gradient accumulation give mathematically identical results to a single large batch, and if not, why not?
• How does gradient accumulation interact with BatchNorm, and why is this not a problem in transformer training?

As asked

Walk me through backpropagation through a single transformer attention block, applying the chain rule at each step. Specifically: what intermediate activations must be stored during the forward pass to compute gradients with respect to Q, K, and V weight matrices, and how does the computational cost of one backward pass compare to one forward pass in a standard transformer?

Sample answer outline

The chain rule states dL/dw = dL/dy * dy/dw for composed functions. In a two-layer network L = loss(f2(f1(x, w1), w2)), backprop computes dL/dw2 = dL/dy2 * df2/dw2 at the output layer, then dL/dw1 = dL/dy2 * df2/dy1 * df1/dw1. To compute these products during the backward pass, the intermediate activations (y1, y2) from the forward pass must be stored because they appear in the Jacobians. This is why memory grows with depth and why gradient checkpointing reduces memory by recomputing rather than storing them.

Expect these follow-ups

• What does it mean for a function to be not differentiable and how does PyTorch handle that, for example with ReLU at zero?
• What is the computational cost ratio of one backward pass to one forward pass in a standard transformer?

As asked

A language model's output head uses a linear projection followed by softmax, not sigmoid. Explain why softmax is used instead of sigmoid for the next-token prediction head, what the paired loss function is and why it is numerically computed via log-softmax rather than softmax then log, and when a research engineer working on reward models or classifiers built on top of an LLM backbone would switch to sigmoid with binary cross-entropy.

Sample answer outline

Softmax with categorical cross-entropy is used for multi-class classification where exactly one class is correct, because softmax forces a valid probability distribution summing to one. Sigmoid with binary cross-entropy is used for multi-label classification where each label is independent, because each output is an independent probability and there is no constraint that they sum to one. Using softmax for multi-label forces an artificial competition between labels, making high probability for one label suppress others even when they are independently true.

Expect these follow-ups

• How does the numerical computation of cross-entropy with softmax differ from sigmoid with BCE in terms of stability?
• In a language model head that predicts the next token, which activation is used and why?

As asked

When you batch sequences of different lengths and pad them to the same length, you need an attention mask. Explain exactly what the mask does in the attention computation, and what goes wrong if you forget to mask padding tokens.

Sample answer outline

The attention mask sets the logits for padding positions to large negative values (typically -1e9 or -inf) before the softmax, so they contribute effectively zero weight in the attention output. If padding tokens are not masked, they receive non-zero attention weight and their embeddings (which are arbitrary random or zero vectors) pollute the attention output, degrading performance especially for tasks where the model processes the entire sequence. For the loss, padding tokens should also be excluded by setting their labels to -100 (the ignore index in PyTorch cross-entropy).

Expect these follow-ups

• In a causal language model, do you need both a causal mask and a padding mask, and how are they combined?
• What is the difference between applying the attention mask at the logit level versus at the key level?

As asked

Define precision and recall, explain the tradeoff between them, and give a concrete example from LLM evaluation or alignment where optimizing one hurts the other.

Sample answer outline

Precision is true positives / (true positives + false positives): among things predicted positive, how many actually are. Recall is true positives / (true positives + false negatives): among actual positives, how many are found. In LLM safety, a safety classifier with high recall catches most harmful outputs but may over-refuse benign requests (low precision). A high-precision classifier rarely over-refuses but misses some harmful outputs. The tradeoff is tuned via the classification threshold, and the choice depends on the relative cost of false positives (over-refusal) versus false negatives (harmful output).

Expect these follow-ups

• How does the F1 score combine precision and recall, and when would you prefer a different combination?
• In information retrieval for RAG, which metric is more important for the retrieval step and why?

As asked

Without looking anything up, walk me through computing the gradient of the attention output O = softmax(QK^T / sqrt(d_k)) V with respect to Q. What shape is the gradient and what is the computational challenge in computing it?

Sample answer outline

The output O has shape (n, d_v). The gradient dL/dQ requires differentiating through the softmax and the matrix products. dL/dA (where A = softmax(S)) uses the Jacobian of softmax: each row of A is a softmax, so the Jacobian is diag(a_i) - a_i a_i^T per row. Then dL/dS = dL/dA times this Jacobian applied row-wise. Finally dL/dQ = dL/dS / sqrt(d_k) times K. The computational challenge is that storing the full softmax Jacobian requires O(n^2) memory, which FlashAttention avoids by fusing the backward pass using the online normalization trick.

Expect these follow-ups

• How does FlashAttention avoid explicitly computing the softmax Jacobian during backpropagation?
• What simplification of the softmax gradient arises when you use the recomputed attention output rather than intermediate activations?