Peter's Weblogs · Deep Learning · Mathematics

Linear Algebra for Deep Learning: Tensors, Matrix Products, Attention and Low-Rank Updates

A neural network looks like a stack of exotic components, but most of what happens inside one is multiplying arrays of numbers. This piece walks through the linear algebra that matters in deep learning, with the shapes and the arithmetic written out so that each idea can be checked by hand.

Technical Guide Mathematics for deep learning ~12 min read

Deep learning frameworks store data and parameters as tensors and spend most of their compute time on matrix multiplication. Knowing a handful of linear algebra ideas, and reading shapes fluently, explains how layers, attention, backpropagation, compression and fine-tuning work. For a curated list of courses, books and libraries, see the Linear Algebra for AI directory on this site; this article is the explanation.

01Scalars, Vectors, Matrices and Tensors

The vocabulary is a ladder of dimensions. A scalar is one number, a vector is a list of numbers, a matrix is a grid, and a tensor is the general n-dimensional array. In deep learning the dimensions carry meaning. A batch of token embeddings is commonly a tensor of shape (batch, sequence, features), and a batch of images is (batch, channels, height, width) or a permutation of it.

The habit worth building is to write the shape next to every quantity. Most bugs in model code are shape errors, and most explanations in papers become clearer once the shapes are written down.

02Vectors, Dot Products and Embeddings

An embedding is a vector that represents a word, image or document. Two operations on vectors do much of the work. The dot product of two vectors multiplies matching components and sums them, and the norm measures length:

a · b = Σ aᵢbᵢ ‖a‖ = √(a · a) cos(a, b) = (a · b) / (‖a‖ ‖b‖)

The dot product is large when two vectors point the same way, which is why it works as a similarity score. Cosine similarity divides out the lengths so that only direction counts. Vector search, retrieval and attention all rest on this one operation, which is why the Vector Databases page on this site is mostly about doing it fast over millions of vectors.

03A Dense Layer Is a Matrix Product

A fully connected layer computes a linear map followed by a nonlinearity. For one input vector x of length d_in and a weight matrix W of shape d_out × d_in:

y = φ(Wx + b)

With a batch of n inputs stacked as the rows of X (shape n × d_in), the whole batch is one matrix product, Y = φ(XWᵀ + b), which is exactly the operation GPUs are built to run quickly. In many frameworks the weight is stored as d_in × d_out so that the product is simply XW; the maths is the same.

The nonlinearity φ is not decoration. Without it, stacking layers collapses: W₂(W₁x) = (W₂W₁)x, so any depth of linear layers is a single linear map. Nonlinearities between the products are what let depth add expressive power.

Counting is easy once you know the shapes. A 4096 × 4096 layer has 4096² = 16,777,216 weights. Multiplying a batch of 32 inputs through it takes about 2 × 32 × 4096 × 4096 = 1,073,741,824 floating-point operations, roughly 1.07 GFLOPs, because an (m × k) by (k × n) product costs about 2mkn operations (one multiply and one add per term).

Input Xn × d_modelQ = XW_Qn × d_kK = XW_Kn × d_kV = XW_Vn × d_vQKᵀ / √d_kn × n scoressoftmax · Vn × d_v· W_On × d_model
One attention head is four matrix products and a softmax. Every arrow changes a shape in a predictable way, and the n × n score matrix in the middle is the one that grows with the square of the sequence length.

04Attention Is Four Matrix Products and a Softmax

In the Transformer (Vaswani et al., 2017), an attention head projects the input X (n × d_model) into queries, keys and values with three weight matrices, then computes:

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

QKᵀ is an n × n matrix whose entry (i, j) is the dot product of query i with key j, so it scores how much token i should look at token j. The softmax turns each row into weights that sum to one, and multiplying by V takes a weighted average of the value vectors.

The division by √d_k has a linear algebra reason. If the components of q and k are independent with zero mean and unit variance, their dot product has variance d_k, so unscaled scores grow with the dimension and push the softmax into regions with tiny gradients. Dividing by √d_k keeps the variance near one.

The n × n score matrix is also the memory problem in long-context models. At n = 8,192 it has 67,108,864 entries per head per layer, about 128 MiB in 16-bit precision if stored in full. Memory-efficient attention kernels exist to avoid materialising it, but the quadratic growth is a property of the algebra. The same computation in code is three lines:

scores  = torch.einsum('bhqd,bhkd->bhqk', Q, K) / d_k ** 0.5
weights = scores.softmax(dim=-1)
out     = torch.einsum('bhqk,bhkd->bhqd', weights, V)

einsum is worth learning because it states the shapes of a product explicitly: letters that repeat on the input side and vanish on the output side are summed over.

05Backpropagation Is Matrix Calculus

Training needs the gradient of the loss with respect to every weight, and the chain rule applied to matrix operations gives clean rules. For Y = XW with upstream gradient G = ∂L/∂Y:

∂L/∂W = Xᵀ G ∂L/∂X = G Wᵀ

So the backward pass of a matrix multiplication is two more matrix multiplications, with transposed operands. This is why training costs a small multiple of inference, and why the same hardware and libraries serve both. In general, frameworks never build the full Jacobian matrix, which would be enormous; they compute vector-Jacobian products, multiplying the upstream gradient by the local Jacobian directly. Parr and Howard's paper, listed below, is a readable guide to the notation.

06Convolutions Are Matrix Products Too

A convolution slides a small kernel over an input. It is a linear operation, so it can be written as a matrix product by unrolling the input patches into columns (a technique often called im2col) and multiplying by the flattened kernels. Many implementations lean on this idea, which is one reason highly tuned matrix multiplication routines matter across all architectures, not only Transformers.

07Eigenvectors, SVD and Low-Rank Structure

Any matrix W can be factored by the singular value decomposition, W = UΣVᵀ, where the diagonal of Σ holds the singular values in decreasing order. Keeping only the largest r of them gives the best rank-r approximation of W in the Frobenius and spectral norms (the Eckart-Young theorem). That fact underlies PCA, latent-factor recommenders, and compression of trained weight matrices by truncation.

The same idea made fine-tuning cheap. LoRA (Hu et al., 2021) freezes the pretrained weight W and learns a low-rank update instead:

W' = W + (α / r) · B A B: d × r A: r × d r ≪ d

For a 4096 × 4096 matrix and rank r = 8, the update has 4096 × 8 + 8 × 4096 = 65,536 trainable values, compared with 16,777,216 in the full matrix, about 0.39%. The paper reports that low-rank updates can match full fine-tuning quality on the tasks it tested while training far fewer parameters; how well that holds depends on the task and the rank chosen.

08Initialisation, Norms and Stability

A deep network multiplies by many matrices in sequence, so the size of those matrices matters. The spectral norm of a matrix, its largest singular value σ_max, bounds how much it can stretch a vector: ‖Wx‖ ≤ σ_max‖x‖. If layers repeatedly stretch signals, activations and gradients explode; if they shrink them, they vanish.

Initialisation schemes are linear algebra applied to this problem. Glorot (Xavier) initialisation sets the weight variance from the layer's fan-in and fan-out, Var(W) = 2 / (fan_in + fan_out), to keep signal variance roughly constant through the network. He initialisation uses Var(W) = 2 / fan_in for ReLU layers, to account for ReLU zeroing about half the activations. Normalisation layers and residual connections address the same problem from different angles.

Condition numbers in practice

The condition number of a matrix is the ratio of its largest to smallest singular value. A badly conditioned problem makes gradient descent slow and sensitive, which is a large part of why optimisers that rescale directions, and techniques that normalise inputs, help training.

09Hardware, Precision and Quantisation

Because so much of the work is dense matrix multiplication, accelerators dedicate hardware to it; NVIDIA's Tensor Cores are one example, and TPUs are built around matrix units. Mixed-precision training stores and multiplies in 16-bit or lower formats while accumulating in higher precision. Quantisation goes further, replacing the entries of weight matrices with low-bit integers plus scale factors, which is why the LLM Model Formats and Quantisation Techniques pages on this site are, underneath, about representing matrices compactly.

10Cheat Sheet: Concept to Where It Appears

ConceptWhere it appearsTypical shape or fact
Dot product, cosine similarityEmbeddings, retrieval, attention scoresTwo vectors of equal length give one number
Matrix multiplicationDense layers, projections, convolutions via im2col(m × k)(k × n) costs about 2mkn operations
TransposeQKᵀ, backward pass∂L/∂W = XᵀG, ∂L/∂X = GWᵀ
Softmax over rowsAttention weights, classificationEach row sums to one
SVD, low rankPCA, compression, LoRARank-r update stores r(d_in + d_out) values
Spectral norm, conditioningInitialisation, stability, optimisers‖Wx‖ ≤ σ_max‖x‖
Tensors, einsumBatched and multi-head operations(batch, heads, sequence, features)

11How to Study It

Start with the geometry (vectors as arrows, matrices as transformations), then practise reading shapes in real code, then move to decompositions. Re-deriving the attention shapes and the backward pass of a dense layer on paper is more valuable than memorising formulas. The sources below include a free video series, an open course, two free books and the original papers.

Sources and further reading: