Linear Algebra for Machine Learning Interviews

Read the right-hand side as: how long is $a$, times how long is $b$, times how aligned are they. Two of those three factors are usually noise when you compare embeddings.

In a sentence-embedding model the vector norm often correlates with how common or how long the input was — not with what it means. Two paraphrases can land in almost the same direction but at different radii. Cosine similarity strips the radii away:

This is also why almost every vector database ships an L2-normalise on write option. Once vectors are unit length, cosine similarity, inner product and (negated) squared Euclidean distance all rank identically — so you can use the fastest kernel your index supports.

Raw dot product ranks the semantically unrelated noise vector almost as highly as an exact match, purely because it is long. Cosine does not. That single failure mode is the whole answer.

Each of the $B\times h$ output entries is a dot product over $d$ terms, i.e. $d$ multiplies plus $d-1$ adds. Engineers round this to 2·B·d·h FLOPs because modern GPUs execute a fused multiply-add (FMA) as one instruction pair.

That last row is the one that separates candidates. Arithmetic intensity — FLOPs per byte moved — decides whether you are compute-bound or bandwidth-bound. During training, $B$ is large (thousands of tokens), intensity is high, and the GPU's tensor cores are the bottleneck. During autoregressive decoding, $B=1$ token: you read the entire weight matrix to do a single skinny matrix-vector product. Intensity collapses and you become memory-bandwidth-bound.

Applying it to a transformer layer

For sequence length $n$, model width $d$ and FFN width $4d$, per layer per token:

• QKV projections: 3 × (2·n·d·d) = 6n·d² FLOPs
• Attention scores $Q{K}^{\top }$: 2·n²·d FLOPs, and an n×n score matrix per head
• Attention output (scores × V): another 2·n²·d
• Output projection: 2n·d²
• FFN (two layers): 2·n·d·4d + 2·n·4d·d = 16n·d²

Total ≈ $24n{d}^2+4n^{2}d$. Attention only dominates once $n\gtrsim 6d$ — for a 4096-wide model that is a ~24k token context. Below that, the FFN is where your FLOPs go, which is why MoE routing targets the FFN and why FlashAttention targets memory rather than FLOPs.

Eigendecomposition: only for square matrices

An eigenvector is a direction the transformation does not rotate — it only stretches it by $\lambda$. This exists only for square matrices, and even then not always with a real, full basis (a rotation matrix in 2D has no real eigenvectors). For symmetric matrices you get the good case: real eigenvalues and an orthonormal $Q$, so $A=Q\Lambda Q^{\top }$.

SVD: works on everything

Geometric reading: any linear map is a rotation ($V^{\top }$), then an axis-aligned scaling ($\Sigma$), then another rotation ($U$). Singular values are non-negative and sorted. The link back to eigenvectors:

• $V$ holds the eigenvectors of $A^{\top }A$
• $U$ holds the eigenvectors of $A{A}^{\top }$
• ${\sigma }_i=\sqrt{{\lambda }_i}$ of either product

PCA is SVD of centred data

Given a centred data matrix $X\in {\mathbb{R}}^{n\times d}$ (n samples, d features), the covariance is $C=\frac{1}{n-1}{X}^{\top }X$. Diagonalising $C$ gives principal components — but you should never form $C$ explicitly. Squaring the data squares the condition number and destroys precision on ill-conditioned features. Instead take the SVD of $X$ directly:

Also be ready for the modern angle: truncated / randomised SVD is how you run PCA on a 10M×1000 matrix, and low-rank factorisation via SVD is the conceptual ancestor of LoRA — you approximate a big weight update with $\Delta W=BA$ where $B$ and $A$ are thin.

Why L1 gives sparsity and L2 does not

Look at the gradients. For L2 the penalty gradient is $2\lambda w$ — it shrinks proportionally, so as $w\to 0$ the pull also $\to 0$ and the weight asymptotes but never arrives. For L1 the subgradient is $\lambda \mathrm{sign}(w)$ — constant magnitude regardless of how small $w$ is. Once the data gradient is smaller than $\lambda$, the weight is pushed to exactly zero and pinned there.

The geometric version, if the interviewer prefers pictures: the L1 constraint region is a diamond with corners on the axes; the loss contours are ellipses; an ellipse first touches a diamond at a corner, and corners have zero coordinates. The L2 region is a smooth ball with no corners.

Conditioning and why it wrecks training

A large condition number means the loss surface is a long, thin valley. Gradient descent takes a huge step along the steep axis and a tiny one along the flat axis, so it bounces across the valley walls instead of walking down it. The convergence rate of gradient descent on a quadratic degrades like ${\left(\frac{\kappa -1}{\kappa +1}\right)}^{2t}$ — at $\kappa =10^4$ that is effectively no progress.

Set the shapes down first; they do most of the work.

Scalar derivation, then lift to matrices

Element-wise, $y_{bj}={\sum }_k{x}_{bk}{W}_{kj}+b_j$. So $\frac{\partial y_{bj}}{\partial W_{kj}}=x_{bk}$, and applying the chain rule while summing over every path through the batch:

Why the interviewer asks this

Three follow-on facts fall straight out of these formulas and are what they are really probing:

1. 1

   Activation memory

   Computing ∂L/∂W needs x. That is why the forward activations must be kept alive until the backward pass, why activation memory scales with batch × sequence × width, and how gradient checkpointing trades recompute for memory.
2. 2

   Vanishing/exploding gradients

   ∂L/∂x = gWᵀ means the gradient is repeatedly multiplied by weight matrices going backwards. If the spectral norms are consistently < 1 it vanishes; > 1 it explodes. Residual connections add an identity path so the product never fully decays.
3. 3

   Bias gradients sum, weights do not average

   ∂L/∂b sums over the batch. Whether your effective learning rate scales with batch size depends on whether your loss uses sum or mean — a classic source of 'it worked at batch 32 but diverged at batch 512'.