Linear algebra of modern ML: autodiff, LoRA/GaLore, tensor decompositions, attention

Before this lesson: lesson 01 (matrices as maps and the views of matrix multiplication), lesson 02 (rank) and lesson 08 (the SVD and Eckart–Young–Mirsky). Time: about 85 minutes, plus the notebook.

---

Why this matters #

Four central tools of modern ML are linear algebra in disguise. Every gradient is a chain of transposed Jacobians applied to a vector (VJPs). LoRA and GaLore exploit low rank, which is Eckart–Young at work. CP and Tucker extend the SVD to data with more than two indices, such as climate fields over space × time × variable × ensemble member. An attention head is a low-rank bilinear form (which token to read from) plus a low-rank linear map (what to copy), acting on a residual stream that is a vector space. The same ideas decide memory and compute when you fine-tune or train a scientific foundation model.

1. Predict first #

2. The picture #

A derivative is a linear map. Near $x$, $f(x + \delta) \approx f(x) + J\delta$, with the Jacobian $J \in \mathbb{R}^{m\times n}$: the matrix of all partial derivatives $J_{ij} = \partial f_i/\partial x_j$. For a chain $f = f_L \circ \dots \circ f_1$ the chain rule gives a product $J = J_L\cdots J_1$. You never need $J$ itself, only $J$ times a vector, and the product can be evaluated from either end:

• right to left with an input direction $v$: $J_L(\cdots(J_1v))$. This is forward mode (JVP);
• left to right with an output weighting $u$: $((u^\top J_L)J_{L-1})\cdots J_1$. This is reverse mode (VJP), and backpropagation is its special case $u = 1$ for a scalar loss.

Each order costs only matrix-vector products. Which one wins depends on the shape of $J$.

Interactive figure: loading…

Why reverse mode wins for a scalar loss: the multiplication order. Matrix multiplication is associative, so $J_3J_2J_1$ is the same matrix whatever pair you multiply first. The cost is not the same. A product of an $a\times b$ matrix and a $b\times c$ matrix needs $abc$ multiplications. Take a loss after two hidden layers of width 1000 on an input of size 1000: $J_3$ is $1\times1000$, $J_2$ and $J_1$ are $1000\times1000$.

1. From the left (reverse). $J_3J_2$ is $(1\times1000)(1000\times1000)$: $10^6$ multiplications, and the result is again a $1\times1000$ row. Then $(J_3J_2)J_1$: another $10^6$. Total $2\times10^6$.
2. From the right (forward). $J_2J_1$ is $(1000\times1000)(1000\times1000)$: $10^9$ multiplications. Then $J_3(J_2J_1)$: $10^6$. Total about $10^9$, 500 times more.

Starting at the scalar loss keeps every intermediate result a single row vector. Starting at the input carries one column per input, so the work grows with the number of inputs.

Interactive figure: loading…

Low rank = few outer products. A rank-$r$ matrix $BA = \sum_{i=1}^r b_ia_i^\top$ is stored with $r(d + k)$ numbers instead of $dk$. LoRA restricts a weight update to this form.

Interactive figure: loading…

A tensor is an array with more indices. The simplest tensor is a rank-1 tensor $a \circ b \circ c$ with entries $a_ib_jc_k$. A sum of $R$ of them is a CP model; a small core with a basis per index is a Tucker model.

Attention moves information between tokens. Token $i$ forms a weighted average of (projected) other tokens. The weights come from a bilinear form "query of $i$ against key of $j$"; the projection decides what is copied and where it is written.

3. Definitions and main results #

JVP and VJP #

Let $f : \mathbb{R}^n \to \mathbb{R}^m$ be differentiable at $x$, with Jacobian $J = \partial f/\partial x \in \mathbb{R}^{m\times n}$.

• The JVP (Jacobian-vector product) maps a tangent $v \in \mathbb{R}^n$ to $Jv \in \mathbb{R}^m$.
• The VJP (vector-Jacobian product) maps a cotangent $u \in \mathbb{R}^m$ to $J^\top u \in \mathbb{R}^n$ (the row vector $u^\top J$).

Notation: for a scalar loss $L$, a bar denotes its gradient with respect to a variable: $\bar y = \partial L/\partial y$. If $y = f(x)$, the chain rule gives $\bar x = J^\top\bar y$: a VJP.

• Forward mode carries (value, tangent) pairs through the computation, in the same order. Each primitive operation has a JVP rule.
• Reverse mode runs the computation once and records the graph and its intermediate values (the tape). A backward pass then applies each primitive's VJP rule in reverse order, starting from $\bar L = 1$. Where a variable feeds several operations, the contributions add.

Theorem (VJP and JVP of a linear layer) #

Let $y = Wx$ with $W \in \mathbb{R}^{m\times n}$, $x \in \mathbb{R}^n$, and $\bar y = \partial L/\partial y$. Then

Proof.

1. Product rule: $dy = dW\,x + W\,dx$. This is the JVP with $(\dot x, \dot W) = (dx, dW)$.
2. The loss changes by $dL = \bar y^\top dy = \bar y^\top W\,dx + \bar y^\top dW\,x$.
3. First term: $\bar y^\top W\,dx = (W^\top\bar y)^\top dx$, so $\bar x = W^\top\bar y$.
4. Second term: $\bar y^\top dW\,x = \operatorname{tr}(\bar y^\top dW\,x) = \operatorname{tr}(x\bar y^\top dW) = \langle\bar yx^\top, dW\rangle_F$, with the Frobenius inner product $\langle P, Q\rangle_F = \operatorname{tr}(P^\top Q)$. So $\bar W = \bar yx^\top$. $\blacksquare$

Check every shape. The safest way to remember these formulas is to write the shape next to every symbol. A gradient always has the shape of its variable, and in each case only one product of the available matrices has that shape.

quantity	shape	formula	shape check
$x$, $\bar x$	$n\times1$	$\bar x = W^\top\bar y$	$(n\times m)(m\times1) = n\times1$
$W$, $\bar W$	$m\times n$	$\bar W = \bar yx^\top$	$(m\times1)(1\times n) = m\times n$
$X$, $\bar X$ (batch of $B$ columns)	$n\times B$	$\bar X = W^\top\bar Y$	$(n\times m)(m\times B) = n\times B$
$W$, $\bar W$ (batch)	$m\times n$	$\bar W = \bar YX^\top = \sum_b\bar y_bx_b^\top$	$(m\times B)(B\times n) = m\times n$

For a batch $Y = WX$ with examples as columns, the weight gradient $\bar W = \bar YX^\top$ is a sum of $B$ outer products, so its rank is at most the batch size $B$. For any product $C = AB$: $\bar A = \bar CB^\top$ and $\bar B = A^\top\bar C$. PyTorch's nn.Linear uses rows, $y = xW^\top + b$; the same rules hold, transposed.

Cost of training. The forward product $Wx$ costs $2mn$ flops per example; the backward pass needs $W^\top\bar y$ and $\bar yx^\top$, another $2mn$ each. So a training step costs about three forward passes, which gives the usual estimate "training flops $\approx 6 \times$ parameters $\times$ tokens".

Forward or reverse mode? #

The full Jacobian needs $n$ JVPs (one per column) or $m$ VJPs (one per row). Each pass costs a small constant times the cost of $f$: below 6 and typically 2–3 (Griewank and Walther, cited by Baydin et al.).

	forward mode (JVP)	reverse mode (VJP)
one pass gives	$Jv$, one input direction	$u^\top J$, one output weighting
full Jacobian	$n$ passes	$m$ passes
extra memory	a tangent per value, no tape	all intermediates of the forward pass
wins when	few inputs, many outputs	few outputs, many inputs

Training has $m = 1$ (a scalar loss) and $n$ up to $10^{12}$ parameters, so one reverse pass gives the whole gradient. The price is memory: reverse mode stores the activations. Checkpointing stores only some of them and recomputes the rest in the backward pass. A Hessian-vector product is the JVP of the gradient (forward over reverse): it costs a few gradient evaluations and never forms the Hessian, which is what Lanczos needs (lesson 14).

LoRA #

LoRA (low-rank adaptation; Hu et al., 2021) freezes a pretrained $W_0 \in \mathbb{R}^{d\times k}$ and learns a low-rank update:

In two plain sentences. $B$ starts at zero, so $BA = 0$ and the model at step 0 is exactly the pretrained model; fine-tuning starts from what the model already knows. The fixed factor $\alpha/r$ scales the update so that, with Adam, the first updates have about the same size for every rank $r$, and you do not need to tune the learning rate again when you change $r$.

The details:

• Parameters: $r(d + k)$ instead of $dk$. Gradients and Adam states exist only for these.
• Rank: $\operatorname{rank}(BA) \le r$ by the outer-product view. Both $A$ and $B$ are trained, so the column space of $BA$ moves during training. Only the rank bound is fixed; LoRA does not fix one low-rank subspace.
• Initialization: $A$ random Gaussian, $B = 0$, so $BA = 0$ and training starts exactly at $W_0$. Let $G = \partial L/\partial W$ be the gradient with respect to the effective weight $W = W_0 + \tfrac{\alpha}{r}BA$. Then $dW = \tfrac{\alpha}{r}(dB\,A + B\,dA)$ and $dL = \langle G, dW\rangle_F$ give

At the start $\bar A = 0$ but $\bar B \ne 0$, so $B$ moves first. If both were zero, both gradients would stay zero forever.

• Scaling $\alpha/r$ keeps the size of the early updates, and so the best learning rate, roughly independent of $r$. Hu et al. set $\alpha$ once and do not tune it.
• Inference: merge $W = W_0 + \tfrac{\alpha}{r}BA$: no extra latency. One frozen base can serve many adapters, and it can be stored quantized (QLoRA).

Why low-rank updates are plausible #

1. Structure of the gradient. By the theorem above, one example gives a rank-1 gradient $\bar yx^\top$, and a batch of $T$ tokens a gradient of rank at most $T$. In large batches $T$ exceeds $d$, so this bound alone proves little; but the singular values of gradients and of fine-tuning updates decay fast in practice.
2. Intrinsic dimension. Fine-tuning RoBERTa inside a random 200-dimensional subspace of parameter space reached 90% of full fine-tuning performance on MRPC (Aghajanyan et al., 2020).
3. The LoRA paper (§7). Very small ranks work ($r = 1$ for GPT-3 when adapting $W_q$ and $W_v$); the update amplifies directions that $W_0$ already contains but does not emphasize.
4. LoRA Without Regret (Thinking Machines, 2025). LoRA applied to all layers, especially the MLPs, matches full fine-tuning on small and medium supervised datasets, and in reinforcement learning even at rank 1 (policy-gradient RL gives only $O(1)$ bits of information per episode). It falls behind when the dataset exceeds the adapter's capacity. Its best learning rate is about 10 times that of full fine-tuning.

These are empirical findings, not theorems.

GaLore: full-rank training with a low-rank optimizer state

GaLore (Zhao et al., 2024) trains the full $W \in \mathbb{R}^{m\times n}$ ($m \le n$) but keeps Adam's state in a low-rank space. Every $T$ steps ($T = 200$ in the paper) it computes the SVD of the current gradient $G$ and sets $P \in \mathbb{R}^{m\times r}$ to the top $r$ left singular vectors. At every step:

$$R = P^\top G \in \mathbb{R}^{r\times n}, \qquad N = \text{Adam step computed on } R, \qquad W \leftarrow W - \eta\,\alpha\,PN.$$

Adam's two moments have size $r \times n$ instead of $m \times n$. Each update has rank $\le r$, but $P$ changes, so the sum of updates can have full rank.

	LoRA	GaLore
what is low rank	the update $\Delta W = BA$	the gradient, for the optimizer state
trainable parameters	$r(d + k)$	all $mn$
saves	gradients, optimizer state, per-task storage	optimizer state
total update	rank $\le r$	can be full rank
typical use	fine-tuning, many adapters	memory-limited pretraining or full fine-tuning

Tensors: fibers, unfoldings, CP and Tucker #

A third-order tensor $\mathcal X \in \mathbb{R}^{I\times J\times K}$ is an array with entries $x_{ijk}$. The outer product $a \circ b \circ c$ has entries $a_ib_jc_k$ (a rank-1 tensor).

• A mode-$n$ fiber fixes all indices but the $n$-th (for $n = 1$: the vector $x_{:jk}$).
• The mode-$n$ unfolding $X_{(n)}$ is the matrix whose columns are the mode-$n$ fibers: $I \times JK$ for $n = 1$. We order the columns with the first remaining index varying fastest (Kolda and Bader). The column order is a convention; it changes neither the rank nor the left singular vectors.
• The mode-$n$ product $\mathcal X \times_n U$ multiplies every mode-$n$ fiber by $U$: $(\mathcal X \times_1 U)_{(1)} = UX_{(1)}$.

One small tensor, step by step. Let $\mathcal X \in \mathbb{R}^{2\times2\times2}$ have the two frontal slices ($k = 1$ and $k = 2$)

1. Mode-1 fibers (index $i$ varies): the columns of each slice. In the order $(j, k) = (1,1), (2,1), (1,2), (2,2)$ they are $(1, 3)$, $(2, 4)$, $(5, 7)$, $(6, 8)$.
2. Mode-2 fibers (index $j$ varies): the rows of each slice. In the order $(i, k) = (1,1), (2,1), (1,2), (2,2)$ they are $(1, 2)$, $(3, 4)$, $(5, 6)$, $(7, 8)$.
3. Mode-3 fibers (index $k$ varies): from the front slice to the back slice. In the order $(i, j) = (1,1), (2,1), (1,2), (2,2)$ they are $(1, 5)$, $(3, 7)$, $(2, 6)$, $(4, 8)$.
4. Put the fibers side by side as columns:

5. A mode product. With $U = (1 \;\; 1)$ ($1\times2$), $\mathcal X \times_1 U$ adds the two entries of every mode-1 fiber: $UX_{(1)} = (4 \;\; 6 \;\; 12 \;\; 14)$. The result is a $1\times2\times2$ tensor with slices $(4 \;\; 6)$ and $(12 \;\; 14)$.

Interactive figure: loading…

CP (CANDECOMP/PARAFAC) of rank $R$:

with factor matrices $A$ ($I\times R$), $B$ ($J\times R$), $C$ ($K\times R$): $R(I + J + K)$ parameters (plus $R$ weights $\lambda_r$ if the columns are normalized). The CP rank is the smallest $R$ with equality.

Tucker with core size $R_1\times R_2\times R_3$:

with $R_1R_2R_3 + IR_1 + JR_2 + KR_3$ parameters. CP is Tucker with a superdiagonal core ($g_{rrr} = \lambda_r$, all other entries 0). For a matrix ($N = 2$ indices), Tucker is $UGV^\top$ and the SVD is the case with a diagonal core. The multilinear rank is $(\operatorname{rank}X_{(1)}, \operatorname{rank}X_{(2)}, \operatorname{rank}X_{(3)})$; unlike row and column rank of a matrix, its entries can differ.

Truncated HOSVD. For each mode, $U_n$ = the top $R_n$ left singular vectors of $X_{(n)}$. Then

with the orthogonal projectors $P_n = U_nU_n^\top$. It is the analogue of $U_kU_k^\top AV_kV_k^\top$ for a matrix. The core is not diagonal in general: a tensor cannot be diagonalized by one orthogonal change of basis per mode.

Theorem (truncated HOSVD is quasi-optimal). For an $N$-th order tensor, let $\mathcal X_{\text{best}}$ be the best approximation with multilinear rank $\le (R_1, \dots, R_N)$. Then

Proof for $N = 3$ (telescoping and Eckart–Young per mode)

1. Telescope: $\mathcal X - \hat{\mathcal X} = \mathcal X \times_1 (I - P_1) + \mathcal X \times_1 P_1 \times_2 (I - P_2) + \mathcal X \times_1 P_1 \times_2 P_2 \times_3 (I - P_3)$.
2. The three terms are orthogonal: the first has mode-1 fibers in the range of $I - P_1$, the other two in the range of $P_1$; the second and third differ in the same way in mode 2. So the squared norms add.
3. Mode products in different modes commute, and an orthogonal projection does not increase the norm. So, for example, $\|\mathcal X \times_1 P_1 \times_2 (I - P_2)\| \le \|\mathcal X \times_2 (I - P_2)\|$.
4. $\|\mathcal X \times_n (I - P_n)\|_F^2 = \|(I - P_n)X_{(n)}\|_F^2 = \sum_{i > R_n}\sigma_i(X_{(n)})^2$, because $P_n$ projects on the top $R_n$ left singular vectors (Eckart–Young, lesson 08). This gives the first inequality.
5. Any $\mathcal Y$ with multilinear rank $\le (R_n)$ has $\operatorname{rank}Y_{(n)} \le R_n$, so by Eckart–Young $\|\mathcal X - \mathcal Y\|_F^2 = \|X_{(n)} - Y_{(n)}\|_F^2 \ge \sum_{i > R_n}\sigma_i(X_{(n)})^2$ for every $n$. Take $\mathcal Y = \mathcal X_{\text{best}}$ and add over $n$. $\blacksquare$

So the truncated HOSVD is within a factor $\sqrt N$ of the best Tucker model, but in general not optimal: there is no Eckart–Young theorem for tensors. HOOI (alternating least squares on the factors, started from the HOSVD) usually improves it.

CP rank is harder than matrix rank. It can depend on the field (a real $2\times2\times2$ tensor can have rank 3 over $\mathbb{R}$ and 2 over $\mathbb{C}$), and computing it is NP-hard. A best rank-$R$ approximation may not exist (rank-2 tensors can converge to a rank-3 tensor: degeneracy), and the best rank-1 term need not be part of the best rank-2 model, so greedy deflation fails. But CP is often essentially unique (Kruskal's condition), so its factors can be interpreted; a matrix factorization $BA$ is not unique.

Attention in matrix form #

We use the column convention of Elhage et al. (2021). Token $i$ is the residual-stream vector $x_i \in \mathbb{R}^d$, and $X = [x_1 \cdots x_n] \in \mathbb{R}^{d\times n}$. One head of width $d_h$ has $W_Q, W_K, W_V \in \mathbb{R}^{d_h\times d}$ and $W_O \in \mathbb{R}^{d\times d_h}$. With $q_i = W_Qx_i$, $k_j = W_Kx_j$, $v_j = W_Vx_j$:

A causal mask sets $s_{ij} = -\infty$ for $j > i$. The layer adds $H$ to the residual stream: $X \leftarrow X + \sum_{\text{heads}} H$ (then an MLP adds its output in the same way).

• $W_{QK}$ (the QK circuit) is a $d \times d$ matrix of rank $\le d_h$. The score is a bilinear form: the query side reads the subspace $C(W_Q^\top)$ of the residual stream, the key side reads $C(W_K^\top)$.
• $W_{OV}$ (the OV circuit) is a $d \times d$ map of rank $\le d_h$. In its SVD $W_{OV} = U\Sigma V^\top$, the right singular vectors say what is read from the source token and the left singular vectors say where it is written in the destination token.
• Only these products matter. For any invertible $M, M' \in \mathbb{R}^{d_h\times d_h}$, replacing $W_Q \to MW_Q$, $W_K \to M^{-\top}W_K$, $W_V \to M'W_V$, $W_O \to W_OM'^{-1}$ changes no output. Queries, keys and values are intermediate results of two low-rank matrices.
• Linear once $A$ is fixed. In vectorized form $\operatorname{vec}(H) = (A \otimes W_{OV})\operatorname{vec}(X)$ (lesson 16 proves the rule $\operatorname{vec}(AXB) = (B^\top \otimes A)\operatorname{vec}X$): $A$ mixes tokens, $W_{OV}$ acts on each token's vector. The softmax inside $A$ is the only nonlinear part.
• Heads are independent and additive. "Concatenate the head results $r^h = \sum_j A^h_{ij}v^h_j$, then multiply by $W_O$" equals $\sum_h W_O^hr^h$, where $W_O^h$ is the block of columns of $W_O$ for head $h$. Every head reads from the residual stream with linear maps and writes back into a subspace of dimension at most $d_h$ (usually $d_h = d/H$ for $H$ heads).
• Low rank is in the logits, not in $A$. $S = X^\top W_{QK}X/\sqrt{d_h}$ ($n\times n$) has rank $\le d_h$, but the entrywise exponential makes $A$ full rank in general.
• Why $\sqrt{d_h}$. If $q$ and $k$ have independent entries with mean 0 and variance 1, then $q^\top k$ has variance $d_h$. Dividing by $\sqrt{d_h}$ keeps the logits of order 1, so the softmax does not saturate.

Interactive figure: loading…

Code uses rows. With $X_r = X^\top \in \mathbb{R}^{n\times d}$ and $Q = X_rW_Q^\top$ (what nn.Linear computes), the logits are $QK^\top/\sqrt{d_h} = X_rW_Q^\top W_KX_r^\top/\sqrt{d_h}$: the same $W_{QK}$. The output rows are $AX_rW_V^\top W_O^\top = AX_rW_{OV}^\top$: the map is the same, written as a right product with the transpose.

Answers to section 1. (1) $L$: reverse mode, one pass ($m = 1$). $g$: forward mode, 3 passes ($n = 3$), instead of $10^6$ reverse passes. (2) $\operatorname{rank}W_{QK} \le 128$. The logit matrix also has rank $\le 128$, but $A$ itself can have full rank 1000.

4. Worked example #

A. One linear layer, both modes. $W = \begin{pmatrix} 2 & 1 \\ 0 & 1 \\ 1 & 3\end{pmatrix}$ ($3\times2$), $x = (1, 2)$, upstream gradient $\bar y = (1, -1, 1)$.

• Forward: $y = Wx = 1\cdot(2, 0, 1) + 2\cdot(1, 1, 3) = (4, 2, 7)$.
• VJP for $x$: $\bar x = W^\top\bar y = (2\cdot1 + 0\cdot(-1) + 1\cdot1,\ 1\cdot1 + 1\cdot(-1) + 3\cdot1) = (3, 3)$. Shape: $(2\times3)(3\times1) = 2\times1$, like $x$.
• VJP for $W$: $\bar W = \bar yx^\top = \begin{pmatrix} 1 & 2 \\ -1 & -2 \\ 1 & 2\end{pmatrix}$, shape $3\times2$ like $W$, rank 1 (every row is a multiple of $x^\top = (1, 2)$).
• Check: for $L = \bar y^\top Wx$, increasing $x_1$ by $10^{-6}$ increases $L$ by $3\times10^{-6} = \bar x_1\cdot10^{-6}$.
• JVP with $\dot x = e_1$, $\dot W = 0$: $\dot y = We_1 = (2, 0, 1)$, the first column of $J = W$. The full $3\times2$ Jacobian takes 2 JVPs or 3 VJPs.

B. LoRA and GaLore memory. A $4096 \times 4096$ attention projection with $r = 16$: LoRA trains $16 \cdot 8192 = 131{,}072$ numbers instead of $16{,}777{,}216$ (0.78%), and Adam keeps $262{,}144$ moment entries instead of $33{,}554{,}432$. GaLore with $r = 128$ on a $4096 \times 11008$ MLP matrix trains all $45{,}088{,}768$ weights, but its Adam moments have $2 \cdot 128 \cdot 11008 = 2{,}818{,}048$ entries, plus $4096 \cdot 128 = 524{,}288$ for $P$: 3.7% of the $90{,}177{,}536$ entries of full Adam.

C. A tensor. For the $2\times2\times2$ tensor of section 3, each of the three unfoldings has rank 2, so the multilinear rank is $(2, 2, 2)$. For counts, take a climate data set with 1000 grid points × 365 days × 10 variables ($3{,}650{,}000$ numbers). CP with $R = 20$ needs $20 \cdot (1000 + 365 + 10) = 20\cdot1375 = 27{,}500$ parameters; Tucker with core $20\times10\times5$ needs $20\cdot10\cdot5 + 1000\cdot20 + 365\cdot10 + 10\cdot5 = 1000 + 20{,}000 + 3{,}650 + 50 = 24{,}700$.

D. One attention head by hand. Take $d = 3$, $d_h = 1$, two tokens $x_1 = e_1$ ("has feature 1") and $x_2 = e_2$ ("looks for feature 1"), and a causal mask. Let $W_Q = e_2^\top$, $W_K = 2e_1^\top$, $W_V = e_1^\top$, $W_O = e_3$. Then

both of rank $1 = d_h$. The logits are $s_{11} = x_1^\top W_{QK}x_1 = 0$, $s_{21} = 2$, $s_{22} = 0$ ($s_{12}$ is masked), so

$W_{OV}x_1 = e_3$ and $W_{OV}x_2 = 0$, so $h_1 = e_3$ and $h_2 = 0.881\,e_3$. Token 2 found the token with feature 1 and copied it into coordinate 3 of its own residual vector. (Token 1 can only attend to itself, so it copies its own feature.) The logit matrix $\begin{pmatrix} 0 & 0 \\ 2 & 0\end{pmatrix}$ (before the mask) has rank 1, while $\det A = 0.119 \ne 0$: $A$ has rank 2.

5. In code #

The worked example A, both modes, and a Hessian-vector product, with NumPy and hand-written JVP and VJP rules. An autodiff library (torch.func.vjp, jax.vjp, jax.jvp) returns the same numbers.

import numpy as np

# Worked example A: y = W x, with the upstream gradient ybar
W = np.array([[2., 1.], [0., 1.], [1., 3.]])
x = np.array([1., 2.])
ybar = np.array([1., -1., 1.])

y = W @ x
xbar = W.T @ ybar                # VJP with respect to x
Wbar = np.outer(ybar, x)         # VJP with respect to W: rank 1, shape of W
print(y, xbar)                   # [4. 2. 7.] [3. 3.]
print(Wbar.tolist(), np.linalg.matrix_rank(Wbar))   # [[1.0, 2.0], [-1.0, -2.0], [1.0, 2.0]] 1

# finite-difference check of xbar for L = ybar^T W x
L = lambda W, x: ybar @ (W @ x)
h = 1e-6
print(round((L(W, x + h * np.array([1., 0.])) - L(W, x)) / h, 6))   # 3.0 = xbar[0]

# forward mode: one JVP with dx = e1, dW = 0 gives the first column of J = W
print(W @ np.array([1., 0.]) + np.zeros_like(W) @ x)   # [2. 0. 1.]

# a small nonlinear map f: R^2 -> R^3 and its Jacobian, by columns (JVPs) or rows (VJPs)
W1 = np.array([[1., -1.], [0.5, 2.]]); W2 = W
def f(x): return W2 @ np.tanh(W1 @ x)
def jvp(x, v):                   # forward mode: push v through the chain
    a = W1 @ x
    return W2 @ ((1 - np.tanh(a) ** 2) * (W1 @ v))
def vjp(x, u):                   # reverse mode: pull u back through the chain
    a = W1 @ x
    return W1.T @ ((1 - np.tanh(a) ** 2) * (W2.T @ u))
J_fwd = np.column_stack([jvp(x, e) for e in np.eye(2)])   # n = 2 JVPs
J_rev = np.vstack([vjp(x, e) for e in np.eye(3)])         # m = 3 VJPs
print(np.allclose(J_fwd, J_rev))                          # True

# Hessian-vector product without forming H: the JVP of the gradient.
# An AD library does it exactly (forward over reverse); here a central difference of the gradient.
S = np.array([[3., 1.], [1., 2.]])
grad = lambda z: S @ z + np.array([np.cos(z[0]), 0.])     # gradient of 0.5 z^T S z + sin(z0)
v = np.array([1., -1.])
eps = 1e-5
Hv = (grad(x + eps * v) - grad(x - eps * v)) / (2 * eps)
H = S + np.diag([-np.sin(x[0]), 0.])
print(np.allclose(Hv, H @ v, atol=1e-8))                  # True

The cost of the two multiplication orders for the chain of section 2:

import numpy as np

def chain_cost(shapes, from_left):
    """Multiplications to multiply a chain of matrices (list of (rows, cols)) from one end."""
    shapes = list(shapes)
    total = 0
    while len(shapes) > 1:
        if from_left:                      # reverse mode: start at the output
            (a, b), (_, c) = shapes[0], shapes[1]
            shapes[:2] = [(a, c)]
        else:                              # forward mode: start at the input
            (a, b), (_, c) = shapes[-2], shapes[-1]
            shapes[-2:] = [(a, c)]
        total += a * b * c
    return total

chain = [(1, 1000), (1000, 1000), (1000, 1000)]            # J3, J2, J1
print(chain_cost(chain, from_left=True))                   # 2000000
print(chain_cost(chain, from_left=False))                  # 1001000000

LoRA gradients at initialization, and why GaLore's total update is not low rank:

import numpy as np

rng = np.random.default_rng(0)
d, k, r, alpha, n_tok = 64, 48, 4, 8.0, 32
W0 = rng.standard_normal((d, k))              # frozen
A = rng.standard_normal((r, k)) / np.sqrt(k)  # random init
B = np.zeros((d, r))                          # zero init
X = rng.standard_normal((k, n_tok))           # 32 input tokens as columns
T = rng.standard_normal((d, n_tok))           # targets

# loss 0.5 ||W X - T||_F^2 with W = W0 + (alpha/r) B A; at init W = W0
G = (W0 @ X - T) @ X.T                        # dL/dW: a sum of 32 outer products
Bbar = (alpha / r) * G @ A.T                  # dL/dB
Abar = (alpha / r) * B.T @ G                  # dL/dA
print(np.abs(Abar).max(), np.abs(Bbar).max() > 0)   # 0.0 True: only B moves at the first step
print(np.linalg.matrix_rank(G))               # 32: rank <= number of tokens (here < min(d, k) = 48)
print(r * (d + k), d * k)                     # 448 trainable numbers instead of 3072

# check dL/dB with a finite difference in one entry
def loss(B):
    W = W0 + (alpha / r) * B @ A
    return 0.5 * np.sum((W @ X - T) ** 2)
E = np.zeros_like(B); E[3, 1] = 1e-6
print(np.isclose((loss(B + E) - loss(B - E)) / 2e-6, Bbar[3, 1]))   # True

# GaLore-style: each update has rank <= r, but the subspace changes, so the sum does not
updates = []
for t in range(4):
    Gt = rng.standard_normal((d, k))                        # stand-in for the gradient at refresh step t
    P = np.linalg.svd(Gt, full_matrices=False)[0][:, :r]    # top-r left singular vectors
    updates.append(P @ (P.T @ Gt))                          # project, (optimizer step), project back
print([int(np.linalg.matrix_rank(u)) for u in updates])     # [4, 4, 4, 4]
print(np.linalg.matrix_rank(sum(updates)))                  # 16: the total update is not rank 4

CP, unfoldings and the truncated HOSVD (no library needed beyond the SVD):

import numpy as np

rng = np.random.default_rng(0)
I, J, K = 30, 20, 10

def unfold(X, n):
    """Mode-n unfolding: mode-n fibers as columns, first remaining index fastest."""
    return np.moveaxis(X, n, 0).reshape(X.shape[n], -1, order="F")

# the 2x2x2 example of section 3 (x[i, j, k], slices along the last index)
X2 = np.zeros((2, 2, 2)); X2[:, :, 0] = [[1, 2], [3, 4]]; X2[:, :, 1] = [[5, 6], [7, 8]]
print([unfold(X2, n).astype(int).tolist() for n in range(3)])
# [[[1, 2, 5, 6], [3, 4, 7, 8]], [[1, 3, 5, 7], [2, 4, 6, 8]], [[1, 3, 2, 4], [5, 7, 6, 8]]]

# A CP tensor of rank 3: x_ijk = sum_r a_ir b_jr c_kr
A, B, C = rng.standard_normal((I, 3)), rng.standard_normal((J, 3)), rng.standard_normal((K, 3))
X_cp = np.einsum("ir,jr,kr->ijk", A, B, C)
print([int(np.linalg.matrix_rank(unfold(X_cp, n))) for n in range(3)])   # [3, 3, 3]
print(3 * (I + J + K), I * J * K)                                         # 180 parameters vs 6000 entries

# Truncated HOSVD of a noisy tensor with multilinear rank (3, 3, 3)
X = X_cp + 0.1 * rng.standard_normal((I, J, K))
ranks = (3, 3, 3)
U = [np.linalg.svd(unfold(X, n), full_matrices=False)[0][:, :ranks[n]] for n in range(3)]
G = np.einsum("ijk,ip,jq,ks->pqs", X, *U)                 # core G = X x1 U1ᵀ x2 U2ᵀ x3 U3ᵀ
X_hat = np.einsum("pqs,ip,jq,ks->ijk", G, *U)             # X_hat = G x1 U1 x2 U2 x3 U3
err = np.linalg.norm(X - X_hat)
tails = [np.sum(np.linalg.svd(unfold(X, n), compute_uv=False)[ranks[n]:] ** 2) for n in range(3)]
lower = np.sqrt(max(tails))     # Eckart-Young per mode: no Tucker model of this size does better
print(G.shape, round(lower, 3), round(err, 3), round(np.sqrt(sum(tails)), 3), round(np.sqrt(3) * lower, 3))
# (3, 3, 3) 7.263 7.606 12.04 12.58: lower <= err <= sqrt(sum of tails) <= sqrt(3) * lower

Attention: multi-head attention computed the way a library does it (rows, heads concatenated, then $W_O$), checked against $W_{QK}$, $W_{OV}$ and the sum over heads:

import numpy as np

rng = np.random.default_rng(0)
d, H, n = 16, 4, 6                     # model width, heads, tokens
dh = d // H
Wq, Wk, Wv, Wo = (rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(4))
Xr = rng.standard_normal((n, d))       # library layout: tokens as rows

def softmax_rows(S):
    E = np.exp(S - S.max(axis=1, keepdims=True))
    return E / E.sum(axis=1, keepdims=True)

# how a library computes multi-head attention: rows, split into heads, concatenate, then W_O
Qr, Kr, Vr = Xr @ Wq.T, Xr @ Wk.T, Xr @ Wv.T
heads = []
for h in range(H):
    c = slice(h * dh, (h + 1) * dh)
    A_h = softmax_rows(Qr[:, c] @ Kr[:, c].T / np.sqrt(dh))
    heads.append(A_h @ Vr[:, c])
out = np.concatenate(heads, axis=1) @ Wo.T            # n x d

# the same layer in the column convention: a sum over heads of W_OV X Aᵀ
X = Xr.T                                               # d x n, column i = token i
total = np.zeros((d, n))
for h in range(H):
    rows = slice(h * dh, (h + 1) * dh)
    W_QK = Wq[rows].T @ Wk[rows]                       # d x d, rank <= dh
    W_OV = Wo[:, rows] @ Wv[rows]                      # d x d, rank <= dh
    S = X.T @ W_QK @ X / np.sqrt(dh)                   # logits s_ij = x_iᵀ W_QK x_j / sqrt(dh)
    A = softmax_rows(S)
    total += W_OV @ X @ A.T                            # head output, d x n
    if h == 0:
        rk = np.linalg.matrix_rank
        print(rk(W_QK), rk(W_OV), rk(S), rk(A))        # 4 4 4 6
        vec = lambda M: M.reshape(-1, order="F")       # vec stacks columns
        print(np.allclose(np.kron(A, W_OV) @ vec(X), vec(W_OV @ X @ A.T)))   # True
        M = rng.standard_normal((dh, dh))              # reparametrize Q and K
        print(np.allclose((M @ Wq[rows]).T @ (np.linalg.inv(M).T @ Wk[rows]), W_QK))  # True
print(np.allclose(total, out.T))      # True: multi-head = sum of the heads' W_OV X Aᵀ

6. Where it shows up #

• Autodiff in science. Variational data assimilation (4D-Var) in weather prediction runs a tangent-linear model (JVPs) and an adjoint model (VJPs): reverse mode, written by hand long before deep learning. Adjoint methods give gradients of a PDE solution with respect to many parameters at the cost of one extra solve. Forward mode fits the sensitivities of a simulator to a few physical parameters.
• Fine-tuning foundation models. Adapting a pretrained Earth-system model to a new region or variable with LoRA keeps one frozen backbone and one small adapter per task.
• Tensors in ML for science. Ensemble forecasts (member × variable × grid point × time) are compressed with Tucker/HOSVD, with an error you can choose. The TFNO model of the neuraloperator library stores the spectral weights of a Fourier neural operator in Tucker form. CP gives interpretable components because it is often essentially unique.
• Attention in practice. The KV cache stores the keys and values of every past token; multi-query and grouped-query attention share them across heads, and DeepSeek's multi-head latent attention (2024) caches a low-rank latent vector instead. Interpretability work reads heads through the SVDs of $W_{OV}$ and $W_{QK}$.

7. Common mistakes #

• "Backprop computes the Jacobian." One reverse pass computes one VJP, $u^\top J$. The full Jacobian needs $m$ passes.
• Wrong shapes in the VJP. $\bar W = \bar yx^\top$ has the shape of $W$ ($m\times n$); $x\bar y^\top$ is its transpose. In row convention ($Y = XW^\top$): $\bar W = \bar Y^\top X$ and $\bar X = \bar YW$.
• Initializing both LoRA factors to zero. Both gradients vanish and nothing trains.
• "GaLore reduces the trainable parameters." It reduces optimizer memory; all weights train.
• Treating the truncated HOSVD as optimal, or truncating a CP model and expecting the best lower-rank CP model. Neither holds for tensors. (Also, CP rank is one number; multilinear rank is one number per mode.)
• "The attention matrix is low rank." The logits are (rank $\le d_h$); $A$ after the softmax is usually not.
• Mixing conventions. $W_Q^\top W_K$ (weights act on columns) and $W_QW_K^\top$ (row-convention weights of shape $d\times d_h$) are the same $W_{QK}$ written two ways; check shapes before you transpose.
• Row-major unfolding. NumPy's default reshape is row-major. With order="C" the columns of $X_{(n)}$ come in another order; the rank and the left singular vectors stay the same, but the matrix differs from the Kolda–Bader convention used in papers.

8. Check yourself #

1. Give the VJP of $C = AB$ ($A \in \mathbb{R}^{m\times p}$, $B \in \mathbb{R}^{p\times n}$) and check the shapes.
   Answer

   $\bar A = \bar CB^\top$ ($m\times n$ times $n\times p$ = $m\times p$) and $\bar B = A^\top\bar C$ ($p\times m$ times $m\times n$ = $p\times n$). From $dL = \langle\bar C, dA\,B + A\,dB\rangle_F$.
2. How do you get $Hv$ for the Hessian $H$ of a scalar loss without forming $H$? What does it cost?
   Answer

   Take the JVP of the gradient function in direction $v$ (forward over reverse): $Hv = \tfrac{d}{dt}\nabla L(x + tv)|_{t=0}$. It costs a small constant times one gradient evaluation, with memory like one backward pass.
3. LoRA with $r = 32$ on a $5120\times5120$ matrix: how many trainable parameters, and what fraction of the matrix?
   Answer

   $32\cdot(5120 + 5120) = 327{,}680$ of $26{,}214{,}400$: 1.25%.
4. Why does LoRA set $B = 0$ and $A$ random, and not $A = B = 0$?
   Answer

   $BA = 0$ makes training start at $W_0$. The gradients are $\bar B = \tfrac{\alpha}{r}GA^\top$ and $\bar A = \tfrac{\alpha}{r}B^\top G$. With random $A$, $\bar B \ne 0$; with $A = B = 0$ both gradients are zero and stay zero.
5. Show that CP is a special case of Tucker, and compare the parameter counts for $I = J = K = 50$, CP rank 5, Tucker core $5\times5\times5$, and the full tensor.
   Answer

   A superdiagonal core ($g_{rrr} = \lambda_r$, other entries 0) turns $\sum_{p,q,s}g_{pqs}a_{ip}b_{jq}c_{ks}$ into $\sum_r\lambda_ra_{ir}b_{jr}c_{kr}$. Counts: CP $5\cdot150 = 750$, Tucker $125 + 750 = 875$, full $125{,}000$.
6. Show that $W_{QK}$ does not change when $W_Q \to MW_Q$ and $W_K \to M^{-\top}W_K$ for an invertible $M \in \mathbb{R}^{d_h\times d_h}$. What is the rank of the logit matrix for $n = 2048$ tokens and $d_h = 64$?
   Answer

   $(MW_Q)^\top(M^{-\top}W_K) = W_Q^\top M^\top M^{-\top}W_K = W_Q^\top W_K$. The logit matrix $X^\top W_{QK}X/\sqrt{d_h}$ has rank at most 64.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What do a JVP and a VJP compute for $f: \mathbb{R}^n \to \mathbb{R}^m$ with Jacobian $J$?

   Show answer

2. Why does training use reverse mode?

   Show answer

3. What does reverse mode cost in memory, and how do you reduce it?

   Show answer

4. Give $\bar x$ and $\bar W$ for $y = Wx$.

   Show answer

5. How many trainable numbers does LoRA with rank $r$ add to a $d\times k$ matrix?

   Show answer

6. Why does LoRA initialize $B = 0$ and $A$ random?

   Show answer

7. What does GaLore make low rank, and what does it save?

   Show answer

8. Define the mode-$n$ unfolding of a tensor.

   Show answer

9. Write a rank-$R$ CP model and its parameter count for an $I\times J\times K$ tensor.

   Show answer

10. What are $W_{QK}$ and $W_{OV}$ of an attention head, and what is their largest possible rank?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why does deep learning use reverse mode?" — the loss is a scalar, so one VJP gives all gradients at a few times the forward cost; forward mode would need one pass per parameter. The cost is memory for the activations.
• "Estimate the training compute of a 7B model on $10^{12}$ tokens." — $6 \times 7\times10^9 \times 10^{12} = 4.2\times10^{22}$ flops (2 per parameter per token forward, 4 backward).
• "LoRA or GaLore for pretraining under a memory limit?" — GaLore: it trains all weights and the total update can have full rank; LoRA's update is limited to rank $r$.
• "What are the QK and OV circuits?" — $W_Q^\top W_K$ decides where a token looks, $W_OW_V$ decides what is moved and where it is written; both have rank $\le d_h$; with $A$ fixed the head is linear, $(A \otimes W_{OV})$.
• "CP or Tucker for a 4-D climate data set?" — Tucker/HOSVD for compression (orthonormal factors from SVDs, error controlled, within $\sqrt N$ of optimal); CP for interpretable components (unique), with harder fitting and no guarantee that a best model exists.

10. Practice #

The notebook for this lesson runs in Colab with NumPy only. It has five exercises, each with a check cell and a folded solution:

1. Write the VJP of a matrix product $C = AB$ and check it against finite differences.
2. Build the Jacobian of a small two-layer map from JVPs (columns) and from VJPs (rows), and count the passes.
3. Compute the LoRA gradients $\bar B$, $\bar A$ from $G$, and see that only $B$ moves at the first step.
4. Write the mode-$n$ unfolding and the truncated HOSVD, and check the quasi-optimality bounds.
5. Write one attention head in the column convention, reproduce worked example D, and check that multi-head attention is a sum of heads.

11. Watch and read #

What is Automatic Differentiation?Ari Seff · 14 min

Lecture 8 Part 2: Automatic Differentiation on Computational GraphsAlan Edelman, Steven G. Johnson · MIT 18.S096 Matrix Calculus (MIT OpenCourseWare) · 65 min

What is Low-Rank Adaptation (LoRA) | explained by the inventorEdward Hu · 7 min

Attention in transformers, step-by-step | Deep Learning Chapter 63Blue1Brown · Deep learning, chapter 6 · 26 min

• Baydin, Pearlmutter, Radul, Siskind, Automatic Differentiation in Machine Learning: a Survey (JMLR 2018), §3.1–3.2 (forward mode, reverse mode, their costs).
• Hu et al., LoRA: Low-Rank Adaptation of Large Language Models (2021), §4.1 (the parameterization and initialization) and §7 (understanding the low-rank updates). Follow-ups: GaLore (Zhao et al., 2024, §3–4) and LoRA Without Regret (Thinking Machines, 2025).
• Kolda and Bader, Tensor Decompositions and Applications (SIAM Review 2009), §2.4–2.5 (unfoldings, mode-$n$ product), §3.1–3.3 (CP rank, degeneracy), §4.1–4.2 (Tucker, HOSVD, HOOI).
• Elhage et al., A Mathematical Framework for Transformer Circuits (Anthropic, 2021), sections "Transformer Overview" (residual stream, heads independent and additive, information movement) and "One-Layer Attention-Only Transformers" (QK and OV circuits).