AI company interview questions

Scaled dot-product attention compares queries with keys, divides the scores by the square root of the key dimension, applies softmax, and mixes values using those weights. Scaling counters dimension-driven score differences that can prematurely saturate attention and weaken learning signals; it does not force uniform attention.

Queries $Q$ describe what to seek, keys $K$ provide matching features, and values $V$ contain information to aggregate. In self-attention, separate learned projections typically give $Q=XW_Q$, $K=XW_K$, and $V=XW_V$. With $n$ queries, $m$ keys, query/key dimension $d_k$, and value dimension $d_v$, scores have shape $n\times m$ and outputs $n\times d_v$. Compute $A=\operatorname{softmax}(QK^\top/\sqrt{d_k})$ row-wise, then $O=AV$. Each weight row sums to one.

The square root follows from variance, the expected squared deviation from a mean. Let $s=q^\top k=\sum_{i=1}^{d_k}q_i k_i$. Under the simplifying assumption that all relevant components are independent, zero-mean, and unit-variance, each product has mean zero and second moment $\mathbb{E}[q_i^2k_i^2]=\mathbb{E}[q_i^2]\mathbb{E}[k_i^2]=1$. Cross-product covariances vanish, so $\operatorname{Var}(s)=d_k$ and its standard deviation is $\sqrt{d_k}$. Since $\operatorname{Var}(s/c)=\operatorname{Var}(s)/c^2$, dividing by $\sqrt{d_k}$ gives unit variance. This concerns distributional scale, not monotonic growth of every dot product; trained correlations can violate the assumptions.

Softmax computes $p_i=e^{z_i}/\sum_j e^{z_j}$; score differences determine sharpness. Its derivative is $\partial p_i/\partial z_j=p_i(\delta_{ij}-p_j)$, where $\delta_{ij}$ is one for equal indices and zero otherwise. These local derivatives become small near one-hot attention. 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]$: ranking is preserved, but the second position retains more weight.

Standard scaling is a simple default that preserves ranking while compensating for dimension. Omitting it may work for small dimensions or otherwise controlled scores, but makes sharpness more sensitive to dimension and vector magnitude. Dividing by $d_k$ instead gives variance $1/d_k$ under the assumptions, potentially making attention increasingly uniform and weakening discrimination.

An additional design option is temperature: $A=\operatorname{softmax}(QK^\top/(\sqrt{d_k}\tau))$, with $\tau>0$. Higher temperature smooths attention; lower temperature concentrates it. Learning temperature adds flexibility but also parameters and validation requirements. Adopt it with training and evaluation evidence, not as an untested replacement in an existing model.

Normalizing queries and keys to use cosine similarity reduces dependence on vector norms, but changes available signals and usually requires scale calibration. Separately, subtracting each row's maximum before softmax prevents exponential overflow without changing probabilities. Because it preserves score differences, this numerical stabilization cannot replace scaling or undo saturation.

The following is recommended engineering practice, not an author-provided deployment specification. Establish batch, head, sequence, and feature axes. Transpose only the final two key axes and scale using per-head $d_k$, not total embedding width or value dimension. Add masks, apply softmax over keys, then multiply by $V$. Causal masks exclude future positions; padding masks exclude invalid tokens.

A manual sequence is scores = (Q @ K.transpose(-2, -1)) / math.sqrt(Q.shape[-1]), add an additive mask, then weights = softmax(scores, dim=-1) and output = weights @ V. Excluded positions may receive negative infinity, but fully masked rows need explicit handling, such as defined zero outputs, to avoid NaNs. Check library defaults for built-in scaling and Boolean mask semantics.

Verify more than shapes. Reproduce the two-key example; assigning values zero and ten gives a scaled output approximately $1.19$. Check valid-row sums near one, zero masked weights, and finite forward and backward results. Generate independent random vectors satisfying the assumptions across several dimensions and compare empirical score variances before and after scaling. Finite samples should approximate, not exactly match, the theoretical values.

These are additional operational recommendations; the sources provide no production measurements. For unstable loss, NaNs, or heads concentrating on one position unusually early, sample query/key norms, pre- and post-scaling score standard deviations, maximum weights, and gradient norms. Attention entropy, $H=-\sum_i p_i\log p_i$, helps describe concentration but must be interpreted alongside valid key counts. Low entropy is not inherently faulty: a task may require focused attention.

Typical faults include missing or duplicate scaling, incorrect head dimensions, reversed masks, fully masked rows, and low-precision multiplication overflow. Stable softmax assumes finite inputs; it cannot repair overflow that already occurred in the dot product. Compare a small case against a higher-precision reference and adjust accumulation precision if necessary. Confirm consistent scaling and masking between training and inference.

When changing fused kernels or head counts, compare outputs and gradients within appropriate tolerances before evaluating task quality, latency, and memory. The formula alone establishes no performance result. Interview recap: compare with $QK^\top$, control score scale using $1/\sqrt{d_k}$, and aggregate $V$ through softmax weights. The square root follows from an assumption-dependent variance argument; reliable execution still requires mask, precision, and distribution checks.