Efficient AI · Lecture notes

Hessian complexity

How expensive is the Hessian of a network? When is the diagonal enough? And for MSE loss, how do we approximate $H$ and $H^{-1}$ with rank-1 updates?

Four parts

Full Hessian computation
Hessian diagonal computation
Approximation of $H$ for MSE loss
Hessian inverse for MSE loss

1. Computation of the Hessian

Q. How much time does it take to compute the Hessian of a network?

Assume a computational graph with $n$ nodes and parameters also scale $O(n)$.

Forward pass: $O(n)$
Backward pass (first derivatives): $O(n)$
Hessian $H = \nabla^2 L$ is an $n \times n$ matrix — $O(n^2)$ entries

Forward pass
Evaluate $F(x)$ once along the graph — $O(n)$.

Backward pass
Same graph, reverse sweep — yields all first derivatives $\partial F/\partial z_i$ in $O(n)$.

“Backward–backward”
Run a forward pass on the backward graph — for each node of interest — to get mixed derivatives such as $\partial^2 F/(\partial z_1\,\partial x)$.

\[ O(n)\ \text{nodes} \times O(n)\ \text{per reverse pass} \;=\; O(n\times n) \;=\; O(n^2). \]

Over $N$ data samples the full Hessian costs $O(N n^2)$ — already too large for modern models.

Takeaway Storing and forming the dense Hessian is $O(n^2)$ (and $O(N n^2)$ with data). Exact OBS-style methods that need the full $H$ (or $H^{-1}$) cannot scale naively.

2. Hessian diagonal.

Q. What if we only care about $\partial^2 L / \partial w_i^2$ (the diagonal of $H$)?

This is exactly what Optimal Brain Damage (OBD) needs: a diagonal Hessian approximation.

$\vec{x}$ and $\vec{w}$ feed $f$; output $\vec{z}=f(\vec{x};\vec{w})$ then enters the loss $L(\vec{z})$.

Diagonal of a weight $w_{ij}$

For a scalar edge $x_i \xrightarrow{w_{ij}} z_j$ with local map $f$,

\[ \frac{\partial L}{\partial w_{ij}} = \frac{\partial L}{\partial z_j}\, \frac{\partial z_j}{\partial w_{ij}} = \frac{\partial L}{\partial z_j}\, f'(w_{ij}). \]

\[ \frac{\partial^2 L}{\partial w_{ij}^2} = \bigl(f'(w_{ij})\bigr)^2\, \frac{\partial^2 L}{\partial z_j^2} + f''(w_{ij})\, \frac{\partial L}{\partial z_j} . \]

Hessian chain rule

Let $\vec{z} = f(\vec{x})$ and loss $L(\vec{z})$. Then

\[ \nabla^2_x L = J_f^\top\,(\nabla^2_z L)\, J_f + \sum_{j=1}^{m} \nabla^2_x z_j \frac{\partial L}{\partial z_j}. \]

First term — global curvature pushed through the local Jacobian $J_f$
$\partial L/\partial z_i$ — global gradient; $\nabla^2_x z_i$ — local curvature
We only need local Hessians to be computed along the graph

Dimensions: $\nabla^2_x L$ is $d\times d$, $J_f$ is $d'\times d$, $\nabla^2_z L$ is $d'\times d'$.

Complexity All diagonal entries of the Hessian over the parameters can be obtained in a single “non-traditional” backward pass — cost $O(n)$, same order as a gradient.

$\Rightarrow$ OBD is feasible: one forward + one specialized backward.

3. Approximation of the Hessian for MSE loss

Specialize to squared error. Let $o = F(\vec{w},\vec{x})$ and target $t$, with

\[ E = \frac{1}{2N} \sum_{i=1}^{N} \bigl(t_i - o_i\bigr)^2, \qquad o_i = F(\vec{w},\vec{x}_i). \]

Gradient ($n\times 1$ if $\vec{w}\in\mathbb{R}^n$):

\[ \frac{\partial E}{\partial w} = \frac{1}{N} \sum_{i=1}^{N} (t_i - o_i)\, \Bigl(-\frac{\partial F}{\partial w}\Bigr) = -\frac{1}{N} \sum_{i=1}^{N} (t_i - o_i)\, \nabla_w F_i. \]

Exact Hessian ($n\times n$), by the product rule:

\[ \frac{\partial^2 E}{\partial w^2} = \frac{1}{N} \sum_{i=1}^{N} \Biggl[ \nabla_w F_i\,(\nabla_w F_i)^\top + (o_i - t_i)\, \nabla_w^2 F_i \Biggr]. \]

Approximation. If the model fits well, $t_i \approx o_i$, so the residual term vanishes:

\[ H \;\approx\; \frac{1}{N} \sum_{i=1}^{N} \nabla_w F_i\,(\nabla_w F_i)^\top \;=\; \frac{1}{N} \sum_{i=1}^{N} X_i X_i^\top. \]

$X_i = \nabla_w F(\vec{w},\vec{x}_i)$ is $n\times 1$; each outer product is $n\times n$
$H$ is the (scaled) sample covariance of the output gradients
Cost to form this approximate Hessian: $O(N n^2)$
Note: this holds for MSE; exercise — what changes for cross-entropy?

4. Hessian inverse for MSE loss

A dense inverse from scratch is generally $O(n^3)$. The MSE approximation $H = \tfrac{1}{N}\sum_i X_i X_i^\top$ is a sum of rank-1 matrices, so we can maintain $H^{-1}$ with rank-1 updates.

Sherman–Morrison (rank-1 update)

For invertible $A$ and vectors $u,v$ with $1 + v^\top A^{-1}u \neq 0$:

\[ (A + uv^\top)^{-1} = A^{-1} - \frac{A^{-1}uv^\top A^{-1}}{1 + v^\top A^{-1}u}. \]

Iterative construction

Write $H_k = H_{k-1} + \tfrac{1}{N} X^{(k)} X^{(k)\top}$ with $X^{(k)} = \nabla_w F(\vec{w},\mathrm{in}^{(k)})$. Then

\[ a = 1 + \tfrac{1}{N}\, X^{(k)\top} H_{k-1}^{-1} X^{(k)}, \]

\[ H_k^{-1} = H_{k-1}^{-1} - \frac{ H_{k-1}^{-1} X^{(k)} X^{(k)\top} H_{k-1}^{-1} }{ N\, a }. \]

Equivalently, with $u = v = (\nabla F)_k$ absorbed into the same rank-1 form:

\[ H_k^{-1} = H_{k-1}^{-1} - \frac{ H_{k-1}^{-1} u\, v^\top H_{k-1}^{-1} }{ 1 + v^\top H_{k-1}^{-1} u }. \]

Initialization and regularization

Start from $H_0 = \lambda I$ with $\lambda$ small (so $H_0^{-1}$ exists)
For small $\lambda$, the final inverse is insensitive
$\lambda \neq 0$ is equivalent to $\ell_2$ weight decay on the loss:

\[ E = \frac{1}{2N} \sum_i \bigl(F(\vec{x}_i,\vec{w}) - t_i\bigr)^2 + \frac{\lambda}{2}\,\|\vec{w}\|_2^2. \]

Complexity Updating $H^{-1}$ over $N$ samples costs $O(N n^2)$ (one outer-product-scale update per sample), avoiding a separate $O(n^3)$ factorization. Contrast with forming $H$ then inverting from scratch: $O(N n^2 + n^3)$.