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

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

Problem Statement
The rapid scaling of Large Language Models (LLMs) has created a significant bottleneck in inference efficiency, necessitating effective compression techniques. Semi-structured sparsity, specifically the (N:M) pattern where $N$ elements are retained out of every $M$ consecutive weights, offers a hardware-friendly solution for acceleration and memory reduction. However, identifying the optimal (N:M) mask is a formidable combinatorial optimization problem. Existing approaches fall into two categories with distinct limitations:

1. Rule-based methods: These employ greedy, layer-wise searches using handcrafted metrics (e.g., minimizing layer-wise reconstruction error). While computationally cheap, they suffer from inherent biases and considerable gaps with end-to-end task performance.
2. Learning-based methods (e.g., MaskLLM): These optimize masks directly via gradient descent on a generation loss. While effective, they incur prohibitive training costs. Specifically, MaskLLM requires storing logits for all possible masks within a group of $M$ weights, leading to a memory complexity of $O(\binom{M}{N} \frac{d}{M})$. For large $M$, this scales exponentially, making it infeasible for large models. Furthermore, the vast combinatorial space induces high variance in policy gradients, causing training instability.

Methodology: MaskPro
The authors propose MaskPro, a novel linear-space probabilistic framework designed to learn (N:M)-sparsity efficiently and stably.

1. Linear-Space Probabilistic Relaxation:
   Instead of defining a categorical distribution over all $\binom{M}{N}$ possible masks for a group of $M$ weights, MaskPro reformulates the problem. It establishes a categorical distribution over the $M$ individual elements within each group. The (N:M) mask is then generated by performing an N-way sampling without replacement from this distribution.

   • Theoretical Basis: The authors prove that any (N:M) mask can be represented as a probabilistic sum of $N$ distinct basis vectors.
   • Memory Efficiency: This reformulation reduces the number of learnable parameters (logits) from $O(\binom{M}{N} \frac{d}{M})$ to a linear $O(d)$, where $d$ is the total number of weights. This eliminates the memory explosion associated with previous learning-based methods.

2. Refined Policy Gradient Estimator (PGE):
   To optimize the logits, MaskPro utilizes policy gradients. The authors identify that standard policy gradients using raw loss values are unstable in LLMs due to the high variance caused by minibatch differences (where loss variation from data sampling often dwarfs loss variation from mask sampling).

   • Loss Residuals: Instead of using the raw loss $f(m \odot w, \xi)$, MaskPro updates logits based on the loss residual relative to an initial baseline mask $m_0$: $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$. This isolates the impact of the mask change.
   • Moving Average Tracker: To further stabilize training and prevent numerical explosions, a moving average tracker $\delta$ is introduced to smooth the baseline. The update rule becomes:
     $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
   • Theoretical Guarantees: The paper provides proofs showing that these refined estimators remain unbiased estimators of the true policy gradient. Furthermore, under the condition that the sampled mask's loss is less than twice the baseline loss (a condition easily met in practice), the variance of the refined estimator is strictly lower than that of the vanilla policy gradient.

Key Contributions

1. Linear-Space Framework: MaskPro reduces the memory requirement for learning (N:M) sparsity from exponential $O(\binom{M}{N} \frac{d}{M})$ to linear $O(d)$ by modeling the selection process as N-way sampling without replacement from a categorical distribution over $M$ elements.
2. Enhanced Policy Gradient: The authors introduce a variance-reduced update mechanism using loss residuals and a moving average tracker, addressing the training instability caused by the high variance of policy gradients in large combinatorial spaces.
3. Theoretical Analysis: Comprehensive proofs are provided regarding the unbiasedness of the estimators and the variance reduction properties of the proposed update rules.

Experimental Results
Extensive experiments were conducted on multiple LLMs (Gemma-7B, Vicuna-7B, LLaMA-2-7B, DeepSeek-7B) across various downstream tasks (e.g., HellaSwag, RACE, PIQA) and sparsity patterns ((2:4), (4:8), (8:16)).

• Performance: MaskPro consistently outperforms state-of-the-art rule-based methods (SparseGPT, Wanda, Pruner-Zero) and achieves performance comparable to the memory-intensive MaskLLM, often with an average improvement of over 2% in accuracy.
• Memory Scalability: While MaskLLM encounters Out-of-Memory (OOM) errors on 8x A100 GPUs for (4:8) sparsity on 7B models, MaskPro maintains stable memory usage (~36-37 GB) regardless of the sparsity pattern.
• Data Efficiency: MaskPro demonstrates remarkable robustness to data volume. It achieves stable performance even when trained with only 1 data sample, whereas MaskLLM requires hundreds of thousands of samples to converge.
• Training Efficiency: The method requires significantly fewer training samples and less training time compared to MaskLLM, while maintaining memory usage comparable to rule-based methods.

Significance and Claims
The paper claims that MaskPro provides a viable path for the practical, scalable local deployment of LLMs by solving the memory and stability bottlenecks of learning-based semi-structured sparsity. The authors emphasize that their approach bridges the gap between the high performance of learning-based methods and the efficiency of rule-based methods. By reducing the memory footprint to linear space and stabilizing the training process through refined policy gradients, MaskPro enables the learning of optimal (N:M) masks on large-scale models without the prohibitive computational overhead previously associated with such tasks. The method is presented as a generalizable framework for probabilistic mask learning that is robust to data scarcity and model scale.