Efficient AI · Lecture notes

Pruning a model

Use the local loss landscape at a converged point — via a quadratic (Hessian) approximation — to choose which weight to remove and how to update the rest.

Setup

$F$ — the model
$\vec{w}$ — flattened parameter vector
$X$ — input data representative of the input distribution

Q. Again, what should we be looking at while pruning?

The only “information” we get about good and bad decisions in ML is through the loss landscape.

$\rightarrow$ Training depends on the loss landscape.

Q. Which model should we start at?

Since the information you can get about the loss landscape is local, you better start from the converged model.

Far from the minimum, local geometry is
not useful for pruning decisions.
Near the converged point, it is.

Local quadratic approximation

Let $\mathcal{L}(w, X)$ be the loss and $w^\star$ the weights at a minimum (converged model). To get analytic information about the landscape we approximate it locally with a Taylor expansion:

\[ \mathcal{L}(w, X) = \mathcal{L}(w^\star, X) + \nabla\mathcal{L}(w^\star, X)^\top \delta w + \tfrac{1}{2}\, \delta w^\top \nabla^2\mathcal{L}(w^\star, X)\, \delta w + O(\|\delta w\|_2^3), \]

where $\delta w = w - w^\star$.

Converged: $\nabla\mathcal{L}(w^\star, X) = 0$, so the linear term vanishes.
Local / quadratic: for small $\delta w$, the $O(\|\delta w\|_2^3)$ remainder is negligible.

Writing $H(w^\star, X) = \nabla^2\mathcal{L}(w^\star, X)$ for the Hessian,

\[ E(\delta w) \;:=\; \mathcal{L}(w, X) - \mathcal{L}(w^\star, X) \;=\; \tfrac{1}{2}\, \delta w^\top H(w^\star, X)\, \delta w . \]

This is just a quadratic function and is therefore easy to analyse.

Which $i$ to prune?

Strategy:

Pick the coordinate with lowest saliency
Update the rest of the weights to compensate

Pruning coordinate $i$: $w \;\mapsto\; \bigl(\hat{w}_{0:i-1},\; 0,\; \hat{w}_{i+1:}\bigr)$.

$N$ subproblems

For each candidate index $i$, find the update that zeros weight $i$ while minimizing the quadratic error:

\[ \delta w = \underset{\delta w_i = -w_i}{\arg\min} \; E(\delta w) \qquad \text{(so the $i$-th weight goes to $0$)}. \]

Lagrange multiplier method

Enforce the constraint with a multiplier $\lambda$ and unit vector $e_i$:

\[ \delta w = \underset{\delta w}{\arg\min} \; \underbrace{ T(\delta w, \lambda) }_{ E(\delta w) + \lambda\, e_i^\top (w + \delta w) } , \qquad E(\delta w) \approx \tfrac{1}{2}\,\delta w^\top H\,\delta w . \]

Stationarity conditions:

\[ \frac{\partial T}{\partial \lambda} = 0 \;\Longrightarrow\; \delta w_i = -w_i \]

\[ \frac{\partial T}{\partial (\delta w)} = 0 \;\Longrightarrow\; H\,\delta w + \lambda\, e_i = 0 . \]

(For matrix identities, see the Matrix Cookbook.)

Solving for $\lambda$ and $\delta w$

The $i$-th entry of $\delta w$ is known ($\delta w_i = -w_i$). From $H\,\delta w + \lambda e_i = 0$,

\[ \delta w + \lambda\, H^{-1} e_i = 0 . \]

$H^{-1} e_i$ is the $i$-th column of $H^{-1}$. Matching the $i$-th coordinate gives

\[ -w_i = \lambda\, [H^{-1}]_{ii} \qquad\Longrightarrow\qquad \lambda = \frac{-w_i}{[H^{-1}]_{ii}} . \]

Hence the optimal update

\[ \delta w = -\lambda\, H^{-1} e_i = \frac{w_i}{[H^{-1}]_{ii}}\, H^{-1} e_i . \]

The prefactor $\tfrac{w_i}{[H^{-1}]_{ii}}$ is a scalar; $\delta w$ is proportional to the $i$-th column of $H^{-1}$.

Saliency

The cost of pruning $i$ is the quadratic error under this optimal update:

\[ \text{Saliency} = \tfrac{1}{2}\, \delta w^\top H\, \delta w . \]

Substitute $\delta w = \tfrac{w_i}{[H^{-1}]_{ii}}\, H^{-1} e_i$:

\[ = \tfrac{1}{2} \frac{w_i^2}{\bigl([H^{-1}]_{ii}\bigr)^2} \bigl(H^{-1} e_i\bigr)^\top H\, H^{-1} e_i \]

\[ = \tfrac{1}{2} \frac{w_i^2}{\bigl([H^{-1}]_{ii}\bigr)^2} \, e_i^\top (H^{-1})^\top e_i = \tfrac{1}{2} \frac{w_i^2}{\bigl([H^{-1}]_{ii}\bigr)^2} \, [H^{-1}]_{ii} \]

\[ \text{Saliency} = \tfrac{1}{2}\, \frac{w_i^2}{[H^{-1}]_{ii}} \,, \qquad \delta w = \frac{w_i}{[H^{-1}]_{ii}}\, H^{-1} e_i . \]

Algorithm

OBS-style prune

Train the full model to convergence.
For each $i$, compute saliency $S_i = \tfrac{1}{2}\, w_i^2 / [H^{-1}]_{ii}$.
- Pick $i$ with least saliency.
- Update $w \leftarrow w + \delta w$.
- If $\min_i S_i$ is large compared to the loss $\rightarrow$ retrain (return to step 1).
- Otherwise redo the local approximation and prune again (loop within step 2).

If $\{S_i\}$ are small, continue pruning under the current local approximation; if the best saliency is large, retrain.

Complexity

One prune

Dominated by forming / inverting the Hessian: $O(N^3)$, where $N$ is the number of parameters.

Exercises

1. Analyse the update if we remove multiple weights at a time. Suppose we zero out the block of coordinates $i{:}i{+}j$ (that is, set $w_i,\ldots,w_{i+j}$ to zero in one step). What should the optimal update $\delta w$ be under the same quadratic approximation $E(\delta w)=\tfrac12\delta w^\top H\,\delta w$? How does the Lagrange / constraint formulation change compared to pruning a single index?

2. Analyse the update if, instead of removing weight $i$ (setting it to zero), we set $w_i$ to a specific value $c$ (i.e. the constraint is $(w + \delta w)_i = c$, or equivalently $\delta w_i = c - w_i$). What is the optimal $\delta w$ in this case? How does saliency change?