Training and optimisation interview questions
35 questions with model answers
- junior 7
- mid 15
- senior 13
35 training and optimisation questions asked in AI engineer, ML engineer and applied AI engineer interviews, each with a model answer you can say out loud in a minute or two: 7 junior, 15 mid-level and 13 senior.Updated
Model answer
Work in log-probs and subtract the row max before the log-sum-exp, so nothing overflows.
import numpy as np
def masked_ce(logits, y, mask, eps=0.1, T=1.0):
# logits (N, V), y (N,) ints, mask (N,) 0/1
z = logits / T
z = z - z.max(-1, keepdims=True)
logp = z - np.log(np.exp(z).sum(-1, keepdims=True))
y = np.where(mask > 0, y, 0) # ignored ids safe
nll = -logp[np.arange(len(y)), y]
loss = (1 - eps) * nll - eps * logp.mean(-1)
m = mask.astype(loss.dtype)
return (loss * m).sum() / max(m.sum(), 1.0)
Smoothing mixes the one-hot target with a uniform one, which is a weighted sum of the NLL and the mean negative log-prob. Interviewers check you average over unmasked tokens only, guard an all-masked batch, and never index with ignore labels like -100.
Model answer
Three stages. Recall: cheap filters for LaTeX or MathML markup, math symbols and known math domains, with extraction that keeps equations, which standard HTML-to-text drops. Precision: a fast classifier such as fastText, trained on seed math pages against random web, rerun with new seeds as DeepSeekMath did. Quality: dedupe, strip benchmark contamination by n-gram matching, and score reasoning content with a small LLM grader.
Validate by training small models on each candidate mix and comparing math benchmarks.
Model answer
KL(P‖Q) is the expected extra log-loss from using Q when data comes from P, and it's asymmetric. Cross-entropy training minimises forward KL from data to model, which punishes missing anything the data covers, so it's mass-covering; reverse KL is mode-seeking. If Q gives zero probability where P doesn't, KL is infinite. In RLHF, a penalty on KL from policy to reference keeps the policy close and limits reward hacking.
Model answer
MSE for regression, cross-entropy for classification. Cross-entropy is the negative log-likelihood of the right class; with softmax its gradient is just predicted minus true, so confident mistakes get big updates, while MSE on probabilities gives them small ones.
LoRA freezes the weights and trains a small low-rank update, so fine-tuning is far cheaper.
Model answer
Adam keeps running averages of each gradient and its square, and divides the first by the root of the second, so every weight gets its own step size. Both start at zero, so bias correction divides by one minus beta to the power t.
def adamw(w, g, m, v, t, lr=1e-3, b1=0.9,
b2=0.999, eps=1e-8, wd=0.01):
m = b1 * m + (1 - b1) * g
v = b2 * v + (1 - b2) * g * g
m_hat = m / (1 - b1 ** t)
v_hat = v / (1 - b2 ** t)
w = w - lr * (m_hat / (v_hat ** 0.5 + eps) + wd * w)
return w, m, v
An L2 term in the loss gets divided by that root too, so high-gradient weights decay less. AdamW decouples it: shrink the weights directly, outside the adaptive step.
Model answer
Policies are deterministic or stochastic, on- or off-policy. Policy gradients raise the log-probability of high-advantage actions. TRPO bounds each update with a KL trust region; PPO gets similar stability from a cheap clipped ratio. GAE weights TD errors exponentially, with lambda trading bias for variance.
Scroll sideways to see the whole diagram
In multi-head attention, each head projects its own queries, keys and values. Every key and value is cached for every token.
Text version
- In multi-head attention, each head projects its own queries, keys and values. Every key and value is cached for every token.
- Grouped-query attention keeps all the query heads but lets each group read one shared key-value head. Multi-query attention is the extreme: one K, V for all.
- The cache shrinks in proportion to the key-value heads, which is why most large open models use GQA.
Ainslie et al. (2023), GQA; Llama 2 70B's configuration: 64 query heads, 8 key-value heads (Touvron et al. 2023, Section 2.2).
Linear attention swaps softmax for a kernel map: linear cost, usually lower quality. GQA keeps exact attention and shares key-value heads across query groups, shrinking the KV cache.
Model answer
Train only on impressions older than 24 hours, so labels are final, or model the delay explicitly. High-cardinality IDs go through hashed or learned embeddings, plus crosses and historical click counts; a wide-and-deep or trees-plus-linear model is a solid start. Keep the 0.7% rate, or downsample negatives and correct the probabilities afterwards.
Assess with log loss and calibration, since bidding uses the probability itself, plus AUC, on a time-based split. Then an online A/B test on revenue and click-through, and monitor calibration drift daily.
Model answer
XGBoost learns a default direction for missing values at every split: it tries sending them left and right and keeps whichever lowers the loss, so missingness itself can carry signal.
Detect overfitting by tracking training against validation loss over boosting rounds, on a time-based or grouped split. Control it with early stopping, shallower trees, larger min child weight, a lower learning rate, row and column subsampling, and L1 or L2 on leaf weights.
Model answer
Each output channel has one kernel spanning all input channels, so weights are in-channels times out-channels times kernel height times width. Bias adds one per output channel.
Example: 3 in, 64 out, 3 by 3 kernel gives 1,728 weights, 1,792 with bias. Groups divide the weight count.
Model answer
PPO is stable and simple on-policy; SAC is off-policy and more sample-efficient, which matters when simulation is expensive.
Reward progress along the route, penalise speed above the limit with a penalty that grows with the excess, and add small penalties for jerk and harsh braking. Make collisions and red lights terminations or hard constraints, for example via a Lagrangian method, not weights the agent can trade away. Against exploits: reward route progress, not raw speed, so circling earns nothing; bound per-step rewards; and inspect rollouts for reward hacking in adversarial simulated scenarios.
Model answer
Target: delivered later than the promised time plus a tolerance agreed with the business. Features may only use what's known at order creation: store prep history, current queue, nearby courier supply, distance, weather, time. Anything later, like actual pickup time, is leakage.
Split by time: train on earlier weeks, validate on later ones, with a gap so labels and aggregates don't overlap, and compute aggregates point-in-time. Metrics: PR-AUC for the rare class, calibration if the probability feeds ETAs, and recall at the alert volume ops can handle.
Model answer
Put preprocessing inside the pipeline, so cross-validation fits imputers and scalers on training folds only.
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import (
StratifiedKFold, cross_val_score)
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import (
OneHotEncoder, StandardScaler)
prep = ColumnTransformer([
("num", make_pipeline(SimpleImputer(strategy="median"),
StandardScaler()), num_cols),
("cat", OneHotEncoder(handle_unknown="ignore"),
cat_cols),
])
model = make_pipeline(
prep, LogisticRegression(class_weight="balanced",
max_iter=1000))
cv = StratifiedKFold(5, shuffle=True, random_state=0)
scores = cross_val_score(model, X, y, cv=cv,
scoring="average_precision")
From this baseline, try gradient boosting and tune with the same CV. For imbalance, class weights or threshold tuning beat blind oversampling. Report PR-AUC, plus precision and recall at a threshold chosen on validation data, then read the errors.
Model answer
Choose by the cluster shape you expect and data size.
K-means minimises within-cluster squared distance, assumes compact similar-size clusters, and needs k. Hierarchical merges by a linkage rule into a dendrogram, suiting small, nested data. DBSCAN grows dense regions and flags outliers, finding any shape but struggling with varying density. Gaussian mixtures fit soft, elliptical clusters by EM. Spectral clustering uses eigenvectors of a similarity graph, handling non-convex shapes on moderate data.
Model answer
Walk it as a pipeline: where the data came from, how it was filtered, deduplicated and labelled, and how you split it. Tie each choice to an effect you saw, for example deduplication preventing train-test leakage, or label noise capping accuracy. End with what you'd change and how you measured it.
Model answer
BatchNorm normalizes each feature across the batch. LayerNorm normalizes each example across its own features. So BatchNorm depends on the batch: training uses batch statistics and updates running averages, inference uses those stored running statistics. LayerNorm computes the same thing in both modes, which is why transformers and small or variable batches use it.
Bugs usually show up as a train-eval gap: a missing model.eval(), tiny or unshuffled batches, stale running statistics after fine-tuning, or the wrong axis. Compare train and eval outputs on one batch and inspect the running statistics.
Model answer
With hidden layers, every unit in a layer computes the same output and gets the same gradient, so they stay identical forever: the layer acts like one unit. Worse, with all weights zero the gradients reaching earlier layers are zero too, so they don't move at all.
Linear and logistic regression are fine. There are no hidden units to be symmetric, the loss is convex, and gradient descent from zero reaches the same optimum as from anywhere. Networks therefore use random initialization scaled to keep activation variance steady, such as Xavier or He initialization.
Model answer
Outcomes aren't exclusive, so give each its own binary head on a shared trunk with sigmoid cross-entropy, not one softmax.
For a funnel like impression, click, purchase, model the chain: p(click) times p(purchase given click) equals p(click and purchase). Supervise p(click) and the product on all impressions, as ESMM does, so the conditional head is learned over the whole space instead of only clicked items, avoiding sample-selection bias. Weight the task losses, watch for one task dominating the shared gradients, and check each head's calibration.
Model answer
A higher learning rate. It scales every update equally, so early layers stay starved while late ones turn unstable. The others attack the cause: skip connections give gradients a direct path, ReLU has slope one when active instead of saturating, and batch normalization keeps activations in a healthy range.
Model answer
Boosting adds weak learners in sequence, each fixing current mistakes. For labels plus or minus one, the margin is label times ensemble score: positive means correct. AdaBoost minimises the exponential loss of the margin, a smooth surrogate for 0-1 error. Large training margins tighten generalization bounds but don't guarantee test accuracy, especially with noisy labels.
Model answer
Focal loss scales cross-entropy by (1 − p_t) to the power gamma, so well-classified examples add little and training focuses on hard ones; alpha rebalances the classes. Work from logits throughout.
import torch
def focal_loss(z, y, alpha=0.25, gamma=2.0):
# z: logits, y: float 0/1 targets, same shape
# stable BCE: max(z,0) - z*y + log(1+exp(-|z|))
ce = (z.clamp(min=0) - z * y
+ torch.log1p(torch.exp(-z.abs())))
p_t = torch.exp(-ce)
a_t = alpha * y + (1 - alpha) * (1 - y)
return (a_t * (1 - p_t) ** gamma * ce).mean()
Interviewers check you never take the log of a sigmoid output, that p_t comes from exp of minus the loss, and that gamma 0 recovers alpha-weighted cross-entropy.
Model answer
Start with the data, since a distribution mismatch splits the curves whatever the model does. Check train and test share normalisation, scale and noise level. Then the loop: the loss compares the output with the clean target, not the noisy input; zero_grad, backward and step run in order; train and eval modes are set. Overfit one batch first: if loss won't fall there, the bug is in the code.
Model answer
Variable length: pad and carry a mask so padding never enters pooling. Order stops mattering: encode each point on its own and pool with a symmetric function, the Deep Sets recipe. Missing features: fill with zero and append a per-feature missing flag, so absence itself is a signal.
import torch
import torch.nn as nn
class SetClassifier(nn.Module):
def __init__(self, d=6, h=64, n_cls=2):
super().__init__()
self.phi = nn.Sequential(
nn.Linear(2 * d, h), nn.ReLU(),
nn.Linear(h, h), nn.ReLU())
self.head = nn.Linear(h, n_cls)
def forward(self, x, keep):
# x: (B, T, 6), NaN where a feature is missing
# keep: (B, T) bool, True for real points
miss = torch.isnan(x).float()
x = torch.nan_to_num(x, nan=0.0)
z = self.phi(torch.cat([x, miss], -1))
z = z * keep.unsqueeze(-1).float()
n = keep.sum(1, keepdim=True).clamp(min=1)
return self.head(z.sum(1) / n)
Interviewers check the masking and a permutation test: shuffled points must give the same logits.
Model answer
Overfitting: low training error, much higher validation error; add data or regularise. Underfitting: both high; add capacity. That's the bias-variance trade-off. Imbalance: check class counts, use precision-recall metrics, re-weight or resample. Missing data: measure missingness per feature, impute, and add a missing flag.
Model answer
import numpy as np
def kmeans(X, k, iters=100, seed=0):
rng = np.random.default_rng(seed)
C = X[rng.choice(len(X), k, replace=False)]
for _ in range(iters):
d = ((X[:, None] - C[None]) ** 2).sum(-1)
a = d.argmin(1)
new = np.array([X[a == j].mean(0)
if (a == j).any() else C[j]
for j in range(k)])
if np.allclose(new, C):
break
C = new
return C, a
It minimises the within-cluster sum of squared distances: assign points to the nearest centroid, move centroids to their means, repeat. It reaches a local optimum, so use k-means++ and restarts. For train-inference mismatch, reuse the training scaler and monitor distances to centroids; drift means refit.
Model answer
import numpy as np
def softmax_ce(logits, y):
z = logits - logits.max(1, keepdims=True)
p = np.exp(z)
p /= p.sum(1, keepdims=True)
n = len(y)
loss = -np.log(p[np.arange(n), y] + 1e-12).mean()
grad = p.copy()
grad[np.arange(n), y] -= 1
return loss, grad / n
def train(X, y, k, lr=0.1, epochs=100):
W = np.zeros((X.shape[1], k))
b = np.zeros(k)
for _ in range(epochs):
loss, g = softmax_ce(X @ W + b, y)
W -= lr * X.T @ g
b -= lr * g.sum(0)
return W, b
Subtract the row maximum before exponentiating, for numerical stability. The gradient with respect to the logits is probabilities minus the one-hot label, averaged over the batch; the chain rule gives X transposed times that for the weights. Interviewers look for the stability trick, the averaging, and a finite-difference gradient check.
Model answer
def kmeans(points, init, max_iter=100, tol=1e-9):
cents = [list(c) for c in init]
for _ in range(max_iter):
groups = [[] for _ in cents]
for p in points:
d = [sum((a - b) ** 2 for a, b in zip(p, c))
for c in cents]
groups[d.index(min(d))].append(p)
new = []
for c, g in zip(cents, groups):
if not g:
new.append(c)
continue
n = len(g)
new.append([sum(x) / n for x in zip(*g)])
shift = max(sum((a - b) ** 2
for a, b in zip(c, m))
for c, m in zip(cents, new))
cents = new
if shift <= tol:
break
return cents, groups
State each rule before coding. Given initial centroids make runs deterministic. A tie goes to the lowest-index centroid, which taking the first minimum does. An empty cluster keeps its previous centroid; reseeding at the farthest point is the alternative. Stop when no centroid moves more than a tolerance, or at the iteration cap. Interviewers weigh these edge cases more than the loop.
Model answer
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import StandardScaler
def baseline(df, num, cat, y):
df = df.copy()
for c in cat:
df[c] = df[c].str.strip().str.lower()
nums = make_pipeline(SimpleImputer(strategy="median"),
StandardScaler())
cats = make_pipeline(
SimpleImputer(strategy="most_frequent"),
OneHotEncoder(handle_unknown="ignore"))
prep = ColumnTransformer([("n", nums, num),
("c", cats, cat)])
clf = LogisticRegression(max_iter=1000)
model = make_pipeline(prep, clf)
return cross_val_score(model, df[num + cat], df[y],
cv=5, scoring="roc_auc")
Look first: types, missing rates, label balance, and columns that leak the target, like anything recorded after the outcome. Normalise strings, and keep imputation, scaling and encoding inside a pipeline so they're fitted only on training folds. Compare against predicting the majority class, with ROC AUC or PR AUC under stratified cross-validation.
Model answer
Datasets build windows lazily from contiguous arrays, with input length, horizon and stride in config. Split by time with a gap so no window straddles the boundary, and fit normalisation on training data only.
Reproducibility: seed every RNG, record config, code version and a data hash, and use deterministic kernels where it matters. Distribute with DDP or FSDP and a distributed sampler. Checkpoints hold model, optimizer, scheduler, RNG states and step, written atomically, so resuming is exact. Evaluate with walk-forward backtests per horizon against naive baselines like last value.
Model answer
Multinomial Naive Bayes over count features, with Laplace smoothing.
import numpy as np
class MultinomialNB:
def __init__(self, alpha=1.0):
self.alpha = alpha
def fit(self, X, y):
# X (N, V) counts, y (N,) labels
self.classes = np.unique(y)
n, V = len(self.classes), X.shape[1]
self.log_prior = np.empty(n)
self.log_lik = np.empty((n, V))
for i, c in enumerate(self.classes):
Xc = X[y == c]
self.log_prior[i] = np.log(len(Xc) / len(X))
cnt = Xc.sum(0) + self.alpha
self.log_lik[i] = np.log(cnt / cnt.sum())
return self
def predict(self, X):
s = X @ self.log_lik.T + self.log_prior
return self.classes[s.argmax(1)]
Interviewers check you sum log-probabilities instead of multiplying, which underflows; that smoothing stops one unseen feature zeroing a class; and that you can name the variant to swap in: Gaussian for continuous features, Bernoulli for binary. For probabilities, normalize the scores with log-sum-exp.
Model answer
import numpy as np
def sigmoid(z):
out = np.empty_like(z, dtype=float)
pos = z >= 0
out[pos] = 1 / (1 + np.exp(-z[pos]))
ez = np.exp(z[~pos])
out[~pos] = ez / (1 + ez)
return out
def train(X, y, lr=0.1, steps=1000):
# X (N, D), y (N,) of 0/1
w, b = np.zeros(X.shape[1]), 0.0
for _ in range(steps):
p = sigmoid(X @ w + b)
err = p - y # dLoss/dz for log loss
w -= lr * X.T @ err / len(y)
b -= lr * err.mean()
return w, b
The usual bugs: a sigmoid that overflows for large negative inputs; a derivative written as s(1 minus z) instead of s(1 minus s); and multiplying by that derivative again, when log loss already gives a logit gradient of exactly p minus y. Also check the sign of the update, dividing by N, and a y of shape (N, 1) silently broadcasting to (N, N). Confirm with a finite-difference gradient check.
Model answer
Sigmoid and tanh saturate, so gradients vanish; ReLU is cheap but units can die; GELU and SiLU are smooth and standard in transformers.
The loss follows the output distribution: MSE for Gaussian-ish regression, Huber against outliers, sigmoid plus binary cross-entropy, softmax plus cross-entropy.
m = b1 * m + (1 - b1) * g
v = b2 * v + (1 - b2) * g * g
m_hat = m / (1 - b1 ** t)
v_hat = v / (1 - b2 ** t)
w -= lr * m_hat / (np.sqrt(v_hat) + eps)
Adam keeps momentum and a per-parameter scale, bias-corrected. Defaults: beta1 0.9, beta2 0.999, eps 1e-8; tune the learning rate first.
Model answer
Deep trees have low bias but high variance, so averaging many helps. For B trees, each with variance σ² and pairwise correlation ρ, the average has variance ρσ² + (1 − ρ)σ²/B. More trees shrink only the second term; ρσ² is a floor.
So forests attack correlation: bootstrap samples, plus a random subset of features at each split, so one dominant feature doesn't make every tree alike. Fewer features per split lowers ρ but weakens each tree.
Model answer
Leakage is training on information unavailable at prediction time: future data, target-derived features, or scalers fitted before splitting. Split first.
For missing data, ask why it's missing, then impute from training statistics and add a missing indicator.
MSE fits the conditional mean and punishes outliers; MAE fits the median. Logistic regression uses log loss: the Bernoulli likelihood, convex in the weights. MSE on a sigmoid is non-convex and its gradients vanish when saturated.
Model answer
T = x.size(1)
causal = torch.ones(T, T, dtype=torch.bool,
device=x.device).tril()
keep = causal[None, None] & pad[:, None, None, :]
att = att.masked_fill(~keep, float("-inf"))
att = att.softmax(-1)
# later: predict x[t+1] from tokens up to t
logits = model(x[:, :-1])
loss = F.cross_entropy(logits.transpose(1, 2),
x[:, 1:], ignore_index=PAD)
Classic bugs: triu where you meant tril; masking after the softmax, or with zero instead of minus infinity; a padding mask on queries instead of keys; a fully masked row turning into NaN; targets not shifted by one; pads counted in the loss. Then the loop: zero_grad, train mode, learning rate.
Model answer
Jobs are declarative specs: image, data version, resources. The scheduler gives each team a guaranteed quota and lets idle capacity be borrowed, preemptibly. Gang scheduling places all of a job's workers at once, topology-aware, so no half-started job holds GPUs.
Datasets are immutable, content-hashed snapshots referenced by version. Tracking records config, commit, data version, metrics and artifacts per run.
Jobs write sharded checkpoints asynchronously to object storage, and preemption triggers one; on node failure, the job restarts from the latest. Monitor GPU utilization and queue time per team.