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

TL;DR

MaskPro introduces a probabilistic learning framework for efficient (N:M)-sparsity in large language models, improving inference speed and memory use with robust training methods.

Contribution

It proposes a novel linear-space probabilistic approach for (N:M)-sparsity, addressing errors and training costs of prior methods.

Findings

MaskPro achieves superior sparsity performance in LLMs.

The method demonstrates excellent scalability and robustness.

Extensive experiments validate its effectiveness.

Abstract

The rapid scaling of large language models~(LLMs) has made inference efficiency a primary bottleneck in the practical deployment. To address this, semi-structured sparsity offers a promising solution by strategically retaining $N$ elements out of every $M$ weights, thereby enabling hardware-friendly acceleration and reduced memory. However, existing (N:M)-compatible approaches typically fall into two categories: rule-based layerwise greedy search, which suffers from considerable errors, and gradient-driven combinatorial learning, which incurs prohibitive training costs. To tackle these challenges, we propose a novel linear-space probabilistic framework named MaskPro, which aims to learn a prior categorical distribution for every $M$ consecutive weights and subsequently leverages this distribution to generate the (N:M)-sparsity throughout an $N$-way sampling without replacement. Furthermore, to mitigate the training instability induced by the high variance of policy gradients in the super large combinatorial space, we propose a novel update method by introducing a moving average tracker of loss residuals instead of vanilla loss. Finally, we conduct comprehensive theoretical analysis and extensive experiments to validate the superior performance of MaskPro, as well as its excellent scalability in memory efficiency and exceptional robustness to data samples. Our code is available at \href{https://github.com/woodenchild95/Maskpro.git}{\ttfamily https://github.com/woodenchild95/Maskpro.git}.