Efficient AI · Lecture notes

Bonsai

A forward-pass-only recipe for structured pruning: sample feasible sub-models, measure their utility, fit module scores, and keep the best components under a sparsity budget.

Five parts

Problem setting
Sampling feasible sub-models
Turning samples into supervision
Fitting module scores
Regularization and complements

1. Problem setting

We want to do structured pruning when the only affordable operation is inference. In particular, assume we can run a reasonable number of forward passes, but we cannot rely on gradients, Hessians, or expensive retraining loops.

Starting point. How do we prune with only this much access?

SparseGPT, Wanda, and magnitude pruning already give useful saliency scores, so one option is one-shot pruning. Bonsai asks whether those saliency scores can instead guide a small set of forward-pass experiments, so we learn a better structured pruning decision.

The core idea is to sample several feasible pruned models $\hat{M}_i$, run each on a task or calibration set, and record its score $s_i$. This produces a small dataset:

\[ \mathcal{D} = \{(\hat{M}_i, s_i)\}_{i=1}^{n}, \]

Two questions. How should we sample the $\hat{M}_i$? Once we have $\mathcal{D}$, how do we turn it into a pruning rule?

2. Sampling feasible sub-models

Uniformly sampling structured sub-models is too wasteful. Some samples are obviously broken; for example, dropping every component in a layer can destroy the model. We therefore keep the sampling layer-aware and bias it using saliency from existing methods.

Notation: sparsity $p$ $p \in (0,1)$ is the target sparsity (fraction pruned). Example: $p = 0.1$ means prune $10\%$ and keep $90\%$ of the weights / modules.

Freeze most of the model Within each layer, keep the most salient $(1-2p)$ fraction of components fixed (always chosen). The remaining $2p$ fraction is the candidate pool. Each sampled mask keeps about half of that pool, giving total keep fraction $(1-2p)+p = 1-p$ and sparsity $p$.

Layer-wise split
Hatched: always kept ($(1-2p)$).
Open: candidate pool (size $2p$); keep half of it ($\Rightarrow$ total keep $1-p$).

The prior controls which candidate modules are more likely to appear in sampled sub-models. Saliency gives a natural prior; pure uniform sampling is usually too noisy.

3. Turning samples into supervision

After sampling, the most direct strategy is to pick the best evaluated sub-model. That wastes most of the information in $\mathcal{D}$: every score says something about many modules at once.

Why not just pick the best sample?

The number of possible structured components (neurons, heads, channels, blocks, etc.) is much larger than the number of forward passes we can afford. Most configurations will never be evaluated. Instead of treating each sample as an isolated candidate, we fit a surrogate that explains each score in terms of the modules that were kept.

Utility model Each forward pass gives an observed utility $s_i = U(\hat{M}_i)$. The surrogate predicts $\hat{s}_i$ from the set of modules kept in $\hat{M}_i$.

For example, $U$ might be validation accuracy, downstream task score, or negative perplexity. The important distinction is that $s_i$ is measured, while $\hat{s}_i$ is predicted:

\[ s_i = U(\hat{M}_i), \qquad \hat{s}_i = \hat{u}\bigl(\{m_j : m_j \in \hat{M}_i\}\bigr) \approx s_i. \]

4. Fitting module scores

Linear model

The simplest surrogate assumes that each module contributes additively. If module $m_j$ is kept, it contributes a score $\beta_j$; if it is pruned, it contributes nothing.

\[ \hat{s}_i = \sum_{m_j \in \hat{M}_i} \beta_j = \beta^\top \alpha_{\hat{M}_i}, \]

where $\alpha_{\hat{M}_i}\in\{0,1\}^N$ has a $1$ in positions of kept modules.

As a linear regression problem

Stack the binary masks into a sparse design matrix. Rows correspond to sampled sub-models, columns correspond to structural components under consideration.

$N$: # modules / structural components under consideration (columns)
$n$: # data samples / # forward passes (rows)

\[ \hat{M}\,\vec{\beta} = \vec{s} \]

\[ \underbrace{ \begin{bmatrix} 1 & 1 & 1 & 0 & 0 & 0 & 1 & \cdots \\ \vdots & & & & & & & \\ \end{bmatrix} }_{\text{sparse mask matrix } \hat{M}\ (n\times N)} \; \underbrace{ \begin{bmatrix} \beta_1 \\ \vdots \\ \beta_N \end{bmatrix} }_{N\times 1} = \underbrace{ \begin{bmatrix} s_1 \\ \vdots \\ s_n \end{bmatrix} }_{n\times 1} \]

Typically $N \gg n$ — the system is underspecified.

The goal is to estimate $\beta$, one score per module. If the linear model is a good approximation, pruning becomes simple: keep the frozen $(1-2p)$ modules, then from the candidate pool keep the highest-$\beta$ modules until the total keep budget $1-p$ is reached.

Summary $\hat{M}\beta = s$ is underspecified when $N \gg n$.

5. Regularization and complements

Usually $N \gg n$, so the regression is underspecified. We therefore fit $\beta$ with squared loss plus $\ell_2$ regularization:

\[ \beta^* = \underset{\beta}{\arg\min} \; \frac{1}{n} \sum_{i=1}^{n} \bigl( s_i - \beta^\top \alpha_{\hat{M}_i} \bigr)^2 + \gamma\,\|\beta\|_2^2 \]

Equivalently $\|\hat{M}\beta - s\|_2^2 + \gamma\|\beta\|_2^2$, where each row of $\hat{M}$ is the binary mask $\alpha_{\hat{M}_i}$.

When constructing each row of $\hat{M}$, draw a mask at sparsity $p$ per layer. This keeps the samples feasible and avoids rows that accidentally remove too much of one layer while leaving another layer untouched.

A remaining issue is coverage. Even with a saliency prior, some module may rarely or never appear in the sampled masks. Then its coefficient $\beta_j$ will be poorly estimated.

Complement trick If you sample $\hat{M}_i$, also sample its complement $\hat{M}_i^{c}$. Flip bits only inside the candidate pool (the bottom $2p$); the frozen $(1-2p)$ part stays unchanged.

\[ \hat{M}_i^{c} = \bigl\{\, m \;\big|\; m \notin \hat{M}_i \text{ and } m \text{ is in the candidate pool} \,\bigr\} \;\cup\; \{\text{frozen modules}\}. \]

Exercise. Prove that having the complement of a sample in the data helps reduce the variance of $\beta$ estimation.

Kolawole et al., Everybody Prune Now: Structured Pruning of LLMs with Only Forward Passes (arXiv:2402.05406). Complements reduce regression variance on binary inputs (cf. Covert & Lee, 2020).