Post-training pruning

Overview

Post-training pruning removes weights from a trained network without a retraining phase. The model is fixed, and the method uses a small amount of generic input text to decide what to remove and how to adjust what remains. The goal is to keep the layer's output close to its dense output on those inputs 2301.00774 §2, eq. 1. The paper that covers this concept in this wiki argues that retraining is too costly at GPT scale, so one-shot methods are the practical option there 2301.00774 §1.

The shared core is a per-layer problem: pick a binary mask and new weights that minimise the squared output error on calibration inputs. The mask and the weights are coupled, so the exact problem is NP-hard 2301.00774 §2. Once the mask is fixed, the best weights have a closed form, but computing it row by row needs a separate matrix inversion for each row [arxiv:2301.00774 §3.1, eq. 2]. Methods in this space differ in how they handle that cost and how they choose the mask.

Only one paper with a card in this wiki covers the concept so far. Its predecessors and comparisons (Optimal Brain Surgeon, OBC, AdaPrune, layer-wise magnitude pruning) are named in its bibliography and have no card here, so they are not described on this page 2301.00774 §2, §4.

Variants

Greedy column-wise pruning with Hessian updates (SparseGPT)

The layer is pruned one column at a time. A removed weight is replaced by an update to the columns that are still free, using the inverse Hessian of the calibration inputs. The inverse is shared by all rows, which is what makes the method fast enough for 175B-parameter models [arxiv:2301.00774 §3.1, eq. 4, eq. 5]. The mask is chosen in blocks of 128 columns from the current, already-updated weights, so the mask reflects the corrections made so far 2301.00774 §3.2. The authors report 50% unstructured sparsity, and the 2:4 and 4:8 patterns, on the OPT family and BLOOM-176B [arxiv:2301.00774 §4.1, tab. 1, tab. 6]. Card: SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot.

Calibration needs 128 random segments of 2048 tokens from C4, and Hessian damping of 1% of the mean diagonal 2301.00774 §4, app. A. The paper reports that the damping value matters little between 0.001 and 0.1 2301.00774 app. A, fig. 9.

Magnitude pruning (baseline, not a method you should pick at 50%)

Keeping the largest weights by absolute value needs no calibration data. The paper's baseline shows that at 50% unstructured sparsity it fails on the models it tests. On OPT-175B it reaches WikiText2 perplexity 4.3e4 against 8.35 dense 2301.00774 tab. 1. It is listed here only as the reference point for the method above 2301.00774 §4.1.

When to use which

The only evidence on this page is from one paper, and every number is author-reported. The table is based on constraints the paper states.

Trade-offs to weigh:

• Speed claims are narrow. The authors' CPU speed-up for OPT-2.7B is 1.82× at 50% 2301.00774 app. E, tab. 7. On GPU they report 1.54× to 1.79× for single matrix products of OPT-175B shapes [arxiv:2301.00774 tab. 8]. No end-to-end GPU speed-up is given 2301.00774 §4.1.
• Later layers are more sensitive. Leaving the last third of blocks dense gave the best accuracy in the paper's partial 2:4 study [arxiv:2301.00774 §4.1, fig. 7].
• Quantization can be combined. The paper's joint 4-bit variant is compared against GPTQ at 3 bits, at the same storage 2301.00774 §4.1, app. C. See post-training quantization for the quantization side.

Implementation sketch

The common core is the layer-wise loop of SparseGPT's unstructured variant 2301.00774 alg. 1. The variants differ in the mask rule and the block sizes. This is a sketch in words; use the cited algorithm for exact steps.

# Post-training pruning of one linear layer  [arxiv:2301.00774 alg. 1]
# Input:  weights W (d_row, d_col); calibration inputs X (d_col, n)
# Output: pruned weights, zero where the mask removes them

def prune_layer(W, X, sparsity):
    # Hessian of the layer input from the calibration set  [§3.1]
    H = X @ X.T
    # Damping keeps the inverse stable: 1% of the mean diagonal  [app. A]
    H = H + 0.01 * mean(diagonal(H)) * I
    # Upper Cholesky factor of H^-1: used for updates and mask scores
    # [alg. 1] [§3.4]
    L = cholesky_upper(inverse(H))
    for each block of columns:
        # Mask from w^2 / [H^-1]_cc^2 in this block  [§3.2]
        keep = score_threshold(W, L, sparsity)
        for each column c in the block:
            # Quantize-or-zero target for this column  [§3.5]
            target = keep[:, c] * W[:, c]
            err = (W[:, c] - target) / L[c, c]
            # Push the error onto the columns still free  [alg. 1]
            W[:, c+1:block_end] -= outer(err, L[c, c+1:block_end])
    return W * keep

Notes for variants:

• Unstructured: mask blocks of 128 columns 2301.00774 §3.2.
• n:m (2:4, 4:8): the mask block equals m, and each group of m weights in a row keeps the highest scores 2301.00774 §3.3.
• Joint quantization: the rounding error replaces the pruning error in the update 2301.00774 §3.5, eq. 7.

Contradictions

• Inside the paper: Hessian inverse. Equation 3 and §3.2 are written in terms of entries of H^-1. Algorithm 1 and the released code use the upper Cholesky factor of H^-1 for both the updates and the mask scores 2301.00774 eq. 3, §3.2, alg. 1. The card records this as a discrepancy inside the paper [card: 2301.00774, Paper vs code].
• Paper text vs released code: Hessian collection. The paper says layers are pruned sequentially 2301.00774 §4. The released code collects Hessians for all linear layers of a block before pruning any of them, and only block outputs propagate forward [card: 2301.00774, Paper vs code]. An implementation that follows one reading will get different masks from one that follows the other.
• Runtime. The §4 text gives about 4 hours on one A100 80GB for 175B; the abstract says under 4.5 hours 2301.00774 §4, abstract.
• Zero-shot result. At 50% the average zero-shot accuracy on OPT-175B is 70.52 against 70.29 dense, and 2:4 beats dense on Lambada [arxiv:2301.00774 tab. 2]. The authors call this dataset-specific [arxiv:2301.00774 §4.1, app. C], so this does not show that sparsity improves the model.

Papers

• SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot: one-shot unstructured, 2:4 and 4:8 pruning of GPT-family layers with a Hessian-based update, and joint 4-bit quantization.