LLM-Pruner: On the Structural Pruning of Large Language Models

In plain words

A large language model is expensive to store and slow to run. This paper shrinks one by cutting out parts of it, the way you might remove whole shelves rather than single books. The hard part is that the parts are wired together: removing one piece of a connected set breaks the others. So the method first maps which pieces must be removed together.

Next it scores each set of pieces by how much the model's answers depend on it, using a small amount of ordinary text and the gradients from it. The lowest-scoring sets are removed. The damaged model is then repaired with a cheap fine-tuning step that trains only small add-on matrices, and those are folded back in afterwards. The whole point is to need little original training data and little time, so the model stays general rather than being tuned for one task.

The authors test it on LLaMA-7B, Vicuna-7B and ChatGLM-6B, checking nine zero-shot benchmarks. Smaller models keep most of their score. Pruning too hard still breaks the model, and long generations can become incoherent.

TL;DR

• Build a dependency graph, so structures that must go together are removed as one group §3.1 app. A.
• Score each group with a first-order (or second-order) gradient-times-weight estimate on 10 BookCorpus samples, then drop the lowest-scoring share in each module §3.2 eq. 4 app. B.1.
• Recover with LoRA (rank 8) for 2 epochs on about 50k Alpaca samples, which the authors report takes about 2.5 hours on one RTX 4090 §3.3 app. B.2.
• On LLaMA-7B at 20% parameters removed, the Element² variant averages 60.07 against a re-run baseline of 63.25 (94.97%) tab. 1. This is author-reported.
• Magnitude-only (L2) pruning collapses near 20%, while LLM-Pruner reaches about 60% at similar perplexity on the authors' curve §4.3 fig. 4.
• At 50% the paper reports a large degradation and calls compression at that rate an open challenge §5 app. C.5.

Problem

Deployment and inference of LLMs are slowed by their size §1. Earlier compression methods for language models either target one task, or need the original pre-training corpus, which can be very large or proprietary, and their post-training takes many GPU days or hours §1 §2. The authors want a method that is task-agnostic, uses little data, and finishes quickly §1. They also argue that pruning without respecting dependencies damages the model beyond repair by fine-tuning §4.3 tab. 6.

Method

The method has three stages: find coupled structures, estimate group importance and prune, then recover with LoRA §3.

Stage 1: dependency groups §3.1. A neuron N_j depends on N_i when N_j's only incoming edge comes from N_i, and the reverse rule applies in the other direction eq. 1 eq. 2. Any neuron can trigger a closure over these rules, and the closure is one group of structures to be pruned together. Searching from every neuron finds all groups. Groups of types A and B form the "Block" strategy, and type C forms the "Channel" strategy §4.1 fig. 2. The paper states the two rules in both directions and the closure procedure, but does not give a pseudocode listing.

Stage 2: group importance §3.2. Let the calibration data be N samples x_i, y_i, with N = 10 §3.2. The change in loss from zeroing a weight is approximated by a Taylor expansion. The gradient term is kept, because the calibration data are not from the training set and the gradient is not near zero §3.2 eq. 3 eq. 4. Three estimates are defined:

• Vector: a structure-level score, |sum over k of g_k W_k|, from the gradient and weight eq. 3.
• Element¹ (first order): per-weight |g W|, summed over each output row eq. 4.
• Element² (second order): per-weight Fisher-style term, the sum over samples of (g_j W)², which approximates the Hessian diagonal eq. 5.

The per-structure scores are combined into a group score in one of four ways: Sum, Production, Max, or Last-only, which uses only the last structure in the group's execution order §3.2. The authors choose Sum §4.3 tab. 7.

Stage 3: pruning §3.2 app. B.1. Groups are ranked by score and the lowest share is removed. The ranking is local: groups are compared inside their own module, not across the whole network app. C.3. The first three layers and the last layer are left unchanged app. B.1; the first and last layers matter most for accuracy §4.2 fig. 3.

Stage 4: recovery §3.3. Each weight W gets a low-rank update ΔW = PQ, with P and Q trained and the rest of the model frozen eq. 6. After training, ΔW is merged back, so the final model has the original shape §3.3.

# LLM-Pruner: group, score, prune, recover  [§3]
# Input:  model (dense LLM), calib (N = 10 short text samples), ratio
# Output: pruned model with the LoRA update merged in

def llm_pruner(model, calib, ratio):
    # 1. Find every set of structures that must be removed together
    #    by letting each neuron trigger the dependency rules  [§3.1] [eq. 1]
    groups = discover_groups(model)

    # 2. Score each structure from gradients on the calibration samples
    #    [§3.2] [eq. 3] [eq. 4] [eq. 5]
    scores = {}
    for structure in all_structures(groups):
        scores[structure] = element_score(structure, calib)

    # 3. Sum the structure scores in each group  [§3.2]
    group_scores = {g: group_score(g, scores) for g in groups}

    # 4. Remove the lowest-scoring share of groups in each module; the first
    #    three and the last layer are kept  [app. B.1]
    pruned = remove_groups(model, groups, group_scores, ratio)

    # 5. Recover with LoRA on Alpaca-style data, then merge  [§3.3] [eq. 6]
    adapters = attach_lora(pruned, rank=8)
    train(adapters, alpaca_samples, epochs=2, lr=1e-4, warmup=100)
    return merge(pruned, adapters)  # [app. B.2]

# Per-weight scores used in step 2 (choose one)  [§3.2]

def element_score(weight, grad):
    # First-order: |g * W| summed over each output row  [eq. 4]
    # [code: LLMPruner/pruner/hf_llama_pruner.py#L266-L278]
    salience = weight * grad            # (out, in)
    return salience.abs().sum(dim=1)    # (out,)

def fisher_score(weight, sample_grads):
    # Second-order: per-sample squared gradients, then (g * W)^2  [eq. 5]
    # [code: hf_prune.py#L130-L142]
    fisher = sum(g * g for g in sample_grads) / N   # (out, in)
    return (weight * weight * fisher).abs().sum(dim=1)

def vector_score(weight, grad):
    # Structure-level: |sum of g * W| over each output row  [eq. 3]
    return (weight * grad).sum(dim=1).abs()

Hyperparameters

Results

Author-reported. Setup: zero-shot classification on BoolQ, PIQA, HellaSwag, WinoGrande, ARC-easy, ARC-challenge and OpenBookQA (average of seven), plus perplexity on WikiText2 and PTB, with the authors' own prompts §4.1. The "Baseline" below is the authors' re-run of the dense model (marked * in the paper), not the official score of the LLaMA paper, which is 68.59 for LLaMA-7B tab. 1.

Notes on the table: the 94.97% and 92.03% are the paper's own ratios of its average to the re-run baseline. The "ChatGLM-6B" baseline is the authors' run, with an average of 47.05 over six tasks tab. 5. The LLaMA-3B row is also from the paper's own comparison, under the same Alpaca tuning tab. 9. For the DistilBERT row, "Baseline" is the comparison model, not the dense LLaMA tab. 8. Dependency-based groups matter: without them, WikiText2 perplexity after tuning is 13307.46 against 17.58 with them tab. 6.

Paper vs code

Code: official repository at commit 128a07d977f9b205d60ab14cfbc6a78f8a8e39d2, Apache-2.0 LICENSE. Discrepancies found:

• Calibration length. Paper: each BookCorpus sample truncated to 128 tokens §4.1. Code: calibration windows of 64 tokens for the pruning step hf_prune.py and the same in the channel path hf_prune.py. The window sampler picks a random 64-token slice of a document with at least 64 tokens example_samples.py.
• Layers kept. Paper: first three and last layer unchanged app. B.1. Code defaults are block ranges 3 to 30 hf_prune.py. The example command in the README and script uses 4 to 30 README.md llama_prune.sh. The paper's appendix text gives 5 to 30 for the 20% setting app. B.1.
• Second-order scale. Paper: ½ Σ_j (g_j W)² eq. 5. Code: mean of squared per-sample gradients, then W² times that mean hf_prune.py hf_llama_pruner.py. The ½ and the 1/N factor are constants across weights, so they do not change the ranking (inferred from the code, not stated in the paper).
• Gradient source for first order. Paper: gradient of the loss over the N samples eq. 4. Code: one forward pass over the 10-sample batch with the HF mean-token loss, then backward hf_prune.py. A constant scale does not change the ranking, but the batch loss is a mean, not a sum (inferred).
• Default importance. README says param_mix is the default README.md. The argparse default is param_first hf_prune.py. The paper's reported results use Element¹ and Element², i.e. first and second order, not mix.
• Group reduction. The code's _reduce offers sum, mean, max, prod, first, second, and None hf_llama_pruner.py. The paper's "Last-only" is not named there; whether first is the same thing is not verified.
• param_mix attention branch. On line 264 the attention branch applies -salience to a dictionary (salience is a dict in that branch), and its sign convention differs from the MLP branch on line 271 hf_llama_pruner.py. This looks like a bug. It is not used for the paper's results (inferred from reading; not run).
• Attention structure. The code sets consecutive group sizes for q_proj to head_dim hf_prune.py and recomputes num_heads from the pruned q_proj size hf_prune.py. This means attention is pruned in head-sized blocks, which the paper describes only as Type B groups fig. 2 (inferred).
• Recovery defaults. Running post_training.py without flags uses learning rate 3e-4 and batch size 128 post_training.py post_training.py. The paper's values (1e-4, batch 64 or 128) are only reached by passing flags as in the script llama_prune.sh app. B.2.

Pitfalls

• Baseline choice. The headline "94.97%" divides by the authors' re-run LLaMA-7B (63.25), not by the official LLaMA-7B score (68.59), because the authors' prompts differ from LLaMA's tab. 1 §4.2. Compare against the same prompt setup.
• Ratio is per group, not per parameter. A 25% group ratio gives about 20% of parameters removed, because groups differ in size and embeddings stay app. B.1. Setting pruning_ratio to 0.20 is not the same as the paper's 20% setting.
• Local ranking is the default. Global ranking is slightly worse in the authors' own test; the authors suspect the varying score magnitudes across layers and modules app. C.3.
• Sum is the chosen aggregator. Max gives the best perplexity but the worst classification; Last-only gives the best classification but the worst generation §4.3 tab. 7.
• First vs second order. Element² is best after tuning on LLaMA and Vicuna, but Element¹ is close and cheaper §4.2. On ChatGLM the per-weight estimates fail, and Vector wins §4.2 tab. 5. Do not assume one estimate transfers to another model.
• Second order is costly. The code does one backward pass per calibration sample for Element² hf_prune.py, plus the batch pass for the first-order term.
• Overfitting during recovery. The text says perplexity falls for about 300 steps and then rises gradually §4.3 fig. 5. The appendix table shows WikiText2 perplexity flat between steps 400 and 1000, and the average peaks at step 600 app. C.4. Choose the step count on general metrics, as that table does.
• Generation quality. Pruned models sometimes produce meaningless or repeated tokens on long outputs §4.3. The paper's generations are shown only as examples app. D tab. 10.
• High pruning ratios fail. At 50% on LLaMA-7B the model degrades strongly, and recovery helps by about 7.39% app. C.5. The authors say this remains open §5.
• Data assumptions. The method needs ten text windows (64 tokens in code, 128 in the paper). The authors note their calibration domains may not match the training data, and argue they still carry useful signal §3.2. The paper's data for ChatGLM is dialogue, not BookCorpus §4.1. Calibration data outside the training domain can shift the importance estimates §3.2.
• Hidden costs. The code sets requires_grad on every parameter before scoring, so the gradient and per-sample accumulators cover every weight hf_prune.py. Latency falls from 69.32 s to 58.55 s on the authors' test tab. 3. Adding LLM.int8() to the pruned model cut memory to 5444 MiB but raised latency from 58.55 s to 63.10 s app. C.2.

Lineage

• Builds on: a dependency graph for structural pruning, similar to the one cited as [fang2023depgraph] in §3.1; the loss-change estimate cited as [lecun1989optimal] in §3.2; LoRA, cited as [hulora] in §3.3.
• Compared against: L2 magnitude pruning and random group pruning on the same groups app. B.1 tab. 1; DistilBERT-style layer dropping §4.3 tab. 8; StableLM-3B trained from scratch §4.3 tab. 9.
• Superseded by: not known from content already in this wiki.