Jacobians and Hessians for Optimization

Hello! Let's continue our journey through the mathematical foundations of AI.

In the last lesson, we explored backpropagation as a systematic application of the chain rule. We saw that it allows us to calculate the gradient ($\nabla L$) of a complex loss function with respect to any parameter in a neural network. We also briefly touched on the idea that for vector-valued functions, the derivatives are no longer simple vectors but become matrices called Jacobians.

Today, we will formalize that concept and go one step further. The learning outcome for this lesson is to calculate Jacobians and Hessians for advanced optimization. We will define what these matrices are, see how to compute them, and, most importantly, understand why they are crucial for optimization methods that are more sophisticated than standard gradient descent.

1. The Jacobian: Generalizing Derivatives for Vector Functions

So far, we have focused on gradients, which are for scalar-valued functions (functions that output a single number, like a loss function $L(\mathbf{w})$). But what if a function outputs a vector?

Consider a function $\mathbf{f}$ that takes an $n$-dimensional vector $\mathbf{x}$ as input and produces an $m$-dimensional vector $\mathbf{y}$ as output.
$\mathbf{y}=\mathbf{f}(\mathbf{x})\text{where }\mathbf{x}\in {\mathbb{R}}^n,\mathbf{y}\in {\mathbb{R}}^m$
How do we describe the rate of change here? We need to capture how each output component changes with respect to each input component. We organize these partial derivatives into an $m\times n$ matrix called the Jacobian matrix, denoted as $\mathbf{J}$.

The element at row $i$ and column $j$ of the Jacobian is $\frac{\partial y_i}{\partial x_j}$.

$$\mathbf{J}=\left(\begin{array}{cccc}\frac{\partial y_1}{\partial x_1}& \frac{\partial y_1}{\partial x_2}& \cdots & \frac{\partial y_1}{\partial x_n}\\ \frac{\partial y_2}{\partial x_1}& \frac{\partial y_2}{\partial x_2}& \cdots & \frac{\partial y_2}{\partial x_n}\\ ⋮& ⋮& \ddots & ⋮\\ \frac{\partial y_m}{\partial x_1}& \frac{\partial y_m}{\partial x_2}& \cdots & \frac{\partial y_m}{\partial x_n}\end{array}\right)$$

You can think of the gradient $\nabla f$ as a special case of the Jacobian where the output dimension $m=1$.

Calculation Example

Let's compute the Jacobian for a function $\mathbf{f}:{\mathbb{R}}^2\to {\mathbb{R}}^2$ defined as:
$\mathbf{f}(x_1,x_2)=\left[\begin{array}{c}{f}_1(x_1,x_2)\\ f_2(x_1,x_2)\end{array}\right]=\left[\begin{array}{c}{x}_1^2\sin (x_2)\\ x_1{x}_2\end{array}\right]$
We need to find the four partial derivatives:

• $\frac{\partial f_1}{\partial x_1}=2x_1\sin (x_2)$
• $\frac{\partial f_1}{\partial x_2}=x_1^2\cos (x_2)$
• $\frac{\partial f_2}{\partial x_1}=x_2$
• $\frac{\partial f_2}{\partial x_2}=x_1$

Arranging them into a 2x2 matrix gives us the Jacobian:

$${\mathbf{J}}_{\mathbf{f}}(x_1,x_2)=\left(\begin{array}{cc}2x_1\sin (x_2)& x_1^2\cos (x_2)\\ x_2& x_1\end{array}\right)$$

The Jacobian and the Chain Rule in Neural Networks

The true power of the Jacobian becomes apparent when we revisit the chain rule from our last lesson. If we have a sequence of vector functions, say $\mathbf{z}=g(\mathbf{y})$ and $\mathbf{y}=f(\mathbf{x})$, the chain rule for finding the derivative of $\mathbf{z}$ with respect to $\mathbf{x}$ is simply the product of their Jacobian matrices:
$\frac{\partial \mathbf{z}}{\partial \mathbf{x}}=\frac{\partial \mathbf{z}}{\partial \mathbf{y}}\frac{\partial \mathbf{y}}{\partial \mathbf{x}}⟹{\mathbf{J}}_{\mathbf{g}\circ \mathbf{f}}={\mathbf{J}}_{\mathbf{g}}\cdot {\mathbf{J}}_{\mathbf{f}}$
This is exactly what happens during backpropagation in a neural network. Each layer is a vector function, and propagating the gradient backward involves multiplying by the Jacobian of each layer.

2. The Hessian: Capturing Curvature with Second Derivatives

While the gradient and Jacobian tell us about the slope (first-order information), the Hessian matrix tells us about the curvature (second-order information) of a function. For a scalar-valued function $f:{\mathbb{R}}^n\to \mathbb{R}$, the Hessian $\mathbf{H}$ is the $n\times n$ matrix of all its second-order partial derivatives.

It can also be defined as the Jacobian of the gradient: $\mathbf{H}(f)=\mathbf{J}(\nabla f)$.

The element at row $i$ and column $j$ of the Hessian is $\frac{{\partial }^{2}f}{\partial x_i\partial x_j}$.

$$\mathbf{H}=\left(\begin{array}{cccc}\frac{{\partial }^{2}f}{\partial x_1^2}& \frac{{\partial }^{2}f}{\partial x_1\partial x_2}& \cdots & \frac{{\partial }^{2}f}{\partial x_1\partial x_n}\\ \frac{{\partial }^{2}f}{\partial x_2\partial x_1}& \frac{{\partial }^{2}f}{\partial x_2^2}& \cdots & \frac{{\partial }^{2}f}{\partial x_2\partial x_n}\\ ⋮& ⋮& \ddots & ⋮\\ \frac{{\partial }^{2}f}{\partial x_n\partial x_1}& \frac{{\partial }^{2}f}{\partial x_n^2}& \cdots & \frac{{\partial }^{2}f}{\partial x_n^2}\end{array}\right)$$

A key property, due to Clairaut's theorem, is that if the second partial derivatives are continuous, the order of differentiation doesn't matter (e.g., $\frac{{\partial }^{2}f}{\partial x_1\partial x_2}=\frac{{\partial }^{2}f}{\partial x_2\partial x_1}$). This means the Hessian matrix is symmetric.

Calculation and Implementation in Python

Given your background in Python, let's see how this is done in code. Manually calculating these derivatives is tedious and error-prone. We can use a symbolic mathematics library like SymPy to do the heavy lifting.

3. Application: Advanced Optimization

Why do we care about curvature? Because it allows for more powerful optimization algorithms. Gradient descent only uses the direction of the slope; it's like walking downhill in a thick fog. Second-order methods use curvature information to take more intelligent steps, like using a topographic map to find the valley floor.

Newton's Method

The most famous second-order optimization algorithm is Newton's method. It works by creating a quadratic approximation of the function at the current point using the second-order Taylor expansion:
$f(\mathbf{w})\approx f({\mathbf{w}}_k)+\nabla f({\mathbf{w}}_k{)}^T(\mathbf{w}-{\mathbf{w}}_k)+\frac{1}{2}(\mathbf{w}-{\mathbf{w}}_k{)}^T\mathbf{H}({\mathbf{w}}_k)(\mathbf{w}-{\mathbf{w}}_k)$
To find the minimum of this approximation, we set its gradient to zero and solve for $\mathbf{w}$. This gives the Newton's method update rule:
${\mathbf{w}}_{k+1}={\mathbf{w}}_k-{\mathbf{H}}_k^{-1}\nabla f({\mathbf{w}}_k)$
Notice the difference from gradient descent (${\mathbf{w}}_{k+1}={\mathbf{w}}_k-\alpha \nabla f({\mathbf{w}}_k)$). Instead of a scalar learning rate $\alpha$, Newton's method uses the inverse of the Hessian matrix, ${\mathbf{H}}^{-1}$, to automatically determine the optimal step size and direction.

Classifying Critical Points

The Hessian also helps us classify critical points (where $\nabla f=0$) found by an optimization algorithm. By examining the eigenvalues of the Hessian at a critical point (which connects to our earlier lesson on eigenvalues):

• All eigenvalues are positive: The Hessian is positive definite. The point is a local minimum.
• All eigenvalues are negative: The Hessian is negative definite. The point is a local maximum.
• Eigenvalues have mixed signs: The Hessian is indefinite. The point is a saddle point.

Saddle points are a major challenge in high-dimensional optimization (like training neural networks), and second-order methods are better equipped to handle them.

Test your understanding!

Consider the function $f(x,y)=x^2+3xy+2y^2$. Calculate its gradient and Hessian. Is the point (0,0) a local minimum, maximum, or saddle point?

Show answer

1. Gradient:
   $\nabla f=\left[\begin{array}{c}\frac{\partial f}{\partial x}\\ \frac{\partial f}{\partial y}\end{array}\right]=\left[\begin{array}{c}2x+3y\\ 3x+4y\end{array}\right]$
   At (0,0), the gradient is $[0,0{]}^T$, so it is a critical point.

2. Hessian:
   $\mathbf{H}=\left(\begin{array}{cc}\frac{{\partial }^{2}f}{\partial x^2}& \frac{{\partial }^{2}f}{\partial x\partial y}\\ \frac{{\partial }^{2}f}{\partial y\partial x}& \frac{{\partial }^{2}f}{\partial y^2}\end{array}\right)=\left(\begin{array}{cc}2& 3\\ 3& 4\end{array}\right)$
   The Hessian is constant for this function.

3. Classify the point: We check the definiteness of the Hessian. We can find its eigenvalues by solving $\det (\mathbf{H}-\lambda \mathbf{I})=0$:
   $(2-\lambda )(4-\lambda )-9=0$
   ${\lambda }^2-6\lambda +8-9=0⟹{\lambda }^2-6\lambda -1=0$
   The roots are $\lambda =\frac{6\pm \sqrt{36-4(-1)}}{2}=3\pm \sqrt{10}$.
   One eigenvalue is $3+\sqrt{10}>0$ and the other is $3-\sqrt{10}<0$. Since the eigenvalues have mixed signs, the Hessian is indefinite, and (0,0) is a saddle point.

4. Practicality: Approximating the Hessian

While powerful, Newton's method has a major drawback for deep learning:

1. Computation: Calculating the Hessian requires computing $O(N^2)$ second derivatives, where $N$ is the number of parameters. For a model like GPT-3 with 175 billion parameters, this is impossible.
2. Storage: Storing the $N\times N$ Hessian matrix is also impossible.
3. Inversion: Inverting the Hessian is an $O(N^3)$ operation, which is prohibitively expensive.

For these reasons, pure Newton's method is not used for training large neural networks. Instead, we use quasi-Newton methods. These methods approximate the Hessian (or its inverse) using more easily available information, like the gradients from previous steps.

A very common and clever approximation in non-linear least squares problems is the Gauss-Newton algorithm, which approximates the Hessian using the Jacobian:
$\mathbf{H}\approx {\mathbf{J}}^T\mathbf{J}$
This is a beautiful connection between our two topics for today. We use the matrix of first derivatives (Jacobian) to approximate the matrix of second derivatives (Hessian), avoiding the expensive computation of the true second derivatives.

Conclusion

Today, we have moved beyond the gradient to explore higher-order derivatives and their role in optimization.

Key Takeaways:

• The Jacobian matrix generalizes the derivative to vector-valued functions, representing all partial derivatives of each output with respect to each input. It is fundamental to the chain rule in backpropagation for vector operations.
• The Hessian matrix is the matrix of second-order partial derivatives of a scalar-valued function. It describes the function's local curvature.
• The Hessian is critical for advanced, second-order optimization methods like Newton's method, which uses curvature information to take more direct steps toward a minimum.
• The properties of the Hessian (its definiteness, checked via eigenvalues) allow us to classify critical points as minima, maxima, or saddle points.
• In practice, computing the full Hessian is often too expensive. Therefore, quasi-Newton methods approximate the Hessian, sometimes using the Jacobian (e.g., $\mathbf{H}\approx {\mathbf{J}}^T\mathbf{J}$).

Preview of the next lesson:
We have now covered the essential calculus needed for understanding how AI models learn and are optimized. We will now shift gears to the second major pillar of machine learning mathematics: probability and statistics. In the next lesson, we will start with one of the most foundational theorems in this field: Bayes' theorem, a principle that underpins our ability to reason about and update our beliefs in the face of new evidence.