Attention, normalisation and convolution from scratch

Scaled dot-product attention

Each query row of Q (n × d) scores every key row of K (m × d) by a dot product. The scores are scaled, turned into weights by a softmax over each row, and used to average the rows of V:

Attention(Q, K, V) = softmax(QKᵀ / √d) V

The √d is there because a dot product of two random d-dimensional vectors with unit-variance entries has variance d. Without the scaling, scores grow with the width of the model, the softmax saturates into a near one-hot vector, and its gradient all but vanishes.

A causal mask stops position i from attending to any position j > i, which is what makes a decoder autoregressive. Set the masked scores to −∞ before the softmax; e^−∞ = 0, so they get exactly zero weight. Setting them to zero instead is a common bug: a score of 0 still gets weight e⁰ = 1.

S = Q @ K.T / np.sqrt(d)
mask = np.tril(np.ones((n, n), dtype=bool))    # True where j <= i
S = np.where(mask, S, -np.inf)

Multi-head attention runs h of these in parallel on slices of width d/h and concatenates the results. The cost is O(n²d) in time and O(n²) in memory for the score matrix, and that quadratic term is the reason long-context models need tricks such as FlashAttention, which never stores the full matrix.

Layer norm and batch norm

Both subtract a mean, divide by a standard deviation, then apply a learned scale γ and shift β. They differ in which axis the statistics are taken over.

• Layer norm normalises each example over its features. It behaves identically in training and inference, does not care about the batch size, and is the standard in transformers.
• Batch norm normalises each feature over the batch. At inference there is no batch to average over, so it keeps running estimates of the mean and variance during training and uses those instead.

Two conventions catch people out. Both layers normalise with the biased variance (divide by n, not n − 1). And PyTorch's batch norm updates its running variance with the unbiased batch variance, while normalising the current batch with the biased one. The ε goes inside the square root, √(σ² + ε), so a constant feature gives zeros rather than a division by zero.

Convolution

What deep learning calls convolution is cross-correlation: the kernel slides over the input without being flipped. With input size H, kernel size k, stride s and zero padding p on each side, the output size is ⌊(H + 2p − k)/s⌋ + 1. Get that formula right before writing any loop; the rest is bookkeeping.

A double loop over output positions is the honest first answer. The fast answer builds every window at once with numpy.lib.stride_tricks.sliding_window_view and contracts it against the kernel with one einsum. That is the same idea as im2col, which turns convolution into one large matrix product and is how most libraries implemented it on GPUs.

What interviewers listen for

• Why √d, in one sentence about variance.
• −∞, not zero, in the mask, and what happens to a row that is entirely masked.
• Which axis each normalisation uses, and what batch norm does at inference.
• The output-size formula for a convolution, before any code.