Research map / Large language models / Reasoning and chain of thought

MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

ICLR 2026 · Large language models › Reasoning and chain of thought · approach: Efficient LLM serving

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 MaskPro, a linear-space probabilistic framework that learns (N:M) sparse masks by per-group categorical distributions with N-way sampling without replacement, trained via a policy gradient using a moving-average-tracked loss residual.

Problem

Finding strict (N:M) semi-structured sparse masks for LLMs: rule-based layerwise greedy methods have errors from handcrafted metrics, while learning-based MaskLLM has prohibitive memory (combinatorial logits) and training cost.

Approach

Model each group of M weights with a single categorical distribution (d logits total) and sample N masked positions by N-way sampling without replacement; train with a policy gradient that uses loss residuals relative to an initial mask, plus a moving-average baseline to reduce variance.

Key results

• On (2:4)-sparsity LLaMA-2-7B, MaskPro reaches WikiText PPL 17.17 vs 21.07 for SparseGPT (dense 8.71) and average zero-shot accuracy among the best of the backprop-free methods.
• Memory for DeepSeek-7B (4:8) training is 36.95 GB for MaskPro vs >640 GB OOM for MaskLLM; (2:4) is 36.82 GB vs 339.56 GB for MaskLLM.
• MaskPro trained with only 1 C4 sample achieves avg zero-shot 49.18 on LLaMA-2-7B vs 49.62 with 320k samples.

Surprising findings

• Vanilla policy gradient updates were almost ineffective, with loss oscillating around zero; only the loss-residual with moving-average tracker trained stably.
• Training with a single calibration sample remains stable, with only slight performance loss.

Limitations

• Training time on large models is dominated by simulating the mask sampling process.
• Initialization of logits is critical; zero or random initialization fails, requiring a well-chosen initial mask and magnitude C.
• Gemma-7B weights are not sufficiently sparse, giving suboptimal and unstable results.

Experimental setup

Datasets: C4

Benchmarks: WikiText perplexity HellaSwag RACE PIQA WinoGrande ARC-E ARC-C OBQA WMDP BoolQ COPA

Metrics: perplexity zero-shot accuracy memory (GB)

Baselines: Magnitude SparseGPT Wanda GBLM-Pruner Pruner-Zero MaskLLM

Builds on: MaskLLM Gumbel-Max Wanda SparseGPT

Closest papers by meaning

• MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models (NeurIPS 2024 · similarity 0.92)
• Determining Layer-wise Sparsity for Large Language Models Through a Theoretical Perspective (ICML 2025 · similarity 0.89)
• Discovering Sparsity Allocation for Layer-wise Pruning of Large Language Models (NeurIPS 2024 · similarity 0.89)
• PROXSPARSE: REGULARIZED LEARNING OF SEMI-STRUCTURED SPARSITY MASKS FOR PRETRAINED LLMS (ICML 2025 · similarity 0.89)
• R-Sparse: Rank-Aware Activation Sparsity for Efficient LLM Inference (ICLR 2025 · similarity 0.89)
• Adaptive Layer Sparsity for Large Language Models via Activation Correlation Assessment (NeurIPS 2024 · similarity 0.89)