Matrix differentials, matrix functions, Kronecker product and vec

Before this lesson: lesson 03 (inverse, determinant, trace), lesson 05 ($A = X\Lambda X^{-1}$ and $e^{Jt}$), lesson 06 ($S = Q\Lambda Q^\top$) and lesson 15 (JVP and VJP). Time: about 85 minutes, plus the notebook.

---

Why this matters #

Many ML objectives are scalar functions of a matrix: a Gaussian log-likelihood contains $\log\det\Sigma$ and $\Sigma^{-1}$, and cross-entropy and attention go through a softmax. Index-by-index derivatives are slow and error-prone; the differential method gives them in a few lines. Matrix functions say what $e^{A}$ or $\sqrt{S}$ means. The matrix exponential solves every linear ODE $x' = Ax$: heat diffusion on a grid, linearized dynamics, the state-space layers of Mamba. The Kronecker product and vec turn linear maps on matrices, such as $X \mapsto AXB$, into ordinary matrices: the algebra behind layer Jacobians, K-FAC and Shampoo, and fast Gaussian processes on space–time grids.

1. Predict first #

2. The picture #

Differentials. Start with one variable. The differential $df = f'(x)\,dx$ is the linear part of the change $f(x + dx) - f(x)$; the rest shrinks faster than $dx$. The matrix case is the same idea. For a function of a matrix $X$, $df$ is the part of $f(X + dX) - f(X)$ that is linear in the small matrix $dX$. We use rules (sum, product, chain) on whole matrices and never write indices. The one new thing: matrices do not commute, so the order of the factors must stay.

A scalar $f$ of an $m\times n$ matrix has a linear part $\sum_{ij} G_{ij}\, dX_{ij}$ for some $m\times n$ matrix $G$, and this sum is $\operatorname{tr}(G^\top dX)$. So if we bring $df$ into the form $\operatorname{tr}(G^\top dX)$, the gradient is $G$. The tool for this is the cyclic property of the trace.

How to test a derivative. Walk from $X$ along a direction $dX$ for a step $t$. A correct derivative predicts the new value up to a second-order term:

So the error of the prediction falls like $t^2$: halve $t$ and the error is divided by 4. A wrong derivative leaves an error of order $t$: halve $t$ and the error is only divided by 2. On a log–log plot the first is a line of slope 2 and the second a line of slope 1. This is the standard gradient check of autodiff libraries, done along one random direction.

Interactive figure: loading…

Matrix functions. Take $S = \operatorname{diag}(4, 1)$. The obvious candidate for $\sqrt{S}$ is $\operatorname{diag}(2, 1)$: it squares to $S$. A symmetric $S = Q\Lambda Q^\top$ is diagonal in its eigen-axes, so we apply $f$ to each axis on its own: $f(S)$ scales the eigen-axis $q_i$ by $f(\lambda_i)$. It is not the entrywise $f(S_{ij})$.

Kronecker and vec. The map $X \mapsto AXB$ is linear in $X$, so it is a matrix acting on the coordinates of $X$. Stack the columns of $X$ into one long vector $\operatorname{vec}(X)$. The matrix of the map is then $B^\top \otimes A$: a block matrix made of copies of $A$, each scaled by an entry of $B^\top$.

Interactive figure: loading…

3. Definitions and main results #

The differential method #

Definition. For a map $f$ from matrices to matrices, the differential $df$ is the linear function of $dX$ with $f(X + dX) = f(X) + df + o(\|dX\|)$.

Rules. Each follows from expanding $f(X + dX)$ and dropping second-order terms.

• constant: $dA = 0$; sum: $d(F + G) = dF + dG$;
• product: $d(FG) = dF\, G + F\, dG$ (keep the order);
• transpose and trace: $d(F^\top) = (dF)^\top$, $d \operatorname{tr} F = \operatorname{tr}(dF)$;
• chain rule: the differential of $g(F(X))$ is the differential of $g$ at $F$, applied to $dF$.

Identification rule. If a scalar $f$ has $df = \operatorname{tr}(G^\top dX)$ for every $dX$, then $\nabla_X f = G$, the matrix of partial derivatives $\partial f/\partial X_{ij}$ (same shape as $X$). Reason: $\operatorname{tr}(G^\top dX) = \sum_{ij} G_{ij}\,dX_{ij}$; put $dX = E_{ij}$, the matrix with a single 1 at $(i, j)$. For vectors, $df = g^\top dx$ gives $\nabla f = g$.

The rule entry by entry, for a $2\times2$ matrix.

1. Multiply out: $G^\top dX = \begin{pmatrix} g_{11} & g_{21} \\ g_{12} & g_{22}\end{pmatrix}\begin{pmatrix} dx_{11} & dx_{12} \\ dx_{21} & dx_{22}\end{pmatrix}$. Its diagonal entries are $g_{11}dx_{11} + g_{21}dx_{21}$ and $g_{12}dx_{12} + g_{22}dx_{22}$.
2. Add them: the trace is the sum of the four terms $g_{11}dx_{11}$, $g_{12}dx_{12}$, $g_{21}dx_{21}$ and $g_{22}dx_{22}$, that is, $\operatorname{tr}(G^\top dX) = \sum_{ij}g_{ij}\,dx_{ij}$. Every entry of $dX$ is multiplied by the entry of $G$ at the same position.
3. Compare with the multivariable chain rule $df = \sum_{ij}\frac{\partial f}{\partial x_{ij}}dx_{ij}$: so $g_{ij} = \partial f/\partial x_{ij}$.

Example. $f(X) = \operatorname{tr}(AX)$ with $A = \begin{pmatrix} 1 & 2 \\ 3 & 4\end{pmatrix}$. The diagonal of $AX$ gives $f = (x_{11} + 2x_{21}) + (3x_{12} + 4x_{22})$, so

The rule gives the same in one line: $df = \operatorname{tr}(A\,dX) = \operatorname{tr}\big((A^\top)^\top dX\big)$.

Trace tools. A scalar equals its own trace; $\operatorname{tr}(M^\top) = \operatorname{tr}(M)$; and $\operatorname{tr}(ABC) = \operatorname{tr}(BCA) = \operatorname{tr}(CAB)$ (cyclic moves only).

Two first examples.

• $f(X) = \operatorname{tr}(AX)$: $df = \operatorname{tr}(A\,dX) = \operatorname{tr}\big((A^\top)^\top dX\big)$, so $\nabla f = A^\top$.
• $f(x) = x^\top A x$: $df = dx^\top A x + x^\top A\, dx$. The first term is a scalar, so it equals its transpose $x^\top A^\top dx$. Hence $df = x^\top (A + A^\top)\, dx$ and $\nabla f = (A + A^\top)x$, which is $2Ax$ for a symmetric $A$.

So the quadratic $f(x) = \tfrac12 x^\top A x - b^\top x$ has $\nabla f = \tfrac12(A + A^\top)x - b$. Only the symmetric part of $A$ enters $f$, because $x^\top K x = 0$ for a skew-symmetric $K$. The shortcut $\nabla f = Ax - b$ is right only when $A = A^\top$.

Interactive figure: loading…

Theorem (derivative of the inverse) #

For an invertible $X$: $\quad d(X^{-1}) = -X^{-1}\, dX\, X^{-1}$.

Proof.

1. Differentiate the identity $X X^{-1} = I$. The right side is constant, so $d(XX^{-1}) = 0$.
2. Product rule, with the order kept: $dX\, X^{-1} + X\, d(X^{-1}) = 0$.
3. Multiply on the left by $X^{-1}$: $d(X^{-1}) = -X^{-1}\, dX\, X^{-1}$. $\blacksquare$

Why $-X^{-2}\,dX$ is wrong. To get $-X^{-2}dX$ you must move $dX$ past $X^{-1}$, and matrices do not commute. A small example shows the difference. Take $X = \operatorname{diag}(1, 2)$ and change one off-diagonal entry: $dX = \begin{pmatrix} 0 & 0.1 \\ 0 & 0\end{pmatrix}$.

• Exact: $(X + dX)^{-1} = \begin{pmatrix} 1 & 0.1 \\ 0 & 2\end{pmatrix}^{-1} = \begin{pmatrix} 1 & -0.05 \\ 0 & 0.5\end{pmatrix}$, so the change is $\begin{pmatrix} 0 & -0.05 \\ 0 & 0\end{pmatrix}$.
• $-X^{-1}dX\,X^{-1}$: the entry $0.1$ is multiplied by $(X^{-1})_{11} = 1$ from the left and by $(X^{-1})_{22} = 0.5$ from the right, which gives $-0.05$. Correct (here even exactly).
• $-X^{-2}dX$: the entry is multiplied by $(X^{-2})_{11} = 1$ only, which gives $-0.1$. Twice too large.

Only when $dX$ commutes with $X$ (for example in the $1\times1$ case) do the two agree.

Example: $f(X) = a^\top X^{-1} b$ has $df = -a^\top X^{-1} dX\, X^{-1} b = -\operatorname{tr}(X^{-1} b\, a^\top X^{-1} dX)$, so $\nabla f = -X^{-\top} a\, b^\top X^{-\top}$.

Theorem (Jacobi's formula and $\log\det$) #

For an invertible $X$: $d\det X = \det X \cdot \operatorname{tr}(X^{-1}dX)$. If $\det X > 0$, then $d\log\det X = \operatorname{tr}(X^{-1}dX)$, so $\nabla_X \log\det X = X^{-\top}$ ($= X^{-1}$ for a symmetric $X$).

Proof.

1. Factor: $X + dX = X(I + E)$ with $E = X^{-1}dX$. The determinant is multiplicative, so $\det(X + dX) = \det X \cdot \det(I + E)$.
2. $\det(I + E) = 1 + \operatorname{tr}E + O(\|E\|^2)$. For a $2\times2$ matrix you can see it directly: $\det(I + E) = (1 + e_{11})(1 + e_{22}) - e_{12}e_{21} = 1 + \operatorname{tr}E + \det E$, and $\det E$ is second order. The fold below shows it for every size.
3. Keep the linear part: $d\det X = \det X \cdot \operatorname{tr}(X^{-1}dX)$.
4. Chain rule with $d\log u = du/u$: $d\log\det X = \operatorname{tr}(X^{-1}dX)$.
5. Identify: $\operatorname{tr}(X^{-1}dX) = \operatorname{tr}\big((X^{-\top})^\top dX\big)$, so the gradient is $X^{-\top}$. $\blacksquare$

Step 2 for every size: $\det(I + E) = 1 + \operatorname{tr}E + O(\|E\|^2)$

In the permutation expansion of $\det(I + E)$, the identity permutation gives $\prod_i (1 + E_{ii}) = 1 + \sum_i E_{ii} + O(\|E\|^2)$. Every other permutation moves at least two indices, so it uses at least two off-diagonal entries, and each is an entry of $E$: that term is $O(\|E\|^2)$. The argument uses no eigenvalues, so it holds for every square $E$, also when $X$ is not diagonalizable.

The same formula holds for $\log|\det X|$ at every invertible $X$ (for $\det X < 0$, apply step 4 to $-\det X$).

The softmax Jacobian #

For $z \in \mathbb{R}^n$, $s = \operatorname{softmax}(z)$ has entries $s_i = e^{z_i}/Z$ with $Z = \sum_j e^{z_j}$.

Result. $ds = J\,dz$ with $J = \operatorname{diag}(s) - ss^\top$.

One entry first. Use the quotient rule on $s_1 = e^{z_1}/Z$, and note that $\partial Z/\partial z_j = e^{z_j}$.

• Same index: $\dfrac{\partial s_1}{\partial z_1} = \dfrac{e^{z_1}Z - e^{z_1}e^{z_1}}{Z^2} = s_1 - s_1^2 = s_1(1 - s_1)$.
• Other index: $\dfrac{\partial s_1}{\partial z_2} = \dfrac{0\cdot Z - e^{z_1}e^{z_2}}{Z^2} = -s_1s_2$.

So $J_{ij} = s_i\delta_{ij} - s_is_j$: the diagonal holds $s_i(1 - s_i)$, the rest $-s_is_j$. In matrix form this is $\operatorname{diag}(s) - ss^\top$. Check: equal logits for $n = 2$ give $s = (0.5, 0.5)$ and

Raising $z_1$ raises $s_1$ and lowers $s_2$ by the same amount, because $s_1 + s_2 = 1$.

The same result without indices.

1. Take logs: $\log s_i = z_i - \log\sum_j e^{z_j}$.
2. Differentiate: $d\log\sum_j e^{z_j} = \sum_j e^{z_j} dz_j / \sum_l e^{z_l} = s^\top dz$. So $d \log s_i = dz_i - s^\top dz$.
3. Use $ds_i = s_i\, d\log s_i$: $ds_i = s_i\,dz_i - s_i\,(s^\top dz)$, that is, $ds = (\operatorname{diag}(s) - ss^\top)\,dz$.

Interactive figure: loading…

Properties.

• $J$ is symmetric, and $J\mathbf{1} = s - s\,(s^\top\mathbf{1}) = 0$ ($\mathbf 1$ is the all-ones vector): adding a constant to all logits does not change the softmax. So $J$ is singular.
• $J$ is PSD: $v^\top J v = \sum_i s_i v_i^2 - \big(\sum_i s_i v_i\big)^2$ is the variance of $v_i$ when the index $i$ is drawn with probabilities $s$. If all $s_i > 0$, it is zero only for constant $v$, so $\operatorname{rank} J = n - 1$.
• Cross-entropy $L = -\log s_y$ for the true class $y$: by step 2, $dL = -(dz_y - s^\top dz) = (s - e_y)^\top dz$, so $\nabla_z L = s - e_y$ ($e_y$ is the one-hot vector of class $y$).
• Backprop never builds $J$. With the upstream gradient $g = \partial L/\partial s$, the VJP is $J^\top g = s \odot \big(g - (s^\top g)\mathbf{1}\big)$ ($\odot$ is the entrywise product): $O(n)$ work, not $O(n^2)$.

Matrix functions via the spectrum #

Definition. Let $S = Q\Lambda Q^\top$ be symmetric (lesson 06), and let $f$ be defined at every eigenvalue. Then

For a diagonalizable $A = X\Lambda X^{-1}$, in the same way $f(A) = X f(\Lambda) X^{-1}$.

Why this is the right definition. $S^k = Q\Lambda^k Q^\top$, because the inner factors $Q^\top Q = I$ cancel. So $p(S) = Q\,p(\Lambda)\,Q^\top$ for every polynomial $p$. A power series such as $e^x = \sum_k x^k/k!$ is a limit of polynomials, so the definition agrees with $e^S = \sum_k S^k/k!$. It also does not depend on the choice of eigenbasis: inside an eigenspace, $f$ multiplies the whole projector onto it by the single number $f(\lambda)$.

The three main functions.

• $e^S$: defined for every symmetric $S$; it is symmetric positive definite, with eigenvalues $e^{\lambda_i} > 0$. The series $e^A = \sum_k A^k/k!$ defines $e^A$ for every square $A$, even a defective one.
• $S^{1/2} = Q\Lambda^{1/2}Q^\top$ for $S \succeq 0$: the unique PSD matrix whose square is $S$. It is used for whitening ($S^{-1/2}x$) and in the Fréchet distance between Gaussians.
• $\log S = Q \log\Lambda\, Q^\top$ for $S \succ 0$: symmetric, with $e^{\log S} = S$. A useful identity: $\log\det S = \sum_i \log\lambda_i = \operatorname{tr}\log S$.

Rules that do not carry over from numbers. $e^{A+B} = e^A e^B$ holds when $AB = BA$ (then the binomial theorem works for $(A + B)^k$). Without this condition it fails in general. In the same way, $\frac{d}{dt}e^{A(t)} = A'(t)\,e^{A(t)}$ holds when $A'(t)$ commutes with $A(t)$, and fails in general.

Theorem (linear ODEs) #

For $A \in \mathbb{R}^{n\times n}$ and $x_0 \in \mathbb{R}^n$, the unique solution of $x'(t) = Ax(t)$ with $x(0) = x_0$ is $x(t) = e^{At}x_0$.

The key step is the derivative of the series: $\frac{d}{dt}e^{At} = Ae^{At}$, just as $\frac{d}{dt}e^{at} = ae^{at}$ for a number $a$.

Proof (convergence, derivative, uniqueness)

1. The series $e^{At} = \sum_{k\ge0} t^k A^k/k!$ converges for every $t$: in the spectral norm ($\|A^k\|_2 \le \|A\|_2^k$) the $k$-th term has norm at most $(|t|\,\|A\|_2)^k/k!$, and these numbers sum to $e^{|t|\,\|A\|_2}$.
2. Differentiate term by term (allowed for a convergent power series): $\frac{d}{dt}e^{At} = \sum_{k\ge1} \frac{k\,t^{k-1}A^k}{k!} = A\sum_{k\ge0}\frac{t^kA^k}{k!} = Ae^{At}$.
3. So $x(t) = e^{At}x_0$ satisfies $x' = Ax$, and $x(0) = Ix_0 = x_0$.
4. Uniqueness. Let $y$ be any solution. $e^{-At}$ commutes with $A$, so $\frac{d}{dt}\big(e^{-At}y\big) = e^{-At}(-Ay + y') = 0$. Hence $e^{-At}y(t) = y(0) = x_0$. As $At$ and $-At$ commute, $e^{At}e^{-At} = e^{0} = I$, so $y(t) = e^{At}x_0$. $\blacksquare$

Reading the solution. If $A = X\Lambda X^{-1}$ with eigenvectors $x_i$ and $c = X^{-1}x_0$, then $x(t) = \sum_i c_i e^{\lambda_i t} x_i$: each eigen-direction grows or decays at its own rate. The solution goes to 0 for every $x_0$ exactly when all $\operatorname{Re}\lambda_i < 0$. A $2\times2$ Jordan block $J_\lambda = \lambda I + N$, with $N = \begin{pmatrix}0&1\\0&0\end{pmatrix}$ and $N^2 = 0$, gives $e^{J_\lambda t} = e^{\lambda t}(I + tN)$, the $t e^{\lambda t}$ terms of lesson 05. Sampling at a step $\Delta$ gives the exact recurrence $x_{k+1} = e^{A\Delta}x_k$.

Interactive figure: loading…

How the figure computes $e^{tA}$, and how libraries do it

For a $2\times2$ matrix the figure computes $e^{tA}$ exactly from the Cayley–Hamilton theorem ($A^2 = 2\tau A - (\det A) I$ with $\tau = \tfrac12\operatorname{tr}A$): with $\delta^2 = \tau^2 - \det A$,

$$e^{tA} = e^{\tau t}\Big(\cosh(\delta t)\,I + \frac{\sinh(\delta t)}{\delta}\,(A - \tau I)\Big),$$

read with $\cosh(\delta t) = \cos(\omega t)$ and $\sinh(\delta t)/\delta = \sin(\omega t)/\omega$ when $\delta^2 = -\omega^2 < 0$, and with $1$ and $t$ when $\delta = 0$. For larger matrices libraries use scaling and squaring: $e^{A} = \big(e^{A/2^s}\big)^{2^s}$, with a Padé or Taylor approximation of $e^{A/2^s}$ (the notebook builds one).

Kronecker product and vec #

Definitions. $\operatorname{vec}(X) \in \mathbb{R}^{mn}$ stacks the columns of $X \in \mathbb{R}^{m\times n}$, the first column on top. For $A \in \mathbb{R}^{m\times n}$ and $B \in \mathbb{R}^{p\times q}$, the Kronecker product is the $mp \times nq$ block matrix

Rules. $(A\otimes B)(C\otimes D) = AC \otimes BD$ (the mixed-product rule, when $AC$ and $BD$ exist); $(A\otimes B)^\top = A^\top\otimes B^\top$; $(A\otimes B)^{-1} = A^{-1}\otimes B^{-1}$ (for invertible $A$, $B$); $\operatorname{vec}(uv^\top) = v\otimes u$; $\operatorname{tr}(A^\top X) = \operatorname{vec}(A)^\top\operatorname{vec}(X)$.

Theorem (vec of a product) #

For $A \in \mathbb{R}^{p\times m}$, $X \in \mathbb{R}^{m\times n}$, $B \in \mathbb{R}^{n\times q}$: $\quad \operatorname{vec}(AXB) = (B^\top \otimes A)\operatorname{vec}(X)$.

Proof.

1. Let $x_j$ be column $j$ of $X$ and $b_{jk}$ the entries of $B$. Column $k$ of $XB$ is $\sum_j b_{jk}\,x_j$.
2. So column $k$ of $AXB$ is $\sum_j b_{jk}\,(Ax_j)$.
3. Block $k$ of $\operatorname{vec}(AXB)$ is column $k$ of $AXB$, and block $j$ of $\operatorname{vec}(X)$ is $x_j$. Step 2 says: block $k$ of the output $= \sum_j (b_{jk}A)\,(\text{block } j \text{ of the input})$.
4. So the matrix of the map has block $(k, j)$ equal to $b_{jk}A = (B^\top)_{kj}A$. This is the definition of $B^\top\otimes A$. $\blacksquare$

Special cases: $\operatorname{vec}(AX) = (I\otimes A)\operatorname{vec}X$ and $\operatorname{vec}(XB) = (B^\top\otimes I)\operatorname{vec}X$. Example: a linear layer $Y = WX$ has $dY = dW\,X$, so the Jacobian of $\operatorname{vec}Y$ with respect to $\operatorname{vec}W$ is $X^\top\otimes I$. A VJP never forms it: for the upstream gradient $G = \partial L/\partial Y$, $(X^\top\otimes I)^\top\operatorname{vec}G = (X\otimes I)\operatorname{vec}G = \operatorname{vec}(GX^\top)$, the familiar weight gradient $GX^\top$.

Theorem (eigenvalues of Kronecker products) #

Let $A$ be $n\times n$ with eigenvalues $\lambda_1, \dots, \lambda_n$, and $B$ be $m\times m$ with eigenvalues $\mu_1, \dots, \mu_m$ (with multiplicity). Then $A\otimes B$ has the $nm$ eigenvalues $\lambda_i\mu_j$, and the Kronecker sum $A\otimes I_m + I_n\otimes B$ has the $nm$ eigenvalues $\lambda_i + \mu_j$.

Why (eigenvectors). If $Ax = \lambda x$ and $By = \mu y$, the mixed-product rule gives $(A\otimes B)(x\otimes y) = Ax\otimes By = \lambda\mu\,(x\otimes y)$ and $(A\otimes I + I\otimes B)(x\otimes y) = \lambda x\otimes y + x\otimes\mu y = (\lambda + \mu)\,(x\otimes y)$. When $A$ and $B$ are diagonalizable, the $nm$ vectors $x_i\otimes y_j$ form a basis, and the theorem follows.

Proof with multiplicities, also for non-diagonalizable matrices (Schur form)

1. Take Schur forms $A = UTU^*$ and $B = VRV^*$, with $U, V$ unitary and $T, R$ upper triangular (diagonals: the eigenvalues).
2. Mixed product: $A\otimes B = (U\otimes V)(T\otimes R)(U\otimes V)^*$, and $U\otimes V$ is unitary. So $A\otimes B$ is similar to $T\otimes R$.
3. $T\otimes R$ is upper triangular: its blocks $t_{ij}R$ below the diagonal are zero, and its diagonal blocks $t_{ii}R$ are upper triangular. Its diagonal entries are $t_{ii}r_{jj} = \lambda_i\mu_j$.
4. In the same way $A\otimes I + I\otimes B$ is similar to $T\otimes I + I\otimes R$, which is upper triangular with diagonal entries $\lambda_i + \mu_j$. $\blacksquare$

Consequences: $\operatorname{tr}(A\otimes B) = \operatorname{tr}A\cdot\operatorname{tr}B$, $\det(A\otimes B) = (\det A)^m(\det B)^n$, the singular values of $A\otimes B$ are $\sigma_i(A)\sigma_j(B)$, so $\kappa(A\otimes B) = \kappa(A)\kappa(B)$.

The Sylvester equation $AX + XB = C$ is linear in $X$. With vec it reads $(I\otimes A + B^\top\otimes I)\operatorname{vec}X = \operatorname{vec}C$. That matrix is a Kronecker sum with eigenvalues $\lambda_i(A) + \mu_j(B)$ ($B^\top$ has the eigenvalues of $B$). So the solution is unique exactly when $\lambda_i(A) + \mu_j(B) \ne 0$ for all $i, j$.

Answers to section 1. (1) No. The derivative is $-X^{-1}\,dX\,X^{-1}$. It equals $-X^{-2}dX$ only when $dX$ commutes with $X$, as in the $1\times1$ case. (2) Not in general; it holds when $AB = BA$. For $A = \begin{pmatrix}0&1\\0&0\end{pmatrix}$ and $B = A^\top$: $e^Ae^B = \begin{pmatrix}1&1\\0&1\end{pmatrix}\begin{pmatrix}1&0\\1&1\end{pmatrix} = \begin{pmatrix}2&1\\1&1\end{pmatrix}$, but $e^{A+B} = \begin{pmatrix}\cosh 1&\sinh 1\\ \sinh 1&\cosh 1\end{pmatrix} \approx \begin{pmatrix}1.543&1.175\\1.175&1.543\end{pmatrix}$.

4. Worked example #

(a) $\log\det$ and the inverse, with numbers. $X = \begin{pmatrix}2&1\\1&2\end{pmatrix}$ has $\det X = 3$ and $X^{-1} = \frac13\begin{pmatrix}2&-1\\-1&2\end{pmatrix}$. Change one entry: $dX = 0.01\,E_{11}$.

• Prediction: $d\log\det X = \operatorname{tr}(X^{-1}dX) = 0.01\cdot\frac23 \approx 0.0066667$. Exact: $\det(X + dX) = 2.01\cdot 2 - 1 = 3.02$ and $\log 3.02 - \log 3 \approx 0.0066445$. The difference, about $2.2\times10^{-5}$, is second order.
• Prediction: $-X^{-1}dX\,X^{-1} = -0.01\,(X^{-1}e_1)(e_1^\top X^{-1}) = -\frac{0.01}{9}\begin{pmatrix}4&-2\\-2&1\end{pmatrix} \approx \begin{pmatrix}-0.004444&0.002222\\0.002222&-0.001111\end{pmatrix}$. Exact: $(X + dX)^{-1} - X^{-1} \approx \begin{pmatrix}-0.004415&0.002208\\0.002208&-0.001104\end{pmatrix}$.

(b) Softmax. $z = (0, \ln 2, \ln 3)$ gives $e^z = (1, 2, 3)$ and $s = (\frac16, \frac13, \frac12)$. Then

For example $J_{11} = \frac16\cdot\frac56 = \frac5{36}$ and $J_{12} = -\frac16\cdot\frac13 = -\frac2{36}$. Each row sums to 0 ($J\mathbf 1 = 0$). With the true class $y = 3$: $\nabla_z L = s - e_3 = (\frac16, \frac13, -\frac12)$.

(c) A square root and a linear ODE. $S = \begin{pmatrix}2&1\\1&2\end{pmatrix}$ has $\lambda = 3$ with $q_1 = (1,1)/\sqrt2$ and $\lambda = 1$ with $q_2 = (1,-1)/\sqrt2$. With the projectors $P_1 = q_1q_1^\top = \frac12\begin{pmatrix}1&1\\1&1\end{pmatrix}$ and $P_2 = q_2q_2^\top = \frac12\begin{pmatrix}1&-1\\-1&1\end{pmatrix}$:

Squaring gives back $S$. In the same way $\log S = \ln 3\cdot P_1$ (because $\ln 1 = 0$), and $\operatorname{tr}\log S = \ln 3 = \log\det S$.

Now heat flow in two cells. Each cell exchanges heat with the other (rate 1) and loses heat to the surroundings (rate 1): $x_1' = (x_2 - x_1) - x_1$, and the same for $x_2$. So $x' = Ax$ with $A = \begin{pmatrix}-2&1\\1&-2\end{pmatrix}$, which has eigenvalue $-1$ for $q_1$ and $-3$ for $q_2$. Start with all heat in cell 1: $x_0 = (1, 0) = \frac1{\sqrt2}(q_1 + q_2)$. Then

The average temperature decays at rate 1; the difference between the cells decays three times faster. At $t = 1$: $x \approx (0.2088, 0.1590)$.

(d) vec and Kronecker, step by step. $A = \begin{pmatrix}1&2\\3&4\end{pmatrix}$, $X = A$, and $B = \begin{pmatrix}0&1\\1&0\end{pmatrix}$ (it swaps columns).

1. Left side: $AX = \begin{pmatrix}7&10\\15&22\end{pmatrix}$ and $AXB = \begin{pmatrix}10&7\\22&15\end{pmatrix}$, so $\operatorname{vec}(AXB) = (10, 22, 7, 15)$ (first column on top).
2. $\operatorname{vec}X = (1, 3, 2, 4)$: block 1 is $x_1 = (1, 3)$, block 2 is $x_2 = (2, 4)$.
3. $B^\top = B$, so block $(k, j)$ of $B^\top\otimes A$ is $b_{jk}A$: blocks $(1,1)$ and $(2,2)$ are $0\cdot A$, blocks $(1,2)$ and $(2,1)$ are $1\cdot A$. So $B^\top\otimes A = \begin{pmatrix}0&A\\A&0\end{pmatrix}$.
4. Block row 1 gives $0\cdot x_1 + A x_2 = A(2, 4) = (10, 22)$; block row 2 gives $Ax_1 + 0\cdot x_2 = A(1, 3) = (7, 15)$.
5. Stacked: $(10, 22, 7, 15)$, the same vector as in step 1. ✓

Eigenvalues: $A_1 = \begin{pmatrix}2&1\\0&3\end{pmatrix}$ (eigenvalues 2, 3) and $B_1 = \begin{pmatrix}5&0\\2&7\end{pmatrix}$ (eigenvalues 5, 7). $A_1\otimes B_1$ has eigenvalues $10, 14, 15, 21$, and $A_1\otimes I + I\otimes B_1$ has $7, 9, 8, 10$. Check: $\det(A_1\otimes B_1) = 6^2\cdot35^2 = 44100 = 10\cdot14\cdot15\cdot21$.

5. In code #

NumPy and SciPy check the hand computations above; finite differences stand in for autodiff. The notebook for this lesson has these cells and the exercises.

import numpy as np
from scipy.linalg import expm, sqrtm, logm
from scipy.special import softmax

X = np.array([[2., 1.], [1., 2.]])
Xi = np.linalg.inv(X)
dX = np.array([[0.01, 0.], [0., 0.]])

# d log det X = tr(X^{-1} dX): equal to first order
print(round(np.linalg.slogdet(X + dX)[1] - np.linalg.slogdet(X)[1], 7),
      round(np.trace(Xi @ dX), 7))                        # 0.0066445 0.0066667
# d(X^{-1}) = -X^{-1} dX X^{-1}
print(np.round(np.linalg.inv(X + dX) - Xi, 6))            # [[-0.004415  0.002208] [ 0.002208 -0.001104]]
print(np.round(-Xi @ dX @ Xi, 6))                         # [[-0.004444  0.002222] [ 0.002222 -0.001111]]

def num_grad(f, X, h=1e-6):
    """Central differences, one entry at a time."""
    G = np.zeros_like(X)
    for idx in np.ndindex(X.shape):
        E = np.zeros_like(X); E[idx] = h
        G[idx] = (f(X + E) - f(X - E)) / (2 * h)
    return G

# the gradient of log det is X^{-T}, not X^{-1} (non-symmetric X here)
Xn = np.array([[2., 1.], [0., 3.]])
g = num_grad(lambda M: np.linalg.slogdet(M)[1], Xn)
print(np.allclose(g, np.linalg.inv(Xn).T, atol=1e-8), np.allclose(g, np.linalg.inv(Xn)))   # True False

# softmax Jacobian = diag(s) - s s^T (columns by finite differences)
z = np.log(np.array([1., 2., 3.]))
s = softmax(z)
J = np.column_stack([(softmax(z + 1e-6 * e) - softmax(z - 1e-6 * e)) / 2e-6 for e in np.eye(3)])
print(np.round(36 * J, 6))                                # [[5 -2 -3] [-2 8 -6] [-3 -6 9]]
print(np.allclose(J, np.diag(s) - np.outer(s, s)), np.allclose(J @ np.ones(3), 0))  # True True

# matrix functions: f(S) = Q f(Lambda) Q^T, not entrywise
S = X
lam, Q = np.linalg.eigh(S)
R = Q @ np.diag(np.sqrt(lam)) @ Q.T
print(np.round(R, 6), np.allclose(R, sqrtm(S)), np.allclose(R @ R, S))   # [[1.366025 0.366025] ...] True True
print(np.allclose(expm(S), Q @ np.diag(np.exp(lam)) @ Q.T), np.allclose(expm(S), np.exp(S)))  # True False
print(np.isclose(np.trace(logm(S)), np.log(np.linalg.det(S))))          # True: tr log S = log det S

# x' = Ax is solved by x(t) = e^{At} x0 (the two-cell heat example)
A = np.array([[-2., 1.], [1., -2.]])
print(np.round(expm(A * 1.0) @ np.array([1., 0.]), 6))   # [0.208833 0.159046]

# e^{A+B} != e^A e^B when AB != BA
N = np.array([[0., 1.], [0., 0.]])
print(np.round(expm(N) @ expm(N.T), 4))                   # [[2. 1.] [1. 1.]]
print(np.round(expm(N + N.T), 4))                         # [[1.5431 1.1752] [1.1752 1.5431]]

# vec(AXB) = (B^T kron A) vec(X); vec stacks COLUMNS, so use order="F"
A2 = np.array([[1., 2.], [3., 4.]]); B2 = np.array([[0., 1.], [1., 0.]])
vec = lambda M: M.reshape(-1, order="F")
print(vec(A2 @ A2 @ B2), np.kron(B2.T, A2) @ vec(A2))    # [10. 22.  7. 15.] twice

# eigenvalues: products for A kron B, sums for the Kronecker sum
A1 = np.array([[2., 1.], [0., 3.]]); B1 = np.array([[5., 0.], [2., 7.]]); I2 = np.eye(2)
print(np.sort(np.linalg.eigvals(np.kron(A1, B1)).real))                      # [10. 14. 15. 21.]
print(np.sort(np.linalg.eigvals(np.kron(A1, I2) + np.kron(I2, B1)).real))    # [ 7.  8.  9. 10.]

The gradient check of section 2, as code: the error of the first-order prediction along a random direction, for the correct gradient $X^{-\top}$ of $\log\det$ and for the wrong $X^{-1}$.

import numpy as np

rng = np.random.default_rng(0)
X = np.array([[2., 1.], [0.5, 3.]])            # not symmetric, so X^{-T} != X^{-1}
D = rng.standard_normal((2, 2)); D /= np.linalg.norm(D)
f = lambda M: np.linalg.slogdet(M)[1]
for G, name in [(np.linalg.inv(X).T, "X^{-T}"), (np.linalg.inv(X), "X^{-1}")]:
    err = [abs(f(X + t * D) - f(X) - t * np.sum(G * D)) for t in (1e-3, 5e-4)]
    print(name, round(err[0] / err[1], 2))      # X^{-T} 4.0 (O(t^2)), X^{-1} 2.0 (O(t))

6. Where it shows up #

• Gaussian maximum likelihood. For $N$ samples with known mean $\mu$, $\ell(\Sigma) = -\frac N2\big(\log\det\Sigma + \operatorname{tr}(\Sigma^{-1}\hat S)\big) + \text{const}$, with $\hat S = \frac1N\sum_n (x_n - \mu)(x_n - \mu)^\top$. The two theorems give $d\ell = -\frac N2\operatorname{tr}\big((\Sigma^{-1} - \Sigma^{-1}\hat S\Sigma^{-1})\,d\Sigma\big)$. This is zero for every symmetric $d\Sigma$ exactly when $\Sigma = \hat S$.
• Gaussian processes and flows. GP hyperparameters need $\partial_\theta\log\det K = \operatorname{tr}(K^{-1}\partial_\theta K)$. Normalizing flows need $\log|\det J|$, so they use triangular Jacobians (the change-of-variables formula).
• Softmax. In attention backward, FlashAttention computes $dS_{ij} = P_{ij}(dP_{ij} - D_i)$ with $D_i = \sum_j P_{ij}\,dP_{ij}$: the VJP $s\odot(g - (s^\top g)\mathbf 1)$, row by row.
• Linear dynamics. $e^{At}$ solves semi-discretized heat equations ($u' = Lu$), linearized climate models and continuous-time Markov chains. Mamba discretizes $h' = Ah + Bx$ with step $\Delta$ by $\bar A = e^{\Delta A}$ (zero-order hold). Its $A$ is diagonal, so here "entrywise" is right.
• SPD square roots and logs. Whitening $\Sigma^{-1/2}x$; the Fréchet distance behind FID, $\|\mu_1-\mu_2\|^2 + \operatorname{tr}\big(\Sigma_1 + \Sigma_2 - 2(\Sigma_1^{1/2}\Sigma_2\Sigma_1^{1/2})^{1/2}\big)$; the log-Euclidean distance $\|\log S_1 - \log S_2\|_F$ between covariance matrices (lesson 19).
• K-FAC and Shampoo. For a layer $s = Wa$ with backpropagated $g = \partial L/\partial s$, the weight gradient is $ga^\top$, and $\operatorname{vec}(ga^\top) = a\otimes g$. K-FAC approximates the Fisher block $\mathbb E[aa^\top\otimes gg^\top]$ by $\bar A\otimes\bar G$ ($\bar A = \mathbb E[aa^\top]$, $\bar G = \mathbb E[gg^\top]$). Then $(\bar A\otimes\bar G)^{-1}\operatorname{vec}(ga^\top) = \operatorname{vec}(\bar G^{-1}\,ga^\top\,\bar A^{-1})$: two small inverses instead of one huge one. Shampoo steps along $H_L^{-1/4}M_tH_R^{-1/4}$ for the gradient matrix $M_t$, with $H_L = \sum_t M_tM_t^\top$ and $H_R = \sum_t M_t^\top M_t$. By the vec rule this is the Kronecker preconditioner $H_R^{-1/4}\otimes H_L^{-1/4}$, built from two matrix functions.
• Separable kernels and grids (ML for science). A space–time GP on a grid often uses $K = K_{\text{time}}\otimes K_{\text{space}}$. Then $K + \sigma^2I$ has eigenvalues $\lambda_i\mu_j + \sigma^2$, so the solve and the $\log\det$ need only the eigendecompositions of the two small factors. The 2-D Poisson matrix is $T_m\otimes I + I\otimes T_m$ ($T_m$ the 1-D second-difference matrix): its eigenvalues are sums of 1-D eigenvalues, which gives its condition number in lesson 14.

7. Common mistakes #

• Losing the order. $d(X^{-1}) \ne -X^{-2}dX$, and $d(X^2) = dX\,X + X\,dX$, not $2X\,dX$.
• $\nabla\log\det X = X^{-1}$ for every $X$. It is $X^{-\top}$. For symmetric $X$, some books give $2X^{-1} - \operatorname{diag}(X^{-1})$: the same differential, written in the free entries only (so $X_{ij}$ and $X_{ji}$ move together). Safer: use the full-matrix gradient with a symmetric $dX$.
• $\nabla(\tfrac12 x^\top Ax - b^\top x) = Ax - b$ for every $A$. It is $\tfrac12(A + A^\top)x - b$.
• Entrywise instead of matrix functions. np.exp(S) and np.sqrt(S) act entry by entry. Use scipy.linalg.expm, sqrtm, logm, or $Qf(\Lambda)Q^\top$.
• $e^{A+B} = e^Ae^B$ without $AB = BA$. The same goes for $\log(AB) = \log A + \log B$.
• Products vs sums. $A\otimes B$ has eigenvalues $\lambda_i\mu_j$; the Kronecker sum has $\lambda_i + \mu_j$.
• Row-major vec. NumPy's ravel() stacks rows; column-stacking vec is X.reshape(-1, order="F"). With row stacking the rule becomes $\operatorname{vec}_r(AXB) = (A\otimes B^\top)\operatorname{vec}_r(X)$.
• Building Kronecker products or full Jacobians. $B^\top\otimes A$ for two $1000\times1000$ matrices has $10^{12}$ entries. Apply it as $AXB$.
• Softmax Jacobian $= \operatorname{diag}(s)$. This misses $-ss^\top$; the true $J$ is singular.
• Trusting a gradient check at one step size. One small error at one $t$ can hide a wrong gradient (for example when $t$ is so small that rounding dominates). Compare two step sizes and look for the factor 4.

8. Check yourself #

1. Gradient of $f(X) = a^\top X b$?
   Answer

   $df = a^\top dX\,b = \operatorname{tr}(b\,a^\top dX)$, so $\nabla f = (ba^\top)^\top = ab^\top$.
2. Gradient of $f(X) = \operatorname{tr}(X^{-1})$?
   Answer

   $df = -\operatorname{tr}(X^{-1}dX\,X^{-1}) = -\operatorname{tr}(X^{-2}dX)$ (cyclic move), so $\nabla f = -(X^{-2})^\top$.
3. Why is $\det e^A = e^{\operatorname{tr}A}$?
   Answer

   With a Schur form $A = UTU^*$, $e^A = Ue^TU^*$ and $e^T$ is upper triangular with diagonal $e^{\lambda_i}$. So $\det e^A = \prod_i e^{\lambda_i} = e^{\sum_i\lambda_i} = e^{\operatorname{tr}A}$.
4. Solve $x' = Ax$ for $A = \begin{pmatrix}0&1\\-1&0\end{pmatrix}$, $x(0) = (1, 0)$. Why does $\|x(t)\|$ stay 1?
   Answer

   $A^2 = -I$, so the series gives $e^{At} = \cos t\,I + \sin t\,A = \begin{pmatrix}\cos t&\sin t\\-\sin t&\cos t\end{pmatrix}$ and $x(t) = (\cos t, -\sin t)$. $A$ is skew-symmetric, so $(e^{At})^\top = e^{-At} = (e^{At})^{-1}$: $e^{At}$ is orthogonal and keeps lengths.
5. $A$ ($2\times2$) has eigenvalues $1, 2$ and $B$ ($2\times2$) has eigenvalues $-1, 4$. Give the eigenvalues of $A\otimes B$ and of $A\otimes I + I\otimes B$. Does $AX + XB = C$ have a unique solution?
   Answer

   Products: $-1, 4, -2, 8$. Sums: $0, 5, 1, 6$. The Sylvester matrix $I\otimes A + B^\top\otimes I$ has the eigenvalues $\lambda_i + \mu_j$, and one of them is $1 + (-1) = 0$. The matrix is singular, so for a given $C$ there is either no solution or infinitely many.
6. Why is the softmax Jacobian PSD, and why is it singular?
   Answer

   $v^\top Jv = \sum_i s_iv_i^2 - (\sum_i s_iv_i)^2$ is a variance under the distribution $s$, so it is $\ge 0$. $J\mathbf 1 = 0$, because adding the same constant to all logits does not change $s$.

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. State the identification rule of matrix calculus.

   Show answer

2. Why must the product rule for matrices keep the order of the factors?

   Show answer

3. What is $d(X^{-1})$?

   Show answer

4. What is $d\log\det X$, and what is the gradient?

   Show answer

5. Compute $\nabla_X\operatorname{tr}(AX)$.

   Show answer

6. What is the Jacobian of $s = \operatorname{softmax}(z)$? Give it for $s = (0.5, 0.5)$.

   Show answer

7. How do you compute $f(S)$ for a symmetric matrix $S$?

   Show answer

8. What solves $x' = Ax$ with $x(0) = x_0$?

   Show answer

9. State the vec rule for $AXB$.

   Show answer

10. What are the eigenvalues of $A\otimes B$ and of $A\otimes I + I\otimes B$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Derive the MLE of a Gaussian covariance." — $d\log\det\Sigma = \operatorname{tr}(\Sigma^{-1}d\Sigma)$ and $d\Sigma^{-1} = -\Sigma^{-1}d\Sigma\,\Sigma^{-1}$ give $\Sigma = \hat S$ (section 6).
• "Jacobian of softmax? Gradient of softmax cross-entropy?" — $\operatorname{diag}(s) - ss^\top$; $s - e_y$.
• "How do you check a hand-written gradient?" — compare $f(X + t\,dX) - f(X)$ with $t\operatorname{tr}(G^\top dX)$ for a random $dX$ and two values of $t$; the error must shrink by 4 when $t$ is halved (central differences shrink it faster).
• "Why is $\log\det$ concave on SPD matrices?" — the second differential is $-\operatorname{tr}(X^{-1}dX\,X^{-1}dX) = -\|X^{-1/2}\,dX\,X^{-1/2}\|_F^2 \le 0$ for symmetric $dX$.
• "How do you solve $x' = Ax$, and when does it decay?" — $x(t) = e^{At}x_0$; it decays for every $x_0$ exactly when all $\operatorname{Re}\lambda_i(A) < 0$.
• "Why is K-FAC cheap to invert?" — $(\bar A\otimes\bar G)^{-1} = \bar A^{-1}\otimes\bar G^{-1}$, applied as $\bar G^{-1}(\text{gradient})\,\bar A^{-1}$ by the vec rule.

10. Practice #

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

1. Gradients by the differential method: $\log\det X$, $a^\top X^{-1}b$ and $\operatorname{tr}(X^{-1})$, checked against finite differences.
2. The softmax Jacobian $\operatorname{diag}(s) - ss^\top$ and its $O(n)$ VJP, and the cross-entropy gradient $s - e_y$.
3. Matrix functions $Qf(\Lambda)Q^\top$ for a symmetric matrix, checked against expm, sqrtm and logm.
4. The matrix exponential by scaling and squaring with a Taylor polynomial, checked against scipy.linalg.expm; then a demo solves the two-cell heat equation with $e^{At}x_0$.
5. Solve the Sylvester equation $AX + XB = C$ with the vec rule and a Kronecker sum, and check it against scipy.linalg.solve_sylvester.

11. Watch and read #

Lecture 1 Part 2: Derivatives as Linear OperatorsAlan Edelman, Steven G. Johnson · MIT 18.S096 Matrix Calculus (MIT OpenCourseWare) · 48 min

Lecture 5 Part 1: Derivative of Matrix Determinant and InverseAlan Edelman, Steven G. Johnson · MIT 18.S096 Matrix Calculus (MIT OpenCourseWare) · 28 min

Lecture 3 Part 1: Kronecker Products and JacobiansAlan Edelman, Steven G. Johnson · MIT 18.S096 Matrix Calculus (MIT OpenCourseWare) · 53 min

How (and why) to raise e to the power of a matrix | DE63Blue1Brown · Differential equations, chapter 6 · 27 min

23. Differential Equations and exp(At)Gilbert Strang · MIT 18.06 Linear Algebra (MIT OpenCourseWare) · 51 min

• Paige Bright, Alan Edelman, Steven G. Johnson, Matrix Calculus (for Machine Learning and Beyond) (MIT 18.S096 notes, free): ch. 2 "Derivatives as linear operators", §3.3 "Kronecker products", ch. 7 "Derivative of matrix determinant and inverse". The course page has all lectures and problem sets.
• K. B. Petersen, M. S. Pedersen, The Matrix Cookbook (2012), free PDF: §2.1–2.2 (derivatives of the determinant and the inverse), §2.5 (traces), §10.1 (functions and series), §10.2 (Kronecker and vec, eq. 519–520).
• Gilbert Strang, MIT 18.06 (2010), lecture 23: Differential equations and exp(At).
• James Martens, Roger Grosse, Optimizing neural networks with Kronecker-factored approximate curvature (ICML 2015), §3 (the Kronecker-factored Fisher).