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

2026-04-21

AI Summary

This work addresses the high memory and latency overhead of large language model inference by proposing a post-pretraining structured pruning framework that jointly prunes FFN channels and KV head groups under a unified global budget. The method employs a lightweight gating mechanism with a projected straight-through estimator to learn pruning scores, while enforcing hard budget constraints and freezing backbone weights during each optimization step. Subsequently, a fine-tuning-free calibration procedure fuses scaling factors to correct magnitude mismatches, introducing zero additional inference cost. Evaluated on LLaMA-2-7B with 50% parameter reduction, the approach achieves a WikiText-2 perplexity of 12.18 and maintains competitive performance across five zero-shot tasks, requiring only a single A100 80GB GPU, 512 unlabeled samples, and four calibration rounds.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

Large language models (LLMs) are expensive to serve because model parameters, attention computation, and KV caches impose substantial memory and latency costs. We present GRASPrune, a structured pruning framework applied after pretraining that jointly prunes FFN channels and KV head groups under a single global budget. Instead of learning importance scores without constraints and applying the budget only after training, GRASPrune learns lightweight gate scores with a projected straight-through estimator that enforces a hard mask satisfying the budget at every step while keeping the backbone weights frozen. After the mask is fixed, we calibrate scaling factors on the retained units to mitigate scale mismatch caused by pruning, and fold these factors into the pruned weights to obtain a smaller dense checkpoint with no extra parameters at inference. On LLaMA-2-7B, GRASPrune removes 50% of parameters and achieves 12.18 perplexity on WikiText-2 while maintaining competitive average zero-shot accuracy on five benchmarks, using four epochs on 512 unlabeled calibration sequences on a single NVIDIA A100 80GB GPU without any full model fine-tuning.

Problem

• structured pruning
• large language models
• budgeted compression
• KV cache
• model serving cost

Innovation

• structured pruning
• global budget
• straight-through estimator
• scale calibration
• LLM compression