GRASPrune: 대규모 언어 모델의 예산 기반 구조적 가지치기를 위한 글로벌 게이팅

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

대규모 언어 모델(LLM)은 모델 파라미터, 어텐션 계산, KV 캐시 등으로 인해 상당한 메모리 및 지연 시간 비용이 발생하여 서비스 운영 비용이 높습니다. 본 논문에서는 사전 학습 후에 적용되는 구조적 가지치기 프레임워크인 GRASPrune을 제시합니다. GRASPrune은 전체 예산을 사용하여 FFN 채널과 KV 헤드 그룹을 동시에 가지치기합니다. 기존 방법과는 달리, GRASPrune은 제약 조건 없이 중요도 점수를 학습하고, 학습 후에만 예산을 적용하는 대신, 투영된 스트레이트-쓰루 추정기를 사용하여 전체 예산을 만족하는 하드 마스크를 모든 단계에서 적용하면서 기반 가중치를 고정합니다. 마스크가 고정되면, 가지치기로 인해 발생하는 스케일 불일치를 완화하기 위해 유지된 유닛에 대한 스케일링 팩터를 조정하고, 이러한 팩터를 가지치기된 가중치에 통합하여 추가 파라미터 없이 더 작은 밀집 체크포인트를 얻습니다. LLaMA-2-7B 모델에서 GRASPrune은 50%의 파라미터를 제거하고 WikiText-2 데이터셋에서 12.18의 퍼플렉시티를 달성했으며, 5개의 벤치마크에서 경쟁력 있는 평균 제로샷 정확도를 유지합니다. 이는 단일 NVIDIA A100 80GB GPU에서 512 길이의 라벨링되지 않은 캘리브레이션 시퀀스를 사용하여 4 에포크 동안 수행되었으며, 전체 모델 미세 조정 없이 가능했습니다.

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.