Efficient AI

Pruning / Sparsity

Aditya Desai

Many thanks to all the resources listed in References on the main website. Almost all figures are taken from these resources; some were generated with Gemini. Some images are taken from Google Images and still need proper citation (in progress).

Size of AI is huge and it is only growing!

Model Sizes Over Time

Deployment targets have tight memory budgets

Device Class Approximate RAM Range
Tiny IoT devices KB to MB
Wearables 0.5–8 GB
Phones & Glasses 4–16 GB
Laptops 8–64 GB
Drones 2–32 GB
Cars 16–512 GB
Edge Servers 64 GB–4 TB

Energy cost of compute vs memory

Energy table for 45nm CMOS: DRAM access is far more expensive than arithmetic
  • Memory access (DRAM) is ~3 orders of magnitude more energy-expensive than simple arithmetic.
  • Compressing models reduces data movement as well as FLOPs — often the bigger win for energy.

what is model compression?

  • Given a trained model, can we reduce its size and computation while maintaining its accuracy?

why model compression?

  • saves memory, compute, energy, money, etc
  • Unlocks new deployment scenarios. e.g. deploying on edge devices -- PCs/laptops/phones/cars/drones/etc

Techniques are transferrable..

KV cache Training memory

We can apply the techniques we will discuss as part of model compression to activations, gradients, optimizer states, etc.

Techniques are transferrable..

Download
  • Pruning is just one structured lens.
  • We can think of other forms of low-dimensional spaces or quantizations of space.
  • The spaces onto which we project are different, but the principles are the same.

Pruning / Sparsity

  • what is pruning?
  • Pruning a neuron
  • Optimal Brain Surgeon
  • Optimal Brain Compression
  • Iterative Magnitude Pruning
  • Pruning in Pre-LLM era
  • Structured Pruning & Hardware Supported Sparsity
  • Sensitivity Analysis for Pruning
  • Lottery Ticket Hypothesis? (If we have time)
  • Pruning LLMs

Two useful identities for inverse of a matrix

Inverse of a rank one update

1) Rank-one update — add outer product $uv^\top$ to $A$:

Relation between: \[ (A + uv^\top)^{-1} \textrm{ and } A^{-1} \]
Sherman Morrison Identity (rank 1 update)
\[ (A + uv^T)^{-1} = A^{-1} - \frac{A^{-1}uv^T A^{-1}}{1 + v^T A^{-1} u} \]

If you do a one-rank update to A, then it provides a relation between inverse of A and the inverse of the updated matrix. O(n2) update instead of O(n3) for inverse of (A+uv^\top)

Woodbury Identity (general rank update)
\[ (A + UCV)^{-1} = A^{-1} - A^{-1}U(C^{-1} + V A^{-1}U)^{-1}V A^{-1} \]

Inverse of a submatrix extraction

\[ A = \begin{pmatrix} a_{11} & a_{12} & a_{13} \\ a_{21} & {\color{#0b6e4f}a_{22}} & {\color{#0b6e4f}a_{23}} \\ a_{31} & {\color{#0b6e4f}a_{32}} & {\color{#0b6e4f}a_{33}} \end{pmatrix} \;\longrightarrow\; B = A[2{:},2{:}] = \begin{pmatrix} {\color{#0b6e4f}a_{22}} & {\color{#0b6e4f}a_{23}} \\ {\color{#0b6e4f}a_{32}} & {\color{#0b6e4f}a_{33}} \end{pmatrix} \] Relation between $A^{-1}$ and $B^{-1}$
Inverse of a submatrix extraction
\[ B^{-1} = \left[A^{-1} - (\frac{1}{A^{-1}_{11}} A^{-1}_{:,1} A^{-1}_{:,1}) \right]_{2:,2:} \]

If you have inverse of a larger matrix, then you can use this to find the inverse of matrix with one row and column removed. O(n2) update instead of O(n3) for inverse of B

Pruning: Sparsity in weights

Snipping a model connections Computational graph view: we remove some of the connections in the model.
Fine-grained (irregular) pruning pattern
weight view: weights are sparsified
Projection View: Pruning is a projection onto co-ordinate aligned subspaces.
Projection View: Pruning is a projection onto co-ordinate aligned subspaces.
  • Removing enough connections can naturally remove some of the neurons and structural components in the model.
Unstructured vs structured pruning Pruning patterns from irregular to regular

Pruning

  • Pruning has a history since 1980s and post deep learning era it is one of the most widely researched techniques
Number of papers on Neural Network Compression vs Neural Network Pruning, 1989–2023

Pruning a neuron

Note Link: notes/pruning_a_neuron.html

what is the complexity of this computation

  • x, w \in R^d, y \in R
  • Single neuron pruning:
  • d × O(d3)
    • d — number of sub problems
    • O(d3) — inverse of the Hessian
    • If pruning k neurons, then the complexity is O(k d4)

A single neuron pruning

  • Pruning a neuron itself is quite expensive (disclaimer: this is not the best algo)
  • Imagine decent sized model pre-LLM (Convolution / MLP etc)
  • Imagine a LLM model!

Pruning a “Model"

(Optimal Brain Surgeon(OBS), Hassibi and Stork, 1993)

Note Link: notes/pruning_a_model.html

All that one needs to know for pruning.

Algorithm for OBS

  1. Train a “reasonably large” network to minimum error.
  2. Compute H−1.
  3. Find the q that gives the smallest saliency Lq = wq2 / (2 [H−1]qq). If this candidate error increase is much smaller than E, delete the q-th weight and go to step 4; otherwise go to step 5.
  4. Use the q from step 3 to update all weights (δw). Go to step 2.
  5. No more weights can be deleted without a large increase in E. (At this point it may be desirable to retrain the network.)

Complexity of OBS?

  • A single step: Dominated by forming / inverting the Hessian: O(N3), where N is the number of parameters.
  • Hessian computation is O(nN2) where n is the number of samples and N is the number of parameters.
  • k step OBS : O(k(N3 + nN2))

How can we reduce the complexity of OBS?

Hint: Use Linear Algebra identities we introduced. (Exercise : We will look at it next class)

How can we simplify

\[ Saliency(q) = L_q = \frac{w_q^2}{2 H^{-1}_{qq}} \] \[ Update(w) = \frac{w_i}{H^{-1}_{qq}} H^{-1} e_i \]

Optimal Brain Damage (OBD)?

assumptions
  • The model is trained to convergence.
  • The model is a quadratic local approximation.
  • Hessian is a diagonal matrix.

Algorithm for OBD

  1. Train a “reasonably large” network to minimum error.
  2. Compute Hessian H
  3. Find the q that gives the smallest saliency Lq = ½ Hqq wq2. Delete few low saliency weights
  4. Update all weights (δw).
  5. Retrain i.e. go to step 1.

Complexity of OBD?

  • We only need diagonal entries of Hessian.

Computing diagonal entries of Hessian

Hint: Can be done in single forward / backward pass. (Exercise : We will look at it next class)

  • The hessian diagonal computation is of the same order as the gradient computation.

Further simplifications: Iterative Magnitude Pruning (IMP)

  • Completely let go of the data consideration
  • Just prune the weights based on their magnitude (equivalent to assuming H_{kk} is same for all k) and rely on training to correct the errors

Algorithm for IMP

  1. Train a dense network to convergence.
  2. Compute the magnitude of each weight: |wi|.
  3. Remove weights with the smallest magnitudes (prune a fixed fraction or all below a threshold).
  4. Retrain the pruned network. Return to step 2 (iterate until target sparsity).

Iterative Magnitude Pruning (IMP)

  • Simple and easy to implement
  • Training is strictly required.

OBS vs OBD vs IMP (w/o re-training)

Comparison of OBS, OBD, and IMP pruning in 2D
  • OBD and IMP may lead to different pruning
  • Implications of using "less" informed pruning methods can be significant when not re-trained.

General Framework for Pruning

  1. Train a dense network to convergence.
  2. Compute the saliency score of each weight
  3. Remove weights with the smallest saliency scores
  4. Retrain the pruned network. Return to step 2 (iterate until target sparsity).

Results from the paper

Learning both Weights and Connections for Efficient Neural Networks" by Han et al. 2015

Results 1
Results 2
Results 3
MNIST Results 4
Results 5
Results 6
Results 7
Results 8
Results 9