MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

Authors: gongfan  fang
hongxu  yin
saurav  muralidharan
greg  heinrich
jeff  pool
jan  kautz
pavlo  molchanov
xinchao  wang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-02 17:00:22

Abstract: Large Language Models (LLMs) are distinguished by their massive parameter counts, which typically result in significant redundancy. This work introduces MaskLLM, a learnable pruning method that establishes Semi-structured (or ``N:M'') Sparsity in LLMs, aimed at reducing computational overhead during inference. Instead of developing a new importance criterion, MaskLLM explicitly models N:M patterns as a learnable distribution through Gumbel Softmax sampling. This approach facilitates end-to-end training on large-scale datasets and offers two notable advantages: 1) High-quality Masks - our method effectively scales to large datasets and learns accurate masks; 2) Transferability - the probabilistic modeling of mask distribution enables the transfer learning of sparsity across domains or tasks. We assessed MaskLLM using 2:4 sparsity on various LLMs, including LLaMA-2, Nemotron-4, and GPT-3, with sizes ranging from 843M to 15B parameters, and our empirical results show substantial improvements over state-of-the-art methods. For instance, leading approaches achieve a perplexity (PPL) of 10 or greater on Wikitext compared to the dense model's 5.12 PPL, but MaskLLM achieves a significantly lower 6.72 PPL solely by learning the masks with frozen weights. Furthermore, MaskLLM's learnable nature allows customized masks for lossless application of 2:4 sparsity to downstream tasks or domains. Code is available at https://github.com/NVlabs/MaskLLM.

Main paper: 
 introduction: 
 Large Language Models (LLMs) have demonstrated remarkable effectiveness across a diverse range of tasks [19, 6, 48, 13] . However, the generality and robustness of LLMs are largely attributed to their vast scale, with parameter counts ranging from one billion to several hundred billion [39, 45, 5] . This substantial model size, in turn, makes it challenging and resource-intensive to deploy LLMs in real-world applications. One effective and practical approach to address this issue is semi-structured pruning [29, 32, 12, 38] , which introduces N:M sparsity into LLMs to improve both memory and computational efficiency. The N:M pattern, with N non-zero values among M consecutive parameters, is typically hardware-friendly to accelerators like GPUs and has thus garnered considerable attention [12, 38, 32] .Despite the simplicity of its core idea, semi-structured pruning still presents considerable challenges within the realm of LLMs. Sparsity aims to identify a subset of parameters that attain a comparable quality to the dense model. Nevertheless, the extensive parameter scale of large language models usually leads to a vast search space. In a fully sparsified LLaMA2-7B model with 2:4 sparsity, for instance, there are 1.6 billion 2:4 masks to be chosen for dense layers. This makes the combinatorial problem of finding the optimal mask set exceedingly challenging. In the literature, leading approaches such as SparseGPT [12] and Wanda [38] , utilize a small calibration set and carefully designed importance criteria to identify redundant parameters. While these techniques have demonstrated remarkable results on several large language models, two substantial challenges remain: Firstly, the small calibration set is insufficient to represent the comprehensive knowledge embedded in LLMs, which are pre-trained on extensive and diverse data domains [39, 7, 31] . As demonstrated in our experiments, hand-crafted importance criteria are only applicable to a compact subset of data, and enlarging the calibration set beyond 256 entries does not improve the resulting quality. This limits the generalizability of pruned LLMs in different domains. Secondly, using handcrafted criteria as a proxy for the true discrepancy inevitably results in errors. A considerable gap remains between the real discrepancy induced by pruning and existing importance indicators, such as gradient information [30] , weight magnitude [15] , and the Hessian Matrix [23, 12] .To tackle the outlined challenges, we propose MaskLLM, a learnable method that facilitates end-toend training of LLM sparsity on large-scale datasets. In the context of N:M sparsity, pruning an LLM involves selecting masks from a discrete and finite set. However, the non-differentiability of mask selection and combination hinders the direct use of backpropagation for mask learning. To address this, our work frames the mask selection problem from a probabilistic perspective, associating each candidate mask with a probability and modeling the mask selection as a stochastic sampling process. We incorporate the Gumbel Softmax [21] for differentiable sampling, which re-parameterizes the randomness of sampling into an independent random variable. This makes the probabilities of each mask candidate optimizable with gradient descent. During training, MaskLLM aims to learn appropriate mask distributions, from which the sampled masks can preserve the original quality of dense LLMs. The differentiable mask offers two advantages in addressing the challenges mentioned above: (1) it effectively scales to large-scale datasets, thereby preserving the rich knowledge in LLMs, and (2) the end-to-end training explicitly optimizes the language modeling loss of LLMs, which exactly measures the discrepancy induced by pruning. Furthermore, inspired by the power of transfer learning, we introduce prior masks, a simple strategy to fully leverage pre-computed masks and enable fast transfer learning of sparsity across domains and tasks as illustrated in Figure 1 To evaluate the proposed method, we conduct experiments on several LLMs including LLaMA-2 7B, LLaMA-2 13B [39] , Nemotron-4 15B [31] , and two in-house LLMs, GPT-3 843M and GPT-3 2B pre-trained using Megatron framework [36] . Our method can learn high-quality masks for pruning through end-to-end training on large-scale datasets. For example, compared to SparseGPT which archives a perplexity (PPL) of 10.42 on LLaMA-2 7B, our method improves the PPL to 6.72, without any update to the LLM parameters. Besides, our method facilitates the learning of domain-specific masks, which can even achieve lossless compression of LLMs on some downstream tasks or domains.The principal contribution of this work lies in a learnable method for semi-structured pruning of LLMs. MaskLLM is designed to fully harness large-scale datasets to learn accurate masks, applicable to both general-purpose and domain-specific pruning. Additionally, the framework facilitates the transfer learning of sparsity patterns across different tasks, enabling efficient training of sparsity. 
related works: 
 Pruning Large Language Models. Network Pruning [15, 30, 17, 18, 41] have been proven an efficient approach to compress pre-trained language models via the removal of redundant parameters.According to the granularity of pruning, existing methods can be classified into three categories: Structured Pruning [26, 43, 24] , Unstructured Pruning [17, 15] , and Semi-Structured Pruning [12, 38, 29, 32, 33] . Structured pruning physically eliminates substructures like attention heads [26] , embeddings or depth [43] in the model, facilitating acceleration independent of specialized hardware or software infrastructure [32] . However, structured approaches typically necessitate huge retraining efforts to recover network quality due to coarse removal of parameters [26, 43, 27, 2] . Conversely, unstructured methods aim to find a sparse model by zeroing out parameters in LLMs, which is characterized by its flexibility and minimal detrimental effect on LLMs' accuracy [12, 38, 20, 42, 44] . The acceleration of sparse models is typically impeded by the irregular nature of the resulting sparse patterns, presenting challenges in achieving computational efficiency. Positioned between structured and unstructured methods, the semi-structured approach introduces hardware-friendly patterns such as N:M sparsity, which leaves only N nonzero values in each group of M values and thereby harmonizes the acceleration benefits of a structured pattern with the flexibility of fine-grained sparsity [32, 33, 12] .In this study, we focus on N:M semi-structured sparsity within Large Language Models and present a learnable framework to obtain high-quality masks via end-to-end training.Learnable Semi-Structured Sparsity. On another hand, a burgeoning interest exists in developing learnable masks [49, 25, 47] , especially in the field of vision models. Markedly contrasted with traditional one-shot pruning methods that rely on a predetermined metric of importance, learnable sparsity can fully leverage the rich information in training data, enabling the identification of more effective sparsity masks. A particularly popular strategy is to directly update the network weight, such as pushing partial weights to zero with Sparse-Refined Straight-Through Estimator (SR-STE) [3, 17, 25] or permuting parameters to achieve better quality [33] . Other methods learn additional indicators to reveal the importance of weight, such as differentiable indexing [35] , optimizable combination [47] ,or decaying [22] . In this work, we make the first attempt to learn N:M masks for frozen LLMs, which is much more challenging due to the huge parameter amount and problem scale.3 MethodQuestion: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: The limitation of this submission was discussed in section K. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
n:m sparsity: 
 We motivate and introduce a learnable framework, MaskLLM, to sparsify Large Language Models (LLMs) for improved inference efficiency. Sparsifying an LLM with N:M patterns imposes the constraint of having (no more than) N non-zero values within each consecutive set of M parameters. This task can be formulated as a mask selection problem with the candidate set of|S| = M N = M ! N !(M -N )! candidates,where |S| denotes the size of the candidate set, and M N represents the combination number of potential N:M masks. For simplicity, this work primarily focuses on 2:4 sparsity, which can be naturally extended to other patterns such as 1:4 and 4:8. Given a parameter block comprising four consecutive parameters, denoted as W ∈ R 1×4 , the goal of sparsification is to identify the optimal binary mask M * ∈ B 1×4 of the same size, ensuring that the pruned weight maintains its behavior on observed data x ∼ p(x). For 2:4 sparsity, the binary mask M must contain exactly two zeros, resulting in a discrete candidate set S 2:4 with |S 2:4 | = 4 2 = 6 candidates:EQUATIONFor an LLM, there exists a substantial number of parameter blocks, denoted as {W i }, each requiring the selection of corresponding masks {M i }. To maintain satisfactory behavior after pruning, it is natural to define the following objective for N:M sparsity:EQUATIONwhere L LM refers to the language modeling loss for pre-training. The operator ⊙ denotes elementwise multiplication, which masks partial parameters for sparsification. However, finding the optimal combination of masks M * can be extremely challenging in the context of LLMs due to the nondifferentiable nature of mask selection and the huge parameter scale. In the following sections, we demonstrate that the mask selection can be transformed into a sampling process.Consider a single parameter block W ∈ R 1×4 consisting of only 4 parameters: directly determining the exact optimal mask for this block is not feasible, since the behavior of the pruned LLM also depends on the pruning of other parameter blocks. Nevertheless, it remains feasible to sample masks independently for each block and assess the overall model quality after pruning. To facilitate random sampling of M, we define a categorical distribution with class probability p 1 , p 2 , . . . p |S| , which satisfy j p j = 1. During the random sampling phase, if a certain mask achieves good quality during pruning, it's reasonable to adjust the categorical distribution by increasing the probability of the sampled mask. With sufficient sampling and updates, this process ends with a distribution where the mask with high probability is more likely to maintain good quality after pruning. Formally, we model the combination problem in Equation 3 from the perspective of random sampling:EQUATIONwhere p(M i ) refers to the categorical distribution of i-th mask M i . If it is feasible to get the gradient w.r.t. the distribution, then the above objective can be optimized with gradient descent as demonstrated in Figure 2 . Nonetheless, drawing samples from a categorical distribution is still non-differentiable.Differentiable Sampling of Masks An effective method to model a sampling operation is Gumbel Max [14] , a re-parameterization trick that disentangles the randomness of sampling into a noise variable. This trick introduces a method to draw samples from the categorical distribution p with an additional noise variable ϵ. It produces the one-hot index y for sampling:EQUATIONwhere ϵ i is a random noise following uniform distribution, and the g i = -log(-log ϵ i ) is known as the Gumbel noise. With the Gumbel Max trick, the randomness of sampling is parameterized to an independent variable g i . The only issue towards differentiable sampling lies in the argmax and onehot operation. To address this, we leverage Gumbel Softmax [21] to approximate the index with Softmax, leading to a soft and differentiable index ỹ = [ỹ 1 , ỹ2 , . . . , ỹ|S| ]:EQUATION)The temperature term τ is a hyper-parameter, controlling the hardness of the sampled index. While τ → 0, the soft index will be more close to a one-hot vector, resulting in ỹi → y i . With the soft index ỹ as a row vector and the mask set S as a matrix where each row i refers to the i-th candidate mask Mi , it's easy to craft a differentiable mask through a simple matrix multiplication:EQUATIONThis operation produces a weighted average of candidate masks according to the soft index. As shown in Figure 3 , we can find all operations, including the sampling and weighted averaging are differentiable, and the gradient w.r.t. the probability p can be easily computed. This allows using the differentiable mask M to optimize the sampling problem defined in Equation 4.Mask Set 𝒮 !:#1 1 0 0 1 0 1 0 1 0 0 1 0 1 1 0 0 1 0 1 0 0 1 1Logits 𝜋 (Learnable)𝑔 ~𝐺𝑢𝑚𝑏𝑒𝑙(0,1)Sampled Index 𝑦 Learning Masks for LLMs Equation 7 provides a differentible mask sampled from the underlying distribution p. The gradient flow can easily reach the probability p i , making it an optimizable variable in the system. Typically, we do not directly learn the probability and instead, learn the logits π i with a scaling factor κ, which produces the probability as p i = exp(πi•κ) j exp(πj •κ) . As will be discussed in Section 4.3, the scaling factor κ will be used to balance the relative magnitude of logits and Gumbel noises, which controls the randomness of sampling. During training, all parameter blocks {W i } are associated with the corresponding distributions {p π (M i )}, and optimal distribution can be learned in an end-to-end manner. However, our empirical experiments on several large language models reveal a new issue with the learnable masks: the gradient may vanish due to the pruning operation that produces zero parameters in the network. This issue will adversely affect downstream transfer and fine-tuning. To address this, we introduce Sparse Weight Regularization, which maintains an appropriately large magnitude in the remaining weights, leading to the following learning objective:! ℳ 1 ! ℳ 2 ! ℳ 3 ! ℳ 4 ! ℳ 5 ! ℳ 6EQUATIONThe regularization term weighted by λ encourages a large magnitude after pruning.Transfer Learning of Sparsity. Transfer learning is one of the most popular paradigms in deep learning. In this section, we show the feasibility of transfer learning in sparsity, which crafts new masks by inheriting pre-computed ones. The pre-computed masks can be obtained with oneshot pruning methods like Magnitude Pruning [15] , SparseGPT [12] and Wanda [38] , or produced by another learning process. Note that given a probability [p 1 , p 2 , . . . , p | S|], the transformation to the final mask is straightforward with a simple argmax. However, if it is possible to map a pre-computed mask back to the class probabilities, then the proposed MaskLLM can begin with a good initialization for sampling. This can hugely improve learning efficiency and quality. To achieve this, we propose Mask Prior, a simple technique to initialize a distribution. Given a prior mask denoted as M 0 , we compute its similarity to all candidate masks as:EQUATIONwhich computes the inner product of two masks and re-centers the results with the mean. Note that for N:M sparsity, the range of M 0 M⊤ i will always be [0, N ], the mean value i (M i M⊤ ) = N/2 is a constant. For candidate masks with high similarity to the prior mask, we increase its probability at the initialization stage with:EQUATIONwhere σ(o) is the standard deviation of logits and α is a hyper-parameter that controls the strength of prior. When α = 0, we learn the differentiable mask without any prior from one-shot methods.Method Summary. The learning process of MaskLLM is straightforward. We begin with randomly initialized logits and update it with prior masks as Equation 10if available. Then we optimize the logits to solve the objective in Equation 8. The mask M i with the largest logits will be taken as the final mask for inference. This process is summarized in Algorithm 1.Algorithm 1 MaskLLM: Learnable Semi-Structured Sparsity for LLMs (2:4) We evaluated MaskLLM on three large language model families, ranging in size from 843M to 15B parameters. This included public models like LLaMA-2 7B and 13B [39] , Nemotron-4 15B [31] , and two in-house models, multilingual GPT-3 843M and 2B [36] . For LLaMA-2 and Nemotron-4, we collected a blended training set following the original papers [36, 31] for training. For the GPT-3 multilingual models, we used the original training set for mask learning. To learn masks, we trained the Gumbel logits for 2,000 steps without updating the LLM parameters. For evaluation, we follow SparseGPT [12] to use C4 dataset [34] for one-shot pruning and Wikitext [28] for evaluation. In addition, we also deploy LM-Eval-Harness [13] for zero-shot evaluation. More details about the models, datasets, training, and evaluation can be found in the appendix.1: procedure DIFFERENTIABLEMASK(π, S) 2: Obtain the soft index ỹ = exp((π i •κ+g i )/τ ) j exp((π j •κ+g j )/τ ) , gi = -log(-log ϵi), ϵi ∼ U (0, 1). 3: Compute the differentiable mask M = ỹ × S = |S| i=0 pi • Mi 4: return M 5: end procedure 6: S = { M1, M2, . . . , M|S| } = {[1, 1, 0, 0], [1, 0, 1, 0], . . . [0, 0, 1, 1]Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Justification: This submission discusses potential societal impacts in the appendix. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
learning 2:4 sparsity in llms: 
 Finding 1. Learnable Sparsity scales effectively to large-scale datasets and can fully leverage computational resources to learn precise masks through end-to-end training.End-to-end training yields accurate masks. In Table 1 , we report the perplexity and accuracies of our method, compared to three 2:4 sparse baselines: Magnitude Pruning [16] , SparseGPT [12] , and Wanda [38] . Previous works can produce satisfactory 2:4 masks efficiently but often suffer from inaccurate estimation of weight importance. The inaccuracy mainly arises from two factors:(1) Accuracy of importance metric: Due to the difficulty of computing the error caused by pruning, existing methods use approximated metrics to estimate weight importance, which inevitably results in errors. (2) Scalability: LLMs are usually pre-trained on large-scale datasets with rich knowledge, but the calibration sets used in existing methods contain very limited samples. With the learnable mask, the above challenges can be naturally addressed through end-to-end training on large-scale datasets, which directly optimizes the language modeling loss. As illustrated in Table 1 , MaskLLM yields superior results compared to existing baselines. For instance, with the LLaMA-2 7B model, the proposed method learns a mask with a PPL of 6.72, which is better than the PPL of 10.42 obtained by SparseGPT with weight update. More results such as comparison to other baselines (Table 13 ) and visualization of mask difference (Figure 8 ) can be found in the appendix.Scaling to large-scale datasets. To further elaborate on the above analysis, we illustrate the relationship between the number of consumed samples and the Wikitext PPL of pruned LLaMA-2 7B in Figure 4 . For one-shot methods such as SparseGPT, all consumed samples are used to compute the Hessian for importance estimation. Increasing the calibration set size from 32 to 256 samples improves the results, but expansion beyond 256 samples yields no notable advantages. More results for Llama-3 [1] or other SOTA methods can be found in Table 12 and 13 of the appendix. In contrast, our proposed learnable method effectively scales to large datasets. Results in Figure 4 show that increasing the number of samples within our framework consistently improves mask quality, with positive results still observable when scaling up to 512k samples. Additionally, our method is also data-efficient and thus applicable to low-resource scenarios with only 1280 samples. With a batch size of 256, the learnable mask is updated for only 5 steps and still produces slightly better masks than SparseGPT. If limited to only one or two steps, the training-based method fails to be comparable to one-shot methods, as this limits the random exploration for finding high-quality masks.Finding 2. Taking pre-computed masks as prior improves training efficiency and mask quality.Transfer Learning with Mask Prior. An important feature of the proposed method lies in transfer learning. We can initialize the Gumbel logits with pre-computed masks, which significantly accelerate the training. In Table 2 , we learn masks using different prior types, including Magnitude prior [15] , SparseGPT prior [12] , and Wanda prior [38] . Firstly, even without any prior, the learnable mask still achieves superior quality compared to the existing baseline methods, demonstrating its capability to independently discover high-quality masks through end-to-end training. However, learning accurate masks in only 2,000 steps can be challenging due to the massive parameter scale of LLMs. Using prior masks pre-computed by one-shot methods can provide substantial benefits. For example, with the Magnitude prior that can be easily pre-computed according to the weight magnitude, we can improve the wikitext perplexity of LLaMA-2 7B from 9.12 to 6.77. The effectiveness of transfer learning with prior masks. We report the Wikitext PPL of both prior and learned masks. The learned masks use the corresponding prior for initialization and refine the logits through end-to-end training. All results are obtained with frozen weights.Finding 3. The randomness of sampling is crucial for mask learning.Encouraging stochastic exploration on candidate masks. At the early stage of mask learning, the optimal mask is unknown. The stochastic sampling with Gumbel softmax allows for the exploration of different candidate masks, which is crucial for effective learning. As mentioned in Section 3.2, the scaling factor κ controls the randomness of sampling. To illustrate this, we visualize the learning process in Figures 5a and 5b , showing the mask difference between adjacent steps and the maximum probability of the learnable distribution, respectively. With a large factor, such as κ=1e5, the Gumbel softmax will be dominated mainly by the logits rather than the Gumbel noises, which produce similar masks with high confidence throughout the training process. In contrast, with a small scaling factor, such as κ=1, the Gumbel noises contribute more to the sampling. As illustrated in Figure 5a , the mask is continuously changing during training, leading to slow convergence. Therefore, selecting an appropriate scaling factor is crucial, which should guarantee sufficient randomness and an acceptable convergence speed. In this work, we use a κ=1e2 and linearly increase it to 5e2 for all experiments.Finding 4. Maintaining a large magnitude of the remaining weights improves downstream tasks. Maintaining a large magnitude of the remaining weights.In Equation 8, we introduce a regularizer in the form of-λ i ∥W i ⊙ Mi ∥ 2 2. This regularizer is crucial for both mask learning and transfer learning, as it directly influences the magnitude of gradients during training. For instance, if certain layers are pruned to a small magnitude, the gradients passed to their inputs will also diminish, thereby impeding mask learning and transfer to downstream tasks. In Table 3 , we demonstrate the effectiveness of weight regularization under different scenarios, such as mask training, LLM fine-tuning after pruning, and transfer learning to downstream tasks. As will be elaborated in subsequent sections, the proposed regularization helps the learning of lossless masks for downstream tasks. We provide more analysis in Section F of the Appendix. Large language models can achieve satisfactory quality across a variety of tasks. In many cases, we are more interested in one particular ability of these large models under a specific task, such as programming or translation, for which an LLM is over-parameterized. This naturally introduces a new problem: can we learn a mask for specific tasks to achieve lossless compression? To evaluate this, we learn masks for 2,000 steps separately on different domains and tasks and report the task-wise PPL in Table 4 . We considered one-shot pruning as baselines, where we collected 256 samples from the task dataset for calibration. Results show that lossless masks can be learned for many tasks with our method.We also evaluated the power of transfer learning for downstream tasks in Table 5 . To deploy sparse LLMs for a single task, we can directly pick the pre-computed general mask from Table 1 , or train an "expert" mask from scratch. However, both strategies show a quality drop compared to the dense model (PPL=7.42) since they either allocate some capacity for other domains (PPL=10.61) or only see limited data from target domains (PPL=7.51). Our work leverages the general mask as prior and transfers it to the downstream tasks, which can produce lossless models (PPL=7.39).Updating parameters for downstream tasks results in additional copies of the model for each task, incurring higher storage costs. Learning masks alone allows for encoding task-specific masks with minimal space while keeping only a single, shared copy of the original parameters. As shown in Table 6 , task-specific masks only need 0.65 bits per parameter for storage on disk using simple arithmetic coding3 with a static, uniform symbol distribution. For BS=1 inference on an A6000 GPU, 2:4 sparsity brings 1.4× acceleration and 27% reduction in the memory footprint (broader speedup results appear in Table 16 in the appendix). 
conclusion: 
 In this work, we present MaskLLM, a learnable pruning method that crafts accurate N:M sparsity in LLMs, thereby reducing computational overhead during inference. Our empirical experiments on several models show the scalability of MaskLLM to large-scale data and the effectiveness of end-toend training for mask learning. Furthermore, we demonstrate that lossless compression with N:M sparsity is attainable in downstream tasks, underscoring its practicality for real-world applications. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This submission is not about data or models. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Assets are properly credited. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In table 11 , we compare the learned masks on the C4 dataset [34] with those on the blended datasets discussed in Section A. Our blended dataset encompasses a broader range of topics and domains compared to the C4 dataset, including coding, different languages, etc. Despite this, the result in table 11 still indicates that MaskLLM is able to learn accurate masks on the C4 dataset, with a minor difference (∆PPL=0.07) compared to the result obtained on the blended dataset. In Table 13 , we compare MaskLLM to several baseline methods that were not implemented using the Megatron framework. We report the official results on Wikitext-2 PPL and LLaMA-2 13B.Even compared to methods that incorporate weight updates, our method achieves superior perplexity results. Weight Norm of Sparse LLMs. In Equation 8, we introduce an additional term to preserve sufficient gradients during training. As shown in Table 14 , a larger weight regularization facilitates large gradients during training, which is beneficial for mask exploration. We also illustrate the weight magnitude of pruned LLMs, obtained using magnitude pruning, SparseGPT (Hessian) and MaskLLM in Figures 6a and 6b . These figures show the relative L1 norm of pruned weights compared to the magnitude pruning baseline, which produces the largest weight norm after pruning. An interesting observation is that, even when initialized with a magnitude prior, learnable method may still select some smaller values during pruning, resulting in a 10% lower norm than magnitude pruning. Introducing sparse weight regularization can effectively improve the weight norm of sparse LLMs and enhance their quality in further transfer learning or fine-tuning for downstream tasks.Regularization Average Gradient Norm 0 0.219 1e-5 0.542 1e-4 0.559 Sensitivity Analysis with Learnable Method. In Figure 7 , we analyze the sensitivity of LLaMA-2 7B using both the learnable and one-shot methods. For efficiency, we update the learnable masks for 500 steps and use Wikitext PPL as the metric. We observe a similar trend in the learned masks and SparseGPT masks, suggesting that a fast one-shot pruning method can reliably indicate sensitivity. Additionally, for 2:4 sparsity, the last layer is typically more sensitive than other layers. To maintain satisfactory results, we can keep the last layer dense, achieving a good trade-off between efficiency and quality. In Table 15 , we report the pruning results when a few layers are kept dense. H Throughput of 2:4 LLaMA-2 7B.In Table 16 , we benchmark the throughput of LLaMA-2 7B with 2:4 sparsity on an A6000 GPU using TensorRT-LLM for a batch size of 1. Throughput is evaluated as the number of tokens processed per second. Over a variety of input and output lengths, 2:4 sparsity achieves an overall acceleration of 1.36× to 1.41× compared to the dense model. In Figure 8 , we visualize the differences between learned masks and one-shot masks, using SparseGPT as the prior for the learnable mask. We observe that SparseGPT and Wanda produce similar masks, with differences typically ranging from 5% to 10%, due to their similar pruning objectives. Our method, however, can produce distinct compared to these baselines, as shown in Figures 8a and 8b . Additionally, we find that weight regularization is crucial for effective mask learning. Without weight regularization, the vanished gradient can hinder mask learning, resulting in only a 2.83% difference from the prior, as shown in Figure 8c . To evaluate the generalizability of learnable semi-structured sparsity, we further extend the proposed method to Vision Transformers (ViTs) [10] . Table 17 presents the top-1 accuracy on ImageNet-1K by pruning an off-the-shelf ViT-B/16. In the one-shot pruning scenario, we randomly selected 128 samples from the training dataset as the calibration set, and no additional fine-tuning was conducted post-pruning. For MaskLLM, we utilized the SparseGPT mask as a prior and directly optimized the learnable mask on the ImageNet dataset while keeping the model weights frozen. Remarkably, with just a single epoch of optimization, the learnable mask achieved an accuracy of 76.23%, significantly outperforming SparseGPT's 71.52%. Moreover, a fully optimized mask with 20 epochs of training1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: This submission introduces a learnable method to learn semi-structured sparsity in large language models. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.https://pypi.org/project/arithmetic-compressor/ (log 2 (6)/4 bits for 6 mask candidates per group) 


Decision: The work formulates and applies end-to-end differentiable masking atop pre-trained LLMs and one-shot pruning methods, significantly narrowing the performance gap with dense models with ~1.5x speedups and high parameter-efficiency (size and compute). Masks were also shown as efficient downstream task adapters.

The four reviewers (selected for adjacency to the topic; avg. confidence 4.5) acknowledged the importance, novelty, and thoroughness of the work, ultimately recommending Accept or Strong Accept. The small set of ablative requests were adequately answered. I found the work comprehensive and well-written.

Due to the significant improvement on existing pruning schemes (at ~15B LLM scale, notable for learned methods) and the informative findings and ablations, I believe this LLM PEFT approach and its core concepts will have broad appeal. **I recommend acceptance w/ special presentation.**

Reviews: review: paper_summary: paper_summary: The authors proposed a novel LLM pruning technique, by modeling the distribution of all possible masks and formulate the selection of optimal masks in a differentiable way.
strength_weakness: strength_weakness: - Unclear what the computational cost of this technique is, and how that compares with the alternative/more straightforward technique of sparse pretraining/finetuning.
- Unclear why the authors do not report speedup for all models/datasets.
- Using perplexity as a proxy for downstream coding task performance is sub-optimal. Using benchmarks such as HumanEval is preferred.
questions: questions: - While this work is interesting and thoroughly executed already, can you comment on how this compares with more straightforward baseline of sparse finetuning/pretraining? What's the compute cost v.s. accuracy trade-off?
- Finding 1. is vague in two ways: 1). unclear what "large scale dataset" mean, 1B tokens is not large scale for LLM, please just state the size of the dataset explicitly. 2). I don't think your experiments are enough to justify "Learnable Sparsity ... fully leverage computational resources to learn precise masks through end-to-end training". Just say your technique better leverages computational resources than prior art, which is what your experiments show.
- What's the calibration set size to produce Table 1.?
- For Table.2 SparseGPT, did you also used the weight updates in Learned Mask setup?
- As for weight magnitude regularization (Finding 4), have you tried tuning the learning rate for downstream finetuning? This sounds like a learning rate issue.
limitations: limitations: all addressed.
paper_summary: paper_summary: The paper introduces MaskLLM, a method for introducing semi-structured sparsity (N:M mask patterns) in LLMs. The authors show that existing model pruning methods such as SparseGPT result in significant loss in model quality at smaller scales (800M ~ 15B parameters) when using semi-structured methods. They formulate the problem of find a good N:M pattern as a mask selection problem from a candidate set of mask $S$. They formulate this further as a sampling problem, whereby they sample masks for all $\mathcal{W}^{1\times4}$ parameters in a layer to measure model quality - since exact mask computation is an intractable problem for large models. They propose using Gumbel-softmax sampling to figure out soft / differentiable masks that can learned through a training process, to optimize the sampling problem. The authors further find that learning of the masks by their proposed method can result in vanishing gradients throughout the network, and enable a sparse weight regularization to promote higher gradient norms through the network. They further find that using methods such as SparseGPT as mask priors enables more effective sampling (ie, learning) of the final masks. They follow these up with different experiments and ablations to highlight the strengths of their proposed method.
strength_weakness: strength_weakness: 1. The best MaskLLM results presented rely on SparseGPT as a prior for mask intialization and then compare against SparseGPT for efficacy. This seems to be an unfair comparison, since you're using the method's best result and then improving on top of that. For e.g, without prior masks, the method seems to be much closer in performance to the SparseGPT method for the LLaMA-2 7B model (9.12 vs 10.42). It will be good to see how the method performs for say the 13B model without using any priors.
2. The authors show that SparseGPT has fundamental limits for improvement based on the number of samples (which is well documented and understood) - but the comparisons shown, for e.g., in Figure 4 are with different datasets? Did the authors test the collated datasets for the MaskLLM training with SparseGPT and what plateaus were observed there?
4. One inherent limitation to the method seems that there are many hyper-parameters to tune to find good masks ($\alpha$ for SparseGPT prior, $\mathcal{k}$ for logit scaling and $\tau$ for softmax temperature, $\lambda$ for the regularization coefficient). Do these hyper-parameters scale optimally for all model scales? Or is more fine-grained optimization needed as models scale up? 
    - Also for the $\lambda$ parameter, the authors show that using $1e^{-4}$ results in highest gradient norm [Table 12], but the hyper-parameter used in Table 7 is $1e^{-5}$, can the authors clarify this aspect?
questions: questions: 1. For the LLaMA-2 models, the authors mention using the official dataset from the paper. However, the paper has no mention of any datasets used for training the models. Can the authors clarify this discrepancy?
2. In Table 11, when the authors mention the score from RIA, does that include the channel permutation method enabled in the paper?
3. Can the authors verify their magnitude pruning results for downstream tasks for Llama2-7B (Table 1). For such a large perplexity difference between magnitude pruning and SparseGPT / Wanda - the downstream results seem too be high?
limitations: limitations: 1. There are some inconsistencies in the links in the paper (for e.g., Appendix D mentioned in page 8 - for weight regularization maps to Section 4.1 in the paper - please fix this).
2. Throughout the paper, it is unclear what results map to what experiments. For example, for Figure 5, which models were used for the ablation? Same for the results in Tables 3, 5, and 13. Scanning the Appendix also did not help clarify this. It makes the results presented somewhat non-trivial to parse.
paper_summary: paper_summary: This paper proposes MaskLLM, a learnable method to craft semi-structured sparsity in Large Language Models (LLMs). The approach involves modeling N:M masks with a categorical distribution, which is can be optimized through gumbel softmax. The key findings in this paper suggest that end-to-end learning of mask parameters can learn more effective and accurate sparsity patterns compared to one-shot methods. Additionally, the porposed MaskLLM also supports transfer learning for downstream tasks, where lossless masks can be learned for deployment.
strength_weakness: strength_weakness: 1. One of my concerns is the use of non-public datasets and LLMs in this study. It would be beneficial for the paper to include more results using publicly available data and open-source models, such as LLaMA-3, to enhance the reproducibility and applicability of the findings.
2. This is a question about the storage cost in Table 6: Could the author clarify the actual storage costs, such as the file size on disk? Additionally, what is the data format used during the fine-tuning process?
3. Table 2 highlights the significance of the prior in mask selection. Does this suggest that the proposed method is only trying to fine-tune the prior mask rather than learn something new? What happens if we don’t have SparseGPT prior? Is the learned mask very similar to its prior?
questions: questions: Please refer to "Weaknesses".
limitations: limitations: Limitations are adequately addressed.
paper_summary: paper_summary: This paper introduces the concept of learnable semi-structured sparsity (N:M sparsity). This extends to practical N:M sparsity to pretrained LLMs and makes the masks learnable. the paper also make the transfer of masks from other sparsity techniques to work for N:M. Finally, the general N:M masks learned can be easily fine-tuned for each task with great performance. the paper also shows extensive evaluations on various practical LLMs with the SOTA performance compared to other sparse LLM methods. 


----------------------
The review will be short and does not reflect the time put in for the review or the quality of the paper. When the ideas are simple and clear -- I tend to write shorter reviews to the point.
strength_weakness: strength_weakness: I do not find any glaring weaknesses concerns or questions about the current version of the paper. However, I might have missed something and would rely on the other reviewers if I was wrong. 

My only suggestion for the paper is to add a more comprehensive related work section to complete the paper. Maybe introducing learning sparsity be it mask learning with STE or STR would be a nice thing to add.
questions: questions: See above.
limitations: limitations: Yes

Comments: title: Response to Rebuttal
Acknowledged
comment: Thanks so much for the encouraging feedback! We will keep refining the draft with all the additional experiments.

Best regards,
Authors
I have read through the response from the authors. The authors provided sufficient experiment results and statistics in the rebuttal to support their claim. My three concerns are adequately addressed. I have also read through other reviews that proves the work is rather solid. I will raise my score.
Thank you for the very positive comments and invaluable suggestions! We will incorporate all the mentioned results into the draft as suggested.

Best regards,
Authors of #1520
Thank you very much for your encouraging feedback! We will ensure that all appropriate fixes and changes are incorporated into the revised version.
We would like to express our sincere gratitude for the insightful comments! We will improve the quality of our draft following the above suggestions.
Thank you for the detailed response to the reviews and additional experimental results. I understand that the 13B model results will take time, thank your for taking a stab at those.

After reading all reviews and responses, I will update my rating to accept (score: 7). Please do incorporate the appropriate fixes / changes in the revised version of the paper.
Acknowledged and will stick to my accept rating.
Thanks for the rebuttal. After looking at other reviews as well, I am in support of accepting this paper. 


Please ensure to update the related work section and incorporate the new experiments presented here in the final paper.