Hessian of a Quadratic

Read the second derivatives

A gradient tells you how a scalar function changes with each input. Its Hessian tells you how those derivatives change. For a vector input with d entries, the Hessian of a scalar function has shape (d, d). Entry (i, j) is the derivative of gradient entry i with respect to input entry j.

import jax
import jax.numpy as jnp

def f(x):
    return x[0]**2 + x[0]*x[1] + 1.5*x[1]**2

H = jax.hessian(f)(jnp.array([3.0, 4.0]))
# [[2.0, 1.0],
#  [1.0, 3.0]]

The gradient is [2*x[0] + x[1], x[0] + 3*x[1]]. Differentiating each entry gives the two rows above. The diagonal entries describe curvature along individual coordinates; the off-diagonal entries describe how the coordinates affect one another's derivatives.

Why symmetry matters for this quadratic

The example can be written as 0.5 * x @ A @ x, with A = [[2, 1], [1, 3]]. For a real symmetric A, the gradient is A @ x, so the Hessian is A at every x.

The factor 0.5 matters: removing it doubles the Hessian. If A were not symmetric, the Hessian of 0.5 * x @ A @ x would instead be (A + A.T)/2. This exercise guarantees symmetry.

What jax.hessian returns

Pass a function to jax.hessian, then call the resulting function at an input. For this scalar loss it returns the full matrix. Unlike jax.grad, jax.hessian also supports array-valued outputs. For an output with shape (m,) and an input with shape (d,), its result has shape (m, d, d): one matrix of second derivatives per output.

A dense Hessian for d scalar parameters contains d*d entries. Two nested differentiation transforms do not make its cost independent of d. If you only need the product with a vector, use an HVP method to avoid constructing the full matrix.

Compute derivatives by following the operations

Automatic differentiation computes derivatives by applying the chain rule to a program's operations. It does not estimate them by evaluating nearby inputs, as finite differences do, or rearrange the whole expression using symbolic algebra.

For f(g(x)), the chain rule combines the derivative of g with the derivative of f at g(x). Repeating this process lets a framework differentiate a long sequence of operations.

Forward mode follows an input change

For a function with n inputs and m outputs, the Jacobian J is an (m, n) matrix of derivatives. Entry (i, j) describes how output i changes with input j.

Forward-mode differentiation computes Jv: the output's rate of change when the input moves in direction v. For example, if f(x) = sum(x_i^2), then at x = [1, 2] the gradient is [2, 4]. In direction v = [3, 0], the rate of change is 2 * 3 + 4 * 0 = 6.

This is a directional derivative. It gives one Jacobian column only when v selects one input coordinate, such as [1, 0].

Reverse mode works backward from an output

Reverse mode computes J^T u, where u weights the outputs. For a scalar loss, choosing u = 1 gives the gradient with respect to all inputs. This makes reverse mode useful for training: a model can have many parameters but just one loss.

Forward mode can be useful when you need sensitivity to only a few input directions. Neither mode needs to construct the full Jacobian for these products. That can save substantial memory. Higher-order differentiation can similarly compute a Hessian-vector product without storing the full matrix of second derivatives.

Changing gradient flow changes what is differentiated

A stop-gradient operation keeps a value in the forward calculation but treats it as constant during differentiation. A custom derivative supplies a derivative rule explicitly. It should match the intended function unless the exercise asks for a surrogate, such as an approximation used during training.