AI company interview questions

Scaled dot-product attention compares queries with keys and mixes values using the resulting weights:

$$\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

Dividing by $\sqrt{d_k}$ offsets dimension-dependent score growth, reducing premature softmax saturation and weak gradients. It does not require uniform attention.

In self-attention, learned projections produce $Q=XW_Q$, $K=XW_K$, and $V=XW_V$ from input representations X. Queries describe what to seek, keys provide matching features, and values supply retrieved content. For n tokens, one head has Q and K shaped $n\times d_k$; V can be $n\times d_v$. Each row of $QK^\top$ compares one query against all keys. Row-wise softmax produces weights summing to one; multiplying by V returns $n\times d_v$ contextual representations.

Let $s=\sum_{i=1}^{d_k}q_i k_i$. Assume the relevant components are mutually independent, zero-mean, and unit-variance. Variance measures spread around the mean. Then $\mathbb{E}[q_i k_i]=0$ and $\mathbb{E}[q_i^2k_i^2]=\mathbb{E}[q_i^2]\mathbb{E}[k_i^2]=1$. Cross-product covariances vanish, giving $\operatorname{Var}(s)=d_k$ and standard deviation $\sqrt{d_k}$. Because $\operatorname{Var}(s/c)=\operatorname{Var}(s)/c^2$, choosing $c=\sqrt{d_k}$ yields unit variance. This is exact under these assumptions, not a guarantee for correlated, trained queries and keys.

Softmax converts scores into weights:

$$p_i=\frac{\exp(z_i)}{\sum_j\exp(z_j)}$$

, depends on score differences, not a shared offset. Large differences can concentrate nearly all weight on one position. Its derivative is:

$$\frac{\partial p_i}{\partial z_j}=p_i(\delta_{ij}-p_j)$$

Here, $\delta_{ij}$ is one for $i=j$ and zero otherwise; most derivatives become small near saturation. In an original illustrative example, $d_k=4$ and scores [4,0] give weights approximately [0.982,0.018]. Dividing by two gives [0.881,0.119]. With scalar values [10,0], the output changes from about 9.82 to 8.81. Ranking remains unchanged, but the weaker match contributes more. These are teaching numbers, not measurements.

Standard scaling is a statistically motivated default. Unscaled attention may behave adequately with small dimensions or small scores, but wider heads increase the risk of excessively sharp distributions at initialization. Sharp attention is not inherently wrong: selecting one position may be appropriate. The concern is concentration driven by scale rather than learned relevance.

Dividing by $d_k$ instead produces variance $1/d_k$ under the same assumptions, increasingly flattening score differences as dimension grows. A fixed divisor accommodates only a particular scale rather than adapting to head width. Smaller scores are therefore not automatically better: excessively flat weights can dilute useful content.

As additional engineering options, a learned temperature can adjust softmax sharpness when architecture experiments have adequate validation resources, but adds tuning and monitoring. Cosine attention normalizes queries and keys to unit length, changing similarity and discarding magnitude information; its temperature needs separate design. The original variance argument no longer applies unchanged. Do not alter an existing checkpoint's scaling without validation.

As recommended practice, first build an inspectable reference. Project inputs, split heads, read $d_k$ from Q's last dimension, and compute scores=(Q @ K.transpose(-2,-1))/sqrt(d_k). Do not substitute the full embedding width or sequence length. Apply masks, run softmax along the key axis, and multiply by V. Query/key feature dimensions and key/value position counts must agree.

Masks exclude padding or future positions, commonly by assigning negative infinity to forbidden scores. Each row needs a valid key, or an explicit policy such as skipping, returning zeros, or rejecting the input. Otherwise, fully masked rows can produce undefined results. Stable softmax subtracts the row maximum to avoid exponential overflow; because score differences remain unchanged, this does not replace scaling.

Reproduce the two-position example, then check output shape, unit weight sums on valid rows, zero masked weights, and finite backward gradients. Synthetic independent standard-normal Q/K samples across dimensions should show raw variance near $d_k$ and scaled variance near one, allowing sampling error. For a fused attention API, verify whether scaling is built in, then compare outputs and gradients against the reference to catch duplicate scaling.

The articles provide no production measurements; these are recommended operational checks. Sample score standard deviations before and after scaling, Q/K norms, maximum attention weights, and gradient norms by layer and head. Entropy $H=-\sum_i p_i\log p_i$ measures weight dispersion. Control the number of valid keys when comparing entropy so sequence-length differences are not mistaken for defects.

If many heads show low entropy and weak gradients early in training, investigate missing scaling, an incorrect embedding-based divisor, growing projection weights, or input distribution changes. Unexpectedly uniform weights suggest duplicate scaling or division by $d_k$. These are diagnostic clues, not verdicts: examine task loss and quality alongside them. Low entropy alone does not prove failure, and scaling does not solve every gradient problem.

For NaNs or infinities, check fully masked rows, nonfinite inputs, low-precision dot-product overflow before softmax, and kernel accumulation precision. When changing precision, fused kernels, or head configurations, compare outputs and gradients on controlled inputs with reasonable floating-point tolerances, and verify causal masks prevent future-information leakage. Interview recap: $QK^\top$ determines matches, softmax determines mixing weights, and V supplies content; $\frac{1}{\sqrt{d_k}}$ stabilizes score scale under explicit assumptions, with correctness and practical behavior verified through tests and observation.