Multivariable and Matrix Calculus

Optimisation theory is calculus with inequalities attached. The assumptions that appear in every convergence theorem later — $L$-smoothness, $\mu$-strong convexity, bounded gradients — are all statements about the objects defined here, and the descent lemma proved below is used in essentially every proof on the gradient descent page.

Derivatives as linear maps

DefinitionDifferentiability

$f : \R^d \to \R^m$ is differentiable at $x$ if there is a linear map $Df(x) : \R^d \to \R^m$ with

$$f(x+h) = f(x) + Df(x)h + o(\norm{h}) \quad \text{as } h \to 0 .$$

In coordinates $Df(x)$ is the Jacobian $J \in \R^{m \times d}$ with $J_{ij} = \partial f_i/\partial x_j$. When $m=1$ we write the derivative as a column vector, the gradient $\nabla f(x) \in \R^d$, so that $Df(x)h = \ip{\nabla f(x)}{h}$.

DefinitionDirectional derivative

The derivative of $f$ at $x$ along $v$ is

$$D_v f(x) = \lim_{t \downarrow 0} \frac{f(x+tv) - f(x)}{t} = \ip{\nabla f(x)}{v},$$

the last equality holding when $f$ is differentiable at $x$. A function can have all directional derivatives at a point without being differentiable there; $f(x_1,x_2) = x_1^3/(x_1^2+x_2^2)$ at the origin is the standard example.

DefinitionHessian

If $\nabla f$ is itself differentiable at $x$, its Jacobian is the Hessian $\nabla^2 f(x) \in \R^{d\times d}$ with $[\nabla^2 f(x)]_{ij} = \partial^2 f/\partial x_i \partial x_j$.

TheoremSymmetry of the Hessian

If the second partial derivatives of $f$ exist and are continuous in a neighbourhood of $x$, then $\nabla^2 f(x)$ is symmetric.

This is Clairaut–Schwarz; the proof is in any analysis text and is not reproduced here. The consequence we use constantly is that the spectral theorem applies to Hessians, so "curvature" is a list of real eigenvalues, and the second-order conditions for a minimum become eigenvalue conditions.

Taylor expansion and the descent lemma

TheoremTaylor's theorem with integral remainder

Let $f : \R^d \to \R$ be twice continuously differentiable. For all $x, y$,

$$f(y) = f(x) + \ip{\nabla f(x)}{y - x} + \int_0^1 (1-t)\,(y-x)^\top \nabla^2 f\big(x + t(y-x)\big)(y-x)\,dt .$$

In particular there is $\xi$ on the segment $[x,y]$ with $f(y) = f(x) + \ip{\nabla f(x)}{y-x} + \tfrac12 (y-x)^\top\nabla^2 f(\xi)(y-x)$.

ProofReveal

DefinitionL-smoothness

$f$ is $L$-smooth if $\nabla f$ is $L$-Lipschitz: $\norm{\nabla f(x) - \nabla f(y)}_2 \le L\norm{x-y}_2$ for all $x,y$. For twice-differentiable $f$ this is equivalent to $-L I \preceq \nabla^2 f(x) \preceq L I$ for all $x$.

LemmaDescent lemma

If $f$ is $L$-smooth then for all $x,y$,

$$\big| f(y) - f(x) - \ip{\nabla f(x)}{y-x} \big| \le \frac{L}{2}\norm{y-x}_2^2 .$$

In particular $f(y) \le f(x) + \ip{\nabla f(x)}{y-x} + \frac{L}{2}\norm{y-x}_2^2$.

ProofReveal

The chain rule and reverse-mode differentiation

TheoremChain rule

If $g : \R^d \to \R^m$ is differentiable at $x$ and $f : \R^m \to \R^k$ is differentiable at $g(x)$, then $f \circ g$ is differentiable at $x$ with

$$D(f\circ g)(x) = Df(g(x)) \cdot Dg(x) .$$

Consider a composition of $n$ maps, as in a feedforward network:

$$f = f_n \circ f_{n-1} \circ \dots \circ f_1, \qquad x_0 = x,\quad x_i = f_i(x_{i-1}) .$$

The chain rule gives $Df(x) = J_n J_{n-1}\cdots J_1$ with $J_i = Df_i(x_{i-1})$. Matrix multiplication is associative, so this product may be evaluated in any order — and the order is exactly the difference between forward and reverse mode.

1. Forward pass. Compute and store $x_1, \dots, x_n$.
2. Seed. Set $\bar{x}_n = \nabla_{x_n} \mathcal{L}$, the gradient of the scalar objective with respect to the output.
3. Backward pass. For $i = n, n-1, \dots, 1$: set $\bar{x}_{i-1} = J_i^\top \bar{x}_i$ and accumulate the gradient with respect to the parameters of layer $i$ as $(\partial x_i/\partial \theta_i)^\top \bar{x}_i$.
4. Return $\bar{x}_0$ and the accumulated parameter gradients.

PropositionCost of reverse mode

For a scalar-valued $f$ built from $n$ elementary operations, reverse mode computes the full gradient at a cost within a constant factor of one function evaluation, using memory proportional to the number of stored intermediates. Forward mode costs one function evaluation per input coordinate.

ProofReveal

Matrix derivative identities

Throughout, $x \in \R^d$, $A \in \R^{d\times d}$, $X \in \R^{n \times p}$, and all gradients are with respect to the indicated variable.

Function	Derivative	Notes
$a^\top x$	$a$
$x^\top A x$	$(A + A^\top)x$	$2Ax$ if $A$ symmetric
$\tfrac12\norm{Ax - b}_2^2$	$A^\top(Ax-b)$	normal equations set this to zero
$\norm{x}_2$	$x/\norm{x}_2$	undefined at $x=0$
$\log\det X$	$X^{-\top}$	$X \succ 0$
$\tr(AX)$	$A^\top$	with respect to $X$
$\tr(X^\top A X)$	$(A + A^\top)X$
$\norm{X}_F^2$	$2X$
$X^{-1}$	$-X^{-1}(dX)X^{-1}$	as a differential
$\sigma(x)$, $\sigma(t)=\tfrac{1}{1+e^{-t}}$	$\diag\!\big(\sigma(x)(1-\sigma(x))\big)$	elementwise
$\operatorname{softmax}(z)$	$\diag(p) - pp^\top$, $p = \operatorname{softmax}(z)$	PSD, singular on $\one$

The softmax Jacobian, since it is used in every LLM derivation

1. $p_i = e^{z_i}/\sum_k e^{z_k}$.

   Definition.

2. $$\frac{\partial p_i}{\partial z_j} = \frac{\delta_{ij}e^{z_i}\sum_k e^{z_k} - e^{z_i}e^{z_j}}{\left(\sum_k e^{z_k}\right)^2} = \delta_{ij}p_i - p_ip_j .$$

   Differentiate the quotient, separating the i = j and i ≠ j cases.

3. $\dfrac{\partial p}{\partial z} = \diag(p) - pp^\top$.

   Collect into matrix form.

4. For any $v$, $v^\top(\diag(p) - pp^\top)v = \E_{i \sim p}[v_i^2] - (\E_{i\sim p}[v_i])^2 = \Var_{i\sim p}(v_i) \ge 0$, and the expression vanishes for $v = \one$.

   The Jacobian is PSD and annihilates the all-ones vector, because shifting all logits by a constant does not change the distribution.

That last line is worth keeping: the softmax Jacobian is a variance under the predicted distribution. It is why gradients vanish when the model is confident — a confident $p$ has small variance — and it reappears as the Fisher information of the categorical family in the language modelling derivations.

Problems

Tightness of the descent lemmaProof · Beginner · 15 min The softmax Jacobian is a covarianceProof · Intermediate · 30 min Gradient and Hessian of log-determinantProof · Advanced · 30 min All directional derivatives exist, yet the function is not differentiableProof · Intermediate · 30 min