Reproduce “MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models” on OpenResearch

About this paper

Gongfan Fang, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jeff Pool, Jan Kautz, Pavlo Molchanov, Xinchao Wang

MaskLLM introduces a learnable framework for N:M semi-structured sparsity in large language models by framing mask selection as a probabilistic, end-to-end trainable problem. This approach, from NVIDIA and NUS, reduces perplexity by a large margin compared to one-shot methods (e.g., LLaMA-2 7B PPL improved from 10.42 to 6.72) and enables lossless task-specific compression, resulting in a 1.4x GPU speed-up and 73% memory reduction.

Key insights

• Learnable sparsity effectively scales with large datasets, yielding higher quality masks that continue to improve with more training data.
• Initializing the mask learning process with pre-computed masks from one-shot methods significantly improves training efficiency and final mask quality.
• The balance between stochastic exploration (controlled by the Gumbel scaling factor "kappa") and convergence speed is crucial for optimal mask learning.
• Maintaining a large magnitude for the remaining active weights through regularization is vital for robust mask learning and preserving downstream task performance.

Summary by alphaXiv.