Matrix Calculus, JVPs, VJPs, and Hessian Products

The differential of a scalar f(x) can be written df = ∇f(x)ᵀdx. A scalar change is produced by pairing the gradient with an input perturbation. For y = Ax, the differential is dy = A dx. For the quadratic form f(x) = xᵀAx, the differential does work a memorised table cannot. Form df = f(x+dx) − f(x), drop the second-order term, read off ∇f = (A + Aᵀ)x. Nothing in that derivation assumes A is symmetric. It is Example 10 of Edelman and Johnson's 2023 MIT 18.S096 matrix-calculus notes. The Matrix Cookbook (2012 edition) records the same result with a linear term added: ∇x(xᵀAx + bᵀx) = (A + Aᵀ)x + b, whose Hessian is A + Aᵀ.

The notes compute that gradient the slow way first, component by component, and then warn the reader off the detour: “However, this approach is awkward, error-prone, labor-intensive, and quickly becomes worse as we move on to more complicated functions.” The differential method leaves you less to memorise. It also makes a transpose mistake easier to see.

A derivative maps a small input perturbation to a first-order output perturbation. Coordinates represent that linear map as a gradient, Jacobian, or higher-order tensor. Matrix calculus becomes confusing when notation hides whether vectors are rows or columns, which variable is differentiated, and whether a derivative acts forward or backward. A reliable approach starts with shapes and differentials. The final array formula should be the last step, not the first.

Treating derivatives as operators rather than as stored arrays pays off, and the payoff has a date on it. Multiplying a stored Hessian by a vector “takes O(n2) time when there are n weights”. Barak Pearlmutter's R-operator technique instead “finds this product in O(n) time and space”, and it makes no approximations: the product is exact, “which takes about as much computation, and is about as local, as a gradient evaluation”. Neural Computation published that in 1994. The 2018 JMLR autodiff survey reports the same result for the reverse-on-forward construction, which “computes Hv with O(n) complexity, even though H is a n × n matrix.” The same n² that makes storage impossible is the n² the operator never touches.

Different books arrange Jacobians differently, and the books themselves say so. Parr and Howard declare their choice partway through their 2018 deep-learning matrix-calculus tutorial: “We are using the so-called numerator layout but many papers and software will use the denominator layout.” The other convention, they add, is “just transpose of the numerator layout Jacobian (flip it around its diagonal)”.

The disagreement is not folklore. It is visible in the printed identities. Minka's 2000 statistics notes declare their own rule — partials with respect to the numerator laid out by the shape of Y, partials with respect to the denominator by the transpose of X — and on that rule his equation 14 gives d/dX tr(AXB) = BA. The Matrix Cookbook prints ∂Tr(AXB)/∂X = AᵀBᵀ as its equation 101. Same derivative, transposed answer. The mathematics can agree while the displayed arrays are transposed. Declare the convention before you copy the formula.

A familiar derivative illustrates why the composed loss is simpler than separate Jacobians suggest. Track the shape at every line and the cancellation appears on its own.

Most frameworks avoid constructing full Jacobians by applying them to vectors, and the cost model behind that choice is published rather than assumed. The 2018 JMLR autodiff survey prices the m × n Jacobian at n·c·ops(f) in forward mode and m·c·ops(f) in reverse mode, “where c is a constant guaranteed to be c < 6 and typically c ∼ [2, 3] (Griewank and Walther, 2008)”. A single direction costs far less than either: “Thus, we can compute the Jacobian–vector product in just one forward pass.” JAX's own documentation puts the same constant at about 3, describing a JVP as “about 3x the cost of just evaluating the function” and a VJP as “only about three times the cost of evaluating f”.

When a full Jacobian really is wanted, the direction is a shape decision, and both frameworks document the same threshold. JAX's Autodiff Cookbook states it directly: “jacfwd uses forward-mode automatic differentiation, which is more efficient for “tall” Jacobian matrices (more outputs than inputs), while jacrev uses reverse-mode, which is more efficient for “wide” Jacobian matrices (more inputs than outputs).” PyTorch's functorch tutorial gives the rule of thumb with letters attached: for f: R^N → R^M, prefer jacfwd when M > N and jacrev otherwise. Reverse mode builds the Jacobian row by row, forward mode column by column.

The HVP is not a matrix in either framework. It is a composition of the other two products. JAX's Autodiff Cookbook builds hvp from grad-of-grad and then from forward-over-reverse, jvp(grad(f), x, v). PyTorch's functorch tutorial defines hvp(f, primals, tangents) as jvp(grad(f), primals, tangents)[1], noting that “it turns out we don't need to materialize the full Hessian to do this”. The MIT 18.S096 notes give the same forward-over-reverse construction and say it can “incur computational cost proportional only to that of a single function evaluation”. The reason to compose rather than store is stated without hedging: “The trick is not to instantiate the full Hessian matrix: if n is large, perhaps in the millions or billions in the context of neural networks, then that might be impossible to store.”

A table of derivative identities can wait until after these six steps. It should also wait until you have checked the table's layout convention against your own, because Minka's equation 14 and the Cookbook's equation 101 print transposes of the same derivative. The steps below are the route Edelman and Johnson take to ∇f = (A + Aᵀ)x: declare shapes, choose the variable, form the differential, collect the perturbation, read the gradient, and verify it numerically.