Research map / Architectures and efficiency / Pruning and efficient tuning
Pruning Large Language Models with Semi-Structural Adaptive Sparse Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Adaptive Sparse Trainer (AST), a retraining framework that learns N:M masks on the fly with annealed SR-STE decay, KL distillation from the dense teacher, and optional low-rank SLoRB adapters, producing 2:4 sparse LLMs with small perplexity and zero-shot gaps.
Problem
One-shot 2:4 semi-structured pruning of LLMs causes large performance degradation, especially on knowledge-intensive and complex reasoning tasks, and retraining pruned billion-parameter models under N:M constraints is under-explored and computationally costly.
Approach
Keep all weights trainable, recompute magnitude-based N:M masks periodically, and gradually decay masked weights with an annealed L2 schedule (Annealing SR-STE), combined with logit-level KL distillation from the dense model and a low-rank SLoRB adapter initialized from pruned-weight group means.
Key results
- On LLaMA2-7B 2:4 sparse, AST-Boosted reaches WikiText perplexity 5.69 vs dense 5.12 (Wanda 11.02), with MMLU 38.2 vs dense 45.3.
- On GPT2-124M, AST-Boosted achieves perplexity 31.13 vs dense 29.95 and SparseGPT 50.09 at 2:4 sparsity.
- TensorRT-LLM end-to-end decoding speedup of the LLaMA2-7B 2:4 sparse model is 1.33x on RTX4090 and 1.78x to 1.83x on L20.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add knowledge distillation during sparse retraining (ablation) | AST-Naive · WikiText-2, GPT2/OPT | perplexity: 40.34 -> 39.29; 29.79 -> 29.08; 28.65 -> 27.21; 39.46 -> 36.97 | Helped |
| use annealing instead of static SR-STE decay (ablation) | AST-Naive · WikiText-2, GPT2/OPT | perplexity: 32.84 -> 32.23; 24.04 -> 23.65; 21.73 -> 21.29; 31.08 -> 30.22 | Helped |
| allow the mask to adapt during training instead of fixing it (ablation) | AST-Naive · WikiText-2, GPT2/OPT | perplexity: 32.93 -> 32.23; 24.18 -> 23.65; 21.95 -> 21.29; 31.06 -> 30.22 | Helped |
| use intermediate attention and hidden-state distillation (ablation) | sparse GPT2 retraining · WikiText-2, GPT2-124M | perplexity: 32.23 -> 42.75 | Hurt |
| use MobileBERT feature-map and attention distillation (ablation) | sparse GPT2 retraining · WikiText-2, GPT2-124M | perplexity: 32.23 -> 44.87 | Hurt |
| add intermediate representation loss to KL distillation (ablation) | sparse GPT2 retraining · WikiText-2, GPT2-124M | perplexity: 32.23 -> 41.19 | Hurt |
| replace forward KL with reverse KL (ablation) | sparse GPT2 retraining · WikiText-2, GPT2-124M | perplexity: 32.23 -> 32.20 | No clear effect |
| initialize SLoRB from the mean of pruned weights rather than zero (ablation) | SLoRB · WikiText-2, GPT2 | perplexity: 32.46 -> 31.13 | Helped |
| train SLoRB together with sparse model weights instead of adding it after retraining (ablation) | AST-Boosted · WikiText-2, GPT2 | perplexity: 32.07 -> 31.13 | Helped |
| use SLoRB parameters (ablation) | AST-Boosted · WikiText-2, GPT2 | perplexity: 32.23 -> 31.13 | Helped |
| tune the KD loss weight to α=2/3 instead of α=0 (ablation) | LLaMA2-7B retraining · WikiText-2, 3B tokens | perplexity: 6.26 -> 6.22 | Helped |
| use only KD loss (α=1) instead of α=0 (ablation) | GPT2 retraining · WikiText-2, GPT2, 5B tokens | perplexity: 32.18 -> 40.34 | Hurt |
Surprising findings
- Training-free methods like Wanda degrade severely on GPT2 (perplexity 115.64 at 124M), and distillation from intermediate hidden states or attention hurts generative LM retraining.
- Reverse KL (MiniLLM) performs similarly to forward KL on logit distillation in this setting.
Limitations
- Retraining was limited to a small number of tokens due to compute constraints, so larger models and longer training are unexplored.
- SLoRB adds about 12.5% extra parameters for k=16 and increases memory use.
- Training-based baselines (IMP, gradual pruning) were reported only for GPT2 and OPT, not LLaMA2-7B.
Experimental setup
Datasets: RedPajama-v1 C4 WikiText-2
Benchmarks: WikiText-2 ARC-e ARC-c OpenBookQA WinoGrande RTE HellaSwag BoolQ MMLU GSM8K MATH
Metrics: perplexity accuracy throughput (tokens/s)
Baselines: SparseGPT Wanda Iterative Magnitude Pruning Gradual Magnitude Pruning Sheared-LLaMA-2.7B TinyBERT-style distillation MiniLLM
Builds on: SR-STE SparseGPT Wanda LoRA AWQ
Closest papers by meaning
- SparseLLM: Towards Global Pruning of Pre-trained Language Models (NeurIPS 2024 · similarity 0.93)
- Compact Language Models via Pruning and Knowledge Distillation (NeurIPS 2024 · similarity 0.93)
- The Journey Matters: Average Parameter Count over Pre-training Unifies Sparse and Dense Scaling Laws (ICLR 2025 · similarity 0.92)
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.92)
- SlimGPT: Layer-wise Structured Pruning for Large Language Models (NeurIPS 2024 · similarity 0.92)
- Adaptive Layer Sparsity for Large Language Models via Activation Correlation Assessment (NeurIPS 2024 · similarity 0.92)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.