Structured (neuron) pruning

The idea

Unstructured pruning leaves zeros that still get stored and multiplied. Structured pruning deletes whole neurons. Removing one hidden unit removes one column of the matrix feeding it and one row of the matrix it feeds, so the matrices really shrink and the model is smaller and faster on any hardware. Which neurons to keep is decided by an importance score. In a conv net the same idea removes channels, which is what the YOLOv8n cockpit does to its detector.

How it works

1. Run the trained net on a probe batch (2,000 training examples here).
2. Score every hidden neuron. l1 is the sum of |w| over its incoming weights. act is its mean |activation|. taylor is the sum over examples of |a * dL/da|, where a is the activation and dL/da the loss gradient with respect to it: a first-order estimate of the loss change if a were zeroed. random is a control.
3. In every hidden layer, keep the k best neurons and delete the rest.
4. Recover by training the smaller net, with labels or by distilling from the original.

The code, line by line

neuron_importance in lab/compress.py returns one score array per hidden layer. For l1, np.abs(net.W[i]).sum(0) sums each column of W[i], the weights going into each neuron. It never looks at data. The other criteria run z, hid, cache = net.forward(X, cache=True), where hid holds each hidden layer's post-ReLU output, and act is np.abs(h).mean(0).

For taylor, dz = (p - np.eye(...)[y]) / len(z) is the gradient of mean cross-entropy with respect to the logits. net.backward(cache, dz, hidden_grads=hg) stores dL/d(hidden[i]) in hg[i], and the score is np.abs(h * g).sum(0). Taking abs per example means a neuron that helps some examples and hurts others still scores high. The group branch is the task-aware variant (see Task-aware pruning).

structured_prune copies the net, takes idx = np.argsort(-scores[i])[:k] per hidden layer, and calls net.keep_neurons(i, idx). Scores are computed once, on the unpruned net.

MLP.keep_neurons in lab/nn.py does the surgery. self.W[layer][:, idx] keeps the chosen output columns, self.b[layer][idx] their biases, and self.W[layer + 1][idx, :] the matching input rows of the next layer, so shapes still line up. Masks are sliced the same way and self.sizes[layer + 1] = len(idx) records the new width. The first line refuses residual layers. On a stream h + f(h), one neuron index is shared by every layer that adds into it, and that dependency group must be cut together.

run_structured in lab/methods.py prunes the 3x256 teacher to [w] * 3, recovers for FT = 500 steps, and compares with a fresh 3xw net distilled for 1,000 steps at the same learning rate.

What the simulations show

From prune_structured (teacher accuracy 0.9139; recovered arms are means of 2 seeds; pruned and fresh nets share optimizer and learning rate; "short" is 750 KD steps and "full" 3,000).

Before recovery:

After recovery with KD:

Halving every layer drops accuracy to 0.30-0.42 before recovery. That is the cliff the cockpit's STATUS.md reports for YOLOv8n, where even a 5% one-shot channel cut sent bird mAP@50 to 0. Recovery repairs a 2x cut completely. With the short budget, the L1-pruned net beats a fresh one at widths 128-32 (0.8927 against 0.8685 at 64). With the full budget the arms are within 0.01 at 128 and 64, L1 still leads at 32 (0.8603 against 0.8469), and at a 16x cut the fresh net wins (0.7007 against 0.6856). The pruned start mostly buys faster recovery. L1, which ignores data, is the best criterion from width 64 down. Taylor trails random at widths 128-32. Recovery with labels instead of KD ends lower in every full-budget arm.

Page runs at the defaults (width 32, taylor, KD): no recovery 0.1175, recovered 0.6835, fresh 0.7794. With l1, 0.7212. With random, 0.6653. With labels, 0.6737. The page gives the fresh net twice the steps, so it favours it.

The toy can't show speed. On Apple's MPS, the cockpit saw a 30% channel cut reduce MACs by 35% while latency got 6% worse.

When it breaks

• The cliff is steep. Never judge a structurally pruned model before fine-tuning it.
• Deep cuts lose to a fresh start. At 16x the pruned weights end worse than a random init.
• Stale scores. Scores come from the unpruned net, so cutting layer 1 changes what layer 2's neurons mean.
• Taylor at a minimum. On a well-trained net dL/da is small and noisy.
• Residual streams and concatenations tie neurons across layers, and keep_neurons refuses them rather than guessing.

Try this

• Set keep width 128, recovery none, and cycle importance. See how close random comes.
• At keep width 128 with recovery kd, compare the recovered row with "fresh 3x128, KD from scratch". Repeat at 16.
• At width 32, compare recovery labels and kd, then importance l1.

Going further

In a detector a "neuron" is an output channel: one conv filter plus the matching input slice of every layer that reads it. Batch-norm, concatenations and residual adds tie channels into dependency groups. The cockpit's harness/prune.py uses torch-pruning's GroupNormPruner with GroupMagnitudeImportance, the conv equivalent of l1 here, and leaves the Detect head alone (see Channel pruning with dependency groups). Its loop matches what these numbers argue for: a ~12% cut fine-tuned on 400 bird images recovered to 0.388 mAP@50, 91% of native 0.428.

Source code

The exact code that runs, read from the repository when this page was built.

def neuron_importance(net: MLP, X: np.ndarray, y: np.ndarray | None = None,
                      method: str = "l1", rng: np.random.Generator | None = None,
                      group: np.ndarray | None = None) -> list[np.ndarray]:
    """One score per hidden neuron, per hidden layer (higher = keep).

    l1        -- L1 norm of the neuron's incoming weights (torch-pruning's
                 GroupMagnitudeImportance, the criterion the YOLOv8n cockpit uses)
    act       -- mean |activation| on data
    taylor    -- |a * dL/da| summed over data: first-order estimate of the loss
                 change if the neuron is removed (data- and task-aware)
    random    -- control
    ``group`` (C x K 0/1 matrix) makes 'taylor' task-aware: the loss becomes CE
    in the grouped label space (e.g. bird-vs-rest), asking "which neurons matter
    for *this* task?" rather than for all 20 classes.
    """
    n_hidden = net.n_layers - 1
    if method == "random":
        rng = rng or np.random.default_rng(0)
        return [rng.random(net.sizes[i + 1]) for i in range(n_hidden)]
    if method == "l1":
        return [np.abs(net.W[i]).sum(0) for i in range(n_hidden)]
    z, hid, cache = net.forward(X, cache=True)
    if method == "act":
        return [np.abs(h).mean(0) for h in hid]
    if method == "taylor":
        p = softmax(z)
        if group is None:
            dz = (p - np.eye(z.shape[1], dtype=F)[y]) / len(z)
        else:  # d/dz of -log q[yg], q = p @ G:  p - p * G[:, yg] / q[yg]
            yg = group[y].argmax(1)
            Gy = group[:, yg].T
            qy = (p * Gy).sum(1, keepdims=True)
            dz = (p - p * Gy / qy) / len(z)
        hg: list = []
        net.backward(cache, dz, hidden_grads=hg)
        return [np.abs(h * g).sum(0) for h, g in zip(hid, hg)]
    raise ValueError(method)

def structured_prune(net: MLP, keep: float | list[int], scores: list[np.ndarray]) -> MLP:
    """Keep the top-scoring neurons in every hidden layer. ``keep`` is a
    fraction (same for every layer) or an explicit per-layer width."""
    net = net.copy()
    for i in range(net.n_layers - 1):
        k = keep[i] if isinstance(keep, (list, tuple)) else max(1, int(round(net.sizes[i + 1] * keep)))
        idx = np.argsort(-scores[i])[:k]
        net.keep_neurons(i, idx)
    return net

    def keep_neurons(self, layer: int, idx: np.ndarray) -> None:
        """Really remove neurons: keep only ``idx`` outputs of hidden layer
        ``layer`` (and the matching inputs of the next layer). The network gets
        genuinely smaller -- the numpy equivalent of torch-pruning's surgery."""
        if self.res[layer] or self.res[layer + 1]:
            raise ValueError("neuron removal across a residual stream needs a "
                             "dependency group (the torch-pruning problem); not supported")
        idx = np.sort(idx)
        self.W[layer] = self.W[layer][:, idx]
        self.b[layer] = self.b[layer][idx]
        self.W[layer + 1] = self.W[layer + 1][idx, :]
        for j in (layer, layer + 1):
            if self.masks[j] is not None:
                self.masks[j] = self.masks[j][:, idx] if j == layer else self.masks[j][idx, :]
        self.sizes[layer + 1] = len(idx)

Simulation figures

Structured (neuron) pruning: criteria, recovery, budget

Delete whole neurons so the net is genuinely smaller. Which importance criterion picks the right ones, and is a pruned teacher a head start over training the small shape fresh, at the same number of training steps?

'short' = 750 steps, 'full' = 3000 steps; pruned and fresh nets use the same optimizer and learning rate, so the only difference is the starting weights.

Full table: Structured (neuron) pruning: criteria, recovery, budget

Papers

• Li et al. Pruning Filters for Efficient ConvNets (2017)
• Molchanov et al. Importance Estimation for Neural Network Pruning (2019)