🤖 AI

GRASPrune: Global Gating for Budgeted Structured Pruning of Large Language Models

GRASPrune is a structured pruning framework that jointly reduces FFN channels and KV head groups in large language models under a global budget by learning lightweight gate scores with a projected straight-through estimator to enforce hard masks during training, followed by scaling calibration to produce smaller, efficient dense checkpoints without full model fine-tuning.

Original authors: Ziyang Wang, Jiangfeng Xiao, Chuan Xiao, Ruoxiang Li, Rui Mao, Jianbin Qin

Published 2026-04-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ziyang Wang, Jiangfeng Xiao, Chuan Xiao, Ruoxiang Li, Rui Mao, Jianbin Qin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, super-intelligent library (a Large Language Model like LLaMA). This library is incredibly smart, but it's also huge. It takes up a giant warehouse (memory), and when you ask it a question, it has to run through thousands of aisles to find the answer, which takes a long time (latency).

Running this library is expensive. You need powerful, expensive computers just to keep the lights on.

The Problem:
You want to shrink the library to make it cheaper and faster to run. But if you just randomly throw away books (parameters), the library stops making sense. It becomes a confused mess.

Existing methods of shrinking libraries usually do one of two things:

The "Layer-by-Layer" approach: They look at each floor of the library and decide, "Okay, we'll cut 10% of the books on this floor, 10% on that floor." But this is rigid. Maybe Floor 3 is full of useless junk, while Floor 4 holds the most important secrets. Cutting them equally is a bad idea.
The "Score First, Cut Later" approach: They grade every book on how important it is, then try to fit the best ones into a smaller box. But by the time they try to fit them in, they realize they've picked too many heavy books and not enough light ones, or they cut the wrong things. They have to go back and re-arrange everything, which takes forever.

The Solution: GRASPrune
The authors of this paper created a new method called GRASPrune. Think of it as a smart, global budget manager for your library.

Here is how it works, using simple analogies:

1. The "Global Budget" (The Wallet)

Instead of giving each floor a fixed budget, GRASPrune gives the entire library one single wallet.

Some books are heavy (like the FFN channels—the "thinking" parts of the model).
Some books are light but take up shelf space (like the KV heads—the "memory" parts of the model).
The goal is to keep the library as smart as possible while staying within the total weight limit of the wallet.

2. The "Smart Gatekeeper" (The Projected STE)

This is the magic trick.

Old way: You ask every book, "How important are you?" and write down a score. Then you try to pick the best ones.
GRASPrune way: It puts a gatekeeper at the door. As the library learns which books are important, the gatekeeper immediately checks the wallet.
- If picking a heavy book would break the budget, the gatekeeper says, "No, you can't go in yet," even if the book is important.
- It forces the library to learn while respecting the budget. It's like training a runner while wearing a heavy backpack; they learn to run efficiently with the weight, not just after taking it off.

3. The "Tuning Knob" (Scaling Calibration)

Once the gatekeeper has decided which books stay and which go, the library might feel a little "off." The remaining books might be shouting too loud or too quiet compared to the original setup.

GRASPrune adds a tiny volume knob (a scaling factor) to the books that stayed.
It turns the volume up or down just enough so the remaining books work together perfectly.
Crucially: It then folds these volume knobs inside the books. So, when you actually use the library later, you don't need the knobs anymore. You just have a smaller, perfectly tuned library with no extra baggage.

Why is this a big deal?

It's Fast: It doesn't need to re-read the whole library or train a new teacher. It just does a quick "gatekeeper" check on a small sample of text (about 6 minutes on a powerful computer).
It's Balanced: It doesn't just cut the "thinking" parts or the "memory" parts. It cuts the right mix of both to save the most space for the least loss in intelligence.
It Works: When they tested it on a 7-billion-parameter model (LLaMA-2-7B), they cut 50% of the size (half the books!) and the library was still almost as smart as the original. It could still write stories, answer questions, and solve puzzles almost as well as the giant version.

In a nutshell:
GRASPrune is like hiring a super-efficient librarian who doesn't just randomly throw away books. Instead, they look at the whole building, manage a strict budget, and decide exactly which books to keep and which to toss, ensuring the library stays small, fast, and cheap to run, without losing its genius.

1. Problem Statement

Large Language Models (LLMs) face significant deployment challenges due to high memory consumption and latency. These costs stem from:

Parameter Count: The sheer size of model weights.
KV Cache: The memory required to store Key-Value pairs during inference, which grows linearly with sequence length.

Existing structured pruning methods often suffer from three main limitations:

Decoupled Pruning: They typically prune Feed-Forward Network (FFN) channels and KV head groups separately, ignoring that they compete for the same global deployment budget.
Heuristic Allocation: Many methods use fixed per-layer sparsity schedules or depth-dependent rules rather than learning a single global allocation strategy.
Score-First, Select-Later: Most pipelines learn importance scores without constraints and apply the budget only after training. This decouples the learning objective from the final discrete selection, often leading to suboptimal allocations or requiring expensive reconstruction/fine-tuning.

2. Methodology: GRASPrune

GRASPrune proposes a global budgeted structured pruning framework that jointly prunes FFN channels and KV head groups under a single resource constraint while keeping the backbone weights frozen.

Core Components:

Global Budget Formulation:
- The problem is formulated as a constrained optimization: $\min L(\theta; z)$ subject to $\sum c_i z_i \leq B$ , where $z$ is a binary mask, $c_i$ is the cost of unit $i$ , and $B$ is the total budget.
- Cost Model: A two-level cost model is used. FFN channels have a unit cost of 1. KV head groups have a cost $\alpha$ derived from their parameter footprint (approximating the ratio of KV projections to FFN projections). This allows FFN and KV units to compete fairly under one global budget.
Projected Straight-Through Estimator (STE) Training:
- Instead of learning unconstrained scores and ranking them later, GRASPrune enforces the budget constraint at every optimization step.
- Forward Pass: Continuous gate probabilities $p$ are projected to a hard binary mask $m$ that satisfies the budget. This is done by sorting units by $p$ and selecting the top units until the budget is exhausted.
- Backward Pass: Gradients are propagated through the soft probabilities $p$ using a Straight-Through Estimator (STE), treating the hard mask as the identity for gradient flow. This ensures the learned scores are optimized specifically for the selection rule used during deployment.
Scaling Calibration:
- After the mask is fixed, a lightweight calibration stage adjusts scaling factors ( $\gamma$ ) for the retained units to mitigate scale mismatches caused by pruning.
- These factors are folded into the pruned weights, resulting in a smaller dense checkpoint with no extra parameters or inference overhead.

3. Key Contributions

Unified Global Formulation: The first framework to treat FFN channel pruning and KV head group pruning as a single global selection problem with heterogeneous unit costs, optimizing a single mask under one budget constraint.
Budget-Feasible Training: Introduces a projected STE loop where the budget constraint is active during training. This prevents the "score-first" mismatch and ensures the learned allocation is directly aligned with the deployable mask.
Efficient Dense Checkpoint Generation: Achieves high compression without full model fine-tuning or task-specific distillation. The method learns only scalar gates on a small unlabeled calibration set (512 sequences) and produces a standard dense model ready for deployment.

4. Experimental Results

The method was evaluated primarily on LLaMA-2-7B and compared against state-of-the-art baselines like LLM-Pruner, SliceGPT, and FLAP.

Performance at High Compression:
- At 50% parameter retention, GRASPrune achieves a WikiText-2 perplexity of 12.18, significantly outperforming baselines (e.g., FLAP: 24.80, SliceGPT: 26.10).
- At 40% retention, it maintains a perplexity of 16.65, whereas baselines degrade drastically (e.g., LLM-Pruner: 2724.8).
- It maintains competitive average zero-shot accuracy across five benchmarks (ARC, HellaSwag, PIQA, etc.) even at high sparsity.
Efficiency and Scalability:
- Training Time: The entire process (4 epochs of gate learning + calibration) takes approximately 6 minutes on a single NVIDIA A100 80GB GPU.
- Memory & Throughput: Pruning KV heads reduces the KV-cache memory and peak prefill memory nearly proportionally to the parameter reduction, enabling longer context lengths or larger batch sizes. Throughput (tokens/sec) improves significantly, especially at larger batch sizes.
Generalization:
- The method generalizes well to newer model families, including LLaMA-3.1-8B, Qwen3-8B, and Qwen3-14B, consistently outperforming FLAP in average downstream accuracy.

5. Significance

GRASPrune addresses a critical gap in LLM compression by demonstrating that global, budget-aware gating is superior to layer-wise or decoupled pruning strategies.

Practical Deployment: It produces a standard dense model that can be served with existing infrastructure without custom kernels or extra parameters.
Cost Efficiency: It eliminates the need for expensive full-model fine-tuning or knowledge distillation, making structured pruning accessible and efficient.
KV Cache Optimization: By explicitly pruning KV heads within the global budget, it directly addresses the memory bottleneck associated with long-context inference, a limitation often ignored by methods focusing solely on parameter count.

In summary, GRASPrune offers a simple, fast, and highly effective mechanism for compressing LLMs while preserving core language modeling capabilities, particularly in high-compression regimes.