Research map / Large language models / Language, safety and interpretability
MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models
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.
Learns 2:4 semi-structured sparsity masks for frozen LLMs by modeling mask choice as a learnable categorical distribution sampled with Gumbel Softmax and trained end-to-end on large data, with optional mask priors for transfer.
Problem
Finding accurate N:M (2:4) semi-structured sparsity masks for LLMs is combinatorially hard and non-differentiable, and one-shot methods rely on small calibration sets and proxy importance criteria that do not scale to large data.
Approach
Model each group of four weights' 2:4 mask as a learnable categorical distribution over the six valid masks, sample differentiably with Gumbel Softmax, and optimize the language modeling loss end-to-end with frozen weights; initialize logits from one-shot prior masks and add a sparse weight regularizer to keep gradients flowing.
Key results
- On LLaMA-2 7B with 2:4 sparsity and frozen weights, MaskLLM reaches 6.72 Wikitext PPL versus 10.42 for SparseGPT (with weight update) and dense 5.12.
- On GPT-3 2B and Nemotron-4 15B, MaskLLM attains Wikitext PPL of 11.42 and 7.31 respectively, versus SparseGPT's 22.14 and 13.38.
- Learned masks for downstream domains reach lossless PPL (7.39 vs dense 7.42 average) and yield 0.65 bits/param storage per task with 1.4x speedup on an A6000 GPU for LLaMA-2 7B.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MaskLLM instead of SparseGPT for 2:4 pruning | LLaMA-2 7B · Wikitext-2 | PPL: 10.42 -> 6.72 | Helped |
| Use MaskLLM instead of the strongest listed one-shot baseline | LLaMA-2 13B · Wikitext-2 | PPL: 8.47 -> 5.85 | Helped |
| Use MaskLLM instead of SparseGPT | Nemotron-4 15B · Wikitext-2 | PPL: 13.38 -> 7.31 | Helped |
| Use MaskLLM instead of SparseGPT | GPT-3 2B · Wikitext-2 | PPL: 22.14 -> 11.42 | Helped |
| Use MaskLLM instead of SparseGPT | GPT-3 843M · Wikitext-2 | PPL: 38.78 -> 15.39 | Helped |
| Use MaskLLM instead of SparseGPT | LLaMA-3 8B · Wikitext-2 | PPL: 17.64 -> 8.50 | Helped |
| Initialize mask learning with a Magnitude prior instead of no prior (ablation) | MaskLLM on LLaMA-2 7B · Wikitext | PPL: 9.12 -> 6.77 | Helped |
| Transfer the general mask as a prior to the downstream task (ablation) | GPT-3 2B · Downstream tasks | Average task PPL: 7.51 -> 7.39 | Helped |
| Use a task-specific general mask rather than the dense model | GPT-3 2B · Downstream tasks | Average task PPL: 7.42 -> 10.61 | Hurt |
| Add sparse weight regularization (1e-5) (ablation) | MaskLLM on GPT-3 2B · Mask training | Average gradient norm: 0.219 -> 0.542 | Helped |
| Add sparse weight regularization (ablation) | MaskLLM on GPT-3 2B · Mask-only, sub-domain, and finetuning | PPL: 11.59 -> 11.42; 7.61 -> 7.39; 10.21 -> 9.96 | Helped |
| Increase Gumbel temperature from 1 → 0.05 to 4 → 0.05 (ablation) | MaskLLM on GPT-3 843M · Wikitext | PPL: 17.52 -> 15.39 | Helped |
Surprising findings
- Without a prior, the learned mask is still better than one-shot baselines, but magnitude-style priors are needed to reach best results in 2,000 steps.
- Weight regularization is needed to prevent vanishing gradients; without it, mask learning barely moves from the prior (only 2.83% difference).
Limitations
- Learnable mask training consumes substantially more compute than one-shot methods.
- Evaluated on a limited set of models; internal GPT-3 models are not publicly available.
- No error bars or statistical significance reported.
Experimental setup
Datasets: C4 Wikitext-2 LM-Eval-Harness tasks ImageNet-1K Blended pretraining-style domain data
Benchmarks: Wikitext-2 HellaSwag RACE PIQA WinoGrande ARC-E ARC-C OBQA ImageNet-1K
Metrics: perplexity accuracy top-1 accuracy throughput
Baselines: Magnitude pruning SparseGPT Wanda SR-STE ADMM-Iter GBLM RIA Pruner-Zero
Builds on: Gumbel Softmax SparseGPT Wanda Megatron-LM
Closest papers by meaning
- SparseLLM: Towards Global Pruning of Pre-trained Language Models (NeurIPS 2024 · similarity 0.93)
- MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs (ICLR 2026 · similarity 0.92)
- Adaptive Layer Sparsity for Large Language Models via Activation Correlation Assessment (NeurIPS 2024 · similarity 0.92)
- Pruning Large Language Models with Semi-Structural Adaptive Sparse Training (AAAI 2025 · similarity 0.91)
- Determining Layer-wise Sparsity for Large Language Models Through a Theoretical Perspective (ICML 2025 · similarity 0.91)
- Discovering Sparsity Allocation for Layer-wise Pruning of Large Language Models (NeurIPS 2024 · similarity 0.91)
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.