Preserve the function a neuron computes under a sparsity budget, then
greedily choose which weight to remove.
Setup
A linear neuron computes an inner product of weights and inputs:
\[
y = \langle w, x \rangle = w \cdot x,
\qquad w, x \in \mathbb{R}^d.
\]
Inputs $x$ weighted by $w$ produce
$y = \vec{w}\cdot\vec{x}$.
Q. What should we try to preserve when pruning?
Behavior of the neuron
What function it was representing
Remark — inference vs training
When pruning at inference time, the goal is to preserve the
function the model already computes. During training, one might instead
care about expressive power (capacity to learn), which leads to
different considerations.
Preserving behavior under compression
We want a sparse (or budget-constrained) weight vector $\hat{w}$ that
matches the original neuron on all inputs:
\[
\hat{w} \in \mathbb{R}^d
\quad\text{s.t.}\quad
\|\hat{w}\|_0 \le k,
\qquad
\hat{w}\cdot x \;\simeq\; w\cdot x
\quad \forall\, x.
\]
Let $X$ be representative data. Then choose $\hat{w}$ by least-squares
reconstruction of the neuron’s outputs:
\[
\hat{w}
=
\underset{\|\hat{w}\|_0 \le k}{\arg\min}
\;
\bigl\|
\hat{w}\cdot X - w\cdot X
\bigr\|_2^2 .
\]
Equivalently, introduce a binary mask $M$ and write
\[
\hat{w},\, M
=
\underset{
\substack{
\hat{w}\in\mathbb{R}^d \\[2pt]
M\in\{0,1\}^d \\[2pt]
\|M\|_0 \le k
}
}{\arg\min}
\;
\bigl\|
(\hat{w} \odot M)\cdot X - w\cdot X
\bigr\|_2^2 ,
\]
where $\odot$ is the elementwise (Hadamard) product.
Jointly optimizing $M$ and $\hat{w}$ is
NP-hard
(Blumensath & Davies, 2008)
$\;\rightarrow\;$ use a greedy algorithm.
Which element to prune first?
Simple heuristics are incomplete:
Magnitude?
$|w_i|$ large does not help if $|x_i|$ is very small.
A joint score $f(|w_i|, |x_i|)$ still
does not consider correlations among coordinates.
True model
\[ y = w \cdot X \]
Approximated model
\[ \hat{y} = \hat{w} \cdot X \]
To remove one weight, solve the constrained reconstruction problem
(at most $d-1$ nonzeros — some $\hat{w}_i = 0$):
\[
\min_{\|\hat{w}\|_0 \le d-1}
\;
\bigl\| \hat{w}\cdot X - w\cdot X \bigr\|_2^2 .
\]
Greedy step: try zeroing each coordinate in turn and keep the best:
\[
\min_{\hat{w}_1=0}
\|\hat{w}\cdot X - w\cdot X\|_2^2,
\quad
\min_{\hat{w}_2=0}
\|\hat{w}\cdot X - w\cdot X\|_2^2,
\quad
\ldots,
\quad
\min_{\hat{w}_d=0}
\|\hat{w}\cdot X - w\cdot X\|_2^2.
\]
Whichever index $i$ gives the best (smallest) loss is the one we prune.
Subproblem for a fixed index $i$
Drop column $i$ of $X$ (write $X_{-i}$) and refit the remaining weights
by linear regression against the original outputs $y = X\cdot w$: