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 2024 spotlight (2024)
Keywords: Sparsity
Pruning
LLMs
Semi-structured Pruning
Primary area: natural_language_processing
TL;DR: Learnable Semi-structured Sparsity for LLMs

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.

MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models
Gongfan Fang*.g
Hongxu Yin Saurav Muralidharan" 9
Greg Heinrich
Jeff Pool
Jan Kautz
Pavlo Molchanov Xinchao Wang**
NVIDIA National University of Singapore*
gongfandu. muw.euuuuuuuueuuu xinchao@mus.c .ww.. edu.
{dannyy, sauravm.gheinniic. j
Abstract
Large Language Models (LLMs) are distinguished by their massive parameter counts, which typically result in significant redundancy. This work introduces MaskLLM, learnable pruning method that establishes Semi-structured (or "N:M"" Sparsity in LLMs, aimed at reducing computational overhead during inference. Instead of developing 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 large datasets and learns accurate masks; 2) Trunsferability 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 15B parameters, and our empirical results show substantial improvements over state-of-the-art methods. For instance, leading approaches achieve perplexity (PPL) of 10 or greater on Wikitext compared to the dense model's 5.12 PPL, but MaskLLM achieves significantly lower 6.72 PPL solely by learning the masks with frozen weights. Furthermore, MaskLLM's s learnable nature allows customized masks for lossless application of 2:4 sparsity downstream tasks or domains. btttp:://gttubb /ggggtttubb..... com /Vlaas
Code is available at wtthtttt//// //tpss:/pp
Introduction
Large Language Models (LLMs) have demonstrated
Leamable
Frozen LLM
LAchater 73% Memory
markable effectiveness across a diverse range of tasks [19]
6.
48
13
However, the generality and robustness of
Werght
LLMs are largely attributed to their vast scale, with pa-
Wis
rameter counts ranging from one billion several hun- dred billion [39
Weiggtt
45 5] This substantial model size, in
Tea 2umma
Weight
turn, makes challenging and resource-intensive to de- ploy LLMs in real-world applications. One effective and practical approach address this issue is semi-structured pruning (29 32 [1] 38), which introduces N:M sparsity into LLMs improve both memory and computational efficiency. The N:M pattern, with N non-zero values
1 1100 10
Dreshvl
No
NIA
100
Lassy (Aug PPL
MasKLLM
Yes
Yos
6.72
Lossless (Avg PPL-77 39)
$ weight
Figure 1: Learnable N:M sparsity for Large Language Models.
Thee work done at NVIDIA
1 Coorespnnding Authors: Pavlo Molchanov, Xinchao Wang
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

among M consecutive parameters, is typically hardware-friende to accelerators like GPUs and has thus garnered considerable attention [22 38 32..
Despite the simplicity of its core idea, semi-structured pruning still presents considerable challenges within the realm of L.LMs. 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 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 [2].
vast search space. In fully sparsified LLaMA2-7B model with 2:4 sparsity, for
To tackle the outlined challenges, we propose MaskLLM, a learnable method that facilitates end-to- end training of LLM sparsity on large-scale datasets. In the context of N:M sparsity, pruning an LLM involves selecting masks from discrete and finite set. However, the non-differentiabiiity 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 stochastic sampling process. We incorporate the Gumbel Softmax (2 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 diffcrentiable 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 D
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-h 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 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 the L.M 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 penerol-purpose and domain-specific pruning. Additionally, the framework facilitates the transfer learning of sparsity patterns across different tasks, enabling efficient training of sparsity.
2 Related Works
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:
Pruning Large Language Models. Network Pruning 15 30 17 18
41
Structured Pruning (26 43 24],, Unstructured Pruning 17 [5], and Semi-Structured Pruning [12
38
29
32
33
Structured pruning physically eliminates substructures like attention heads (26].
embeddings or depth [3) in the model, facilitating acceleration independent of specialized hardware
or software infrastructure
32). However, structured approaches typically necessitate huge retraining
2

---

efforts to recover network quality due to coarse removal of parameters 226 43 27 unstructured methods aim find a sparse model by zeroing out parameters LLMs, which is characterized by its flexibility and minimal detrimental effect on LLMs' accuracy 112 The acceleration of sparse models 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- faiinndyy 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 [2] In this study, we focus on N:M stmi-stru sparsity within Large Language Models and present learnable framework to obtain high-quality masks via end-to-end training.
2]. Conversely,
38 20
42
44).
Learnable Semi-Structured Sparsity. On another hand, 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 as pushing partial weights to zero with Sparse-Refined Straight-Thre
directly update the network weight, such 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 learn N:M masks for frozen LLMs, which
is much more challenging due to the huge parameter amount and problem scale.
3 Method
3.1 N:M Sparsity
We motivate and introduce learnable framework, MaskLLM, 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
mask selection problem with the candidate set of |S| (2)
N!(M-NI candidates, where |S] denotes the size of the candidate set, and (W) represents the combination number of potential N:M masks. For simplicity, this work primarily focuses on 2:4 sparsity, which can be naturally extended other patterns such 1:4 and 4:8. Given a parameter block comprising four consecutive parameters, denoted as W E 1xx4 the goal of sparsification is to identify the optimal binary mask M E B of the same size, ensuring that the pruned weight maintains its behavior on observed data æ ~ p(x). For 2:4 sparsity, the binary mask M must contain exactly two zeros, resulting in discrete candidate set 2 with [S2:4| = (2) 6 candidates:
M1
2 = (M E IB X4 E = 23 = (M M Ma. Ma. Ms.Mal Ma}
= {[1, 1,0, 0), [1,0, oJ, 1I, (0, 1, O, 1], [0, 1, 1, 00. [0, 0, 1,
(2)
For an LLM, there exists substantial number of parameter blocks, denoted as (W,). each requiring the selection of corresponding masks {M,}. To maintain satisfactory behavior after pruning, it is natural define the following objective for N:M sparsity:
(M;}
argmin Expp(a) [CLM(<;; (W, O M.))), (M. M.ESP/4)
where LLas refers to the language modeling loss for pre-training. The operator @ denotes element- wise 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 non- differentiable nature of mask selection and the huge parameter scale. In the following sections, we demonstrate that the mask selection can be transformed into sampling process.
3.2 MaskLLM: Learnable Semi-Structured Sparsity
Consider single parameter block W E R1x1 consisting only 4 parameters: directly determining the exact optimal mask for this block is not feasible, since the behavior of the pruned LLM also
3

---

Learning Sparsity
Lessless Masks Downstream Tasks
Domzin Misk foe Freach
Output
Bonjour
Rakooad
Weights
Differcntiable Mark - Mask Distribatiin 4
-_
i.r-
Transfer
Domain Meak A2 for HTML
</>
=n
Monury
Figure 2: This work introduces learnable semi-structured sparsity for LLMs. MaskLLM models mask selection as distribution learning problem, enabling the creation of accurate masks through end-to-end training on large-scale datasets. The learned and general mask can further transferred to downstream tasks or domains, achieving lossless compression.
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 categorical distribution with class probability Pl.P2... PIs), which satisfy Σ, Pj = 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 distribution where the mask with high probability is more likely to maintain good quality after pruning. Formally, we model the combination problem Equation 3 from the perspective of random sampling:
{o" (M.)) argmin (p(M:))
[LLM (z; (W, O M:))),
where p(M;) refers to the categorical distribution of i-th mask M,, If it is feasible to get the gradient w.r.t. the distribution, then the above objective be optimized with gradient descent as demonstrated in Figure 2 Nonetheless, drawing samples from a categorical distribution still non-differentiable.
Differentiable Sampling of Masks An effective method to model sampling operation is Gumbel Max [14, re-parameterizztion trick that disentangles the randomness of sampling into noise variable. This trick introduces method to draw samples from the categorical distribution p with an additional noise variable It produces the one-hot index y for sampling:
(accho((rggaxx(ogg// + 9:l), 9s = log(- log ei), (O, 1),
(5)
where is a random noise following uniform distribution, and the gi = log(- log <g) is known as the Gumbel noise. With the Gumbel Max trick, the randomness of sampling parameterized to an independent variable g,. The only issue towards differentiable sampling lies in the argmax and onchot operation. To address this, leverage Gumbel Softmax (2T] to approximate the index with Softmax, leading to soft and differentiable index ý (jss):
exp((log(r:) + g;)/T)
0 Σ, exp((lo((p)) + +g))/))
(6)
The temperature term 7 is hyper-parameter controlling the hardness of the sampled index. While 0, the soft index will be more close to one-hot vector, resulting in yı- With the soft index y as row vector and the mask set S as a matrix where each row refers to the i-th candidate mask M, it's easy to craft a differentiable mask through a simple matrix multiplication:
S
M M= M x S
Σ U.
M,
This operation produces 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.f. the probability p can be easily computed. This allows using the differentiable mask M to optimize the sampling problem defined in Equation

---

52:4
Logits (Leamuhlee) Sampled Index y
Soft Mask (For Training)
Hard (For Inference)
Parameters W
0,02
DHT
0.02
Weegher awerighd
0.95
0.05
0.93
0.00
M.
B 0.00
M.
0.02
0.01
Figure 3: Drawing random mask from the learnable distribution with Gumbel Softmax Each consecutive M parameters are associated with learnable distribution for candidate masks. All illustrated computations, including Gumbel Softmax, and the weighted averaging are differentiable.
Learning Masks for LLMs Equation provides differentible mask sampled from the underlying distribution p. The gradient flow can easily reach the probability ps, making an optimizable variable in the system. Typically, we do not directly learn the probability and instead, learn the logits with exp(TT.x) scaling factor K, which produces the probability as p. ΣS exp(xx-x)* As will be discussed in Section the scaling factor k. 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:} are associated with the corresponding distributions (M,)), 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 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:
min Ex M. (M.) ( (.M.O.M.) (M.))
Σ II M.II.
(8)
The 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, 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 [P1. P2..... final mask is straightforward with simple argmax. However, it possible map 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, simple technique to initialize distribution. Given prior mask denoted as Mo. we compute its similarity all candidate masks as:
P) S]]. the transformation to the
1 EMMMTT) M.M -MMM (N/2), S
sim(Mo, MaMM.
which computes the inner product of two masks and re-centers the results with the mean. Note that for N:M sparsity, the range of Mo will always be [0, N), the mean value (MM,MTT) N/2 is a constant. For candidate masks with high similarity to the prior mask, we increase its probability at the initialization stage with:
7 - + o sim M,) a,
(10)
where (o) is the standard deviation of logits and a is hyper-parameter that controls the strength of prior. When a = 0, we learn the differentiable mask without any prior from one-shot methods.
Method Summary. The learning process of MaskLLM straightforward. We begin with randomly initialized logits and update it with prior masks Equation 10 available. Then we optimize the logits to solve the objective in Equation g The mask M; with the largest logits will be taken as the final mask for inference. This process is summarized in Algorithm
5

---

Algorithm MaskLLM: Learnable Semi-Structured Sparsity for LLMs (2:4)
1: procedure DIFFERENTABLEMASKIS S)
2: Obtain the soft index y = " "5t5)//)) 9, = -loo((-- log ~~ U(0, 1).
3: Compute the differentiable mask M = y x - MM
4: return M
5: end procedure
6: {Mi, M2, Mss {[L, 0,000,00100 1, 00,0, 1}}
7: Parallel for all parameter blocks W:
8: Initialize logits n; (0, 0)) for the parameter block W
9: Incorporate prior oं = n. + (xT) (Mm * a with prior mask Mo
10: while Training not terminated do
.M- = DifferentiaheeMss((r,,5
12: Update logits with (I; W Mi
end while
14: Get the index t: = argmax(#)
Obtain mask M MM for pruning
4 Experiments
4.1 Implementation Details.
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,
follow SParseGPT [12] to use C4 dataset (34) for one-shot pruning and Wikitext [28] for evaluation. In addition, we also deploy LM-Eval- Ma--Evl--Haress [3] for zero-shot evaluation. More details about the models, datasets, training, and evaluation can be found in the appendix.
4.2 Learning 2:4 Sparsity in LLMs
Finding I. Learnable Sparsity scales effectively to large-scale datasets and can fully leverage computational resources learn precise masks through end-to-end training.
End-to-end training yields accurate masks. In Table 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: (I) Accuracy of importance metric: Due to the difficulty computing the error caused by pruning, existing methods approximated metrics 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 T MaskLLM yields superior results compared to existing baselines. For instance, with the LLaMA-2 7B model, the proposed method learns mask with 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, 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.

---

Method
Wikitext PPL
HellaS.
RACE
PIQA
WinoG.
ARC-E
ARC-C
OBQA
Avg.
1.LaMA-2 7B (39
5.12
57.03
44.11
78.07
69.39
75.38
42.92
33.20
57.16
Magnitude
54.71
44,60
33.01
68.93
61.56
60.23
31 40
23.60
46.19
SparseGPT
10.42
43.36
36.84
71.38
63.69
62.84
29.18
47.16
Wanda
11.29
41.05
35.02
70.78
62.67
61.99
27.56
22.80
45.98
MaskLLM
6.72
50.91
40.77
74.92
64.48
69.57
36,00
28.00
52.09
LLaMA-2 13B 39
4.57
60.15
44,59
79.27
72.45
78.93
47.18
34,60
59.60
Magnitude
8.32
48.69
38.47
70.24
59.67
61.32
29.69
22,00
47.15
SparseGPT
62
39.62
74.54
70.00
70.29
36,00
27.20
52.32
-Wanda
8.47
46.96
38.09
74.05
66.69
68.64
34.81
25,00
50.61
MaskLLM
5.85
55.09
41.24
77.69
67.80
73.15
40.44
30,00
$6.74
Nemotron-4 15B 31
5.78
62.60
47.75
81.34
77.11
77.69
50.77
33.00
61..7
Magnitude
2.78E+03
26.30
21911
54.62
50.67
29.29
18.52
15.60
30,98
SparseGPT
13.38
47.06
40.86
75.73
68.90
66.96
31.83
26.60
51.13
Wandd
25.05
41.13
34.16
71.71
61.72
58.46
29.78
23.80
45.82
MaskLLM
7.31
55,92
45,45
76.22
69.14
75.93
43.94
30,60
S6.74
GPT3 2B 36
9.35
47.74
36.94
75.73
61.09
63.22
29.78
27.80
48,90
Magnitude
6.02E+004
28.52
24.50
57.62
51,93
33.33
21.67
15.40
33,28
SParseGPT
22.14
34,93
29.28
66.60
54.62
53.07
22.10
16.20
39.54
-Wadda
29.76
39.11
MaskLLM
11.42
42.64
33.88
73.34
58.72
57.37
26.02
21.80
44.82
GPT3 843M (36
12.42
39.24
33.59
70.02
54.30
53.20
21.67
21.40
41.92
Magnitude
1.15E+04
25.84
21.82
55.66
50.75
28.41
19.20
15.20
30,98
SParseGPT
38.78
30.26
28.33
63.38
51.62
39.27
18.86
14.80
35.22
-Wanda
51.37
30.44
28.52
61.64
49.96
40.82
18.17
14.80
34.91
MaskLLM
15.39
34.65
30.91
66.76
51.86
49.07
20.48
20.00
39.10
Table Evaluation 2:4 Sparsity with frozen weights (SparseGPT does perform the weight update
step). One-shot pruning methods are calibrated with C4 and evaluated on Wikitext-2 following 12].
More results for Llama-3 other SOTA methods can be found Table 12 and 13 the appendix.
In contrast, our proposed learnable method effec- tively scales to large datasets. Results in Figure 4 show that increasing the number of samples within framework consistently improves mask quality, with positive results still observ- able when scaling up to 512k samples. Addi- tionally, our method also data-efficient and thus applicable 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 SPTrseGPT. If limited to only one two steps, the training- traseing-based method fails to be comparable to one-shot methods, as this limits the random exploration for finding high-quality masks.
12
Wiicitext22 PPL
11
10
318
SparseGPT
Ours
Dense
10
10
10'
10*
Consumed Samples (Unique)
Figure 4: Consumed samples vs. PPL on LLaMA- 2 7B. MaskL.LM requires 128 samples for the prior and outperforms SparseGPT after 1280 samples.
4.3 How to Learn Good Mask for LIMMs
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 we learn masks using different prior types, including Magnitude prior [15], SparseGPT prior [2]. 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 only 2,000 steps can be challenging due to 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. to 6.77.
7

---

0.5
x-1
1,0
Convergento
k=1u2
Maxt Differecoo (13)
0.4
0.8
x-165
0.3
Softma
0.6
0.2
0:1
0.4
xelel
x=1e5
Comvergence
Staget
Stajo
(a) The mask difference between adjacent steps
(b) The Maximum probability of mask distribution
Figure 5: (a) The LI distance of sampled masks between adjacent training steps. (b) The maxi- mum probability of mask distribution, serving randomness of mask sampling is regulated by the scaling factor s. too-small k introduces huge randomness, resulting in slow convergence shown in (b). And inappropriately large will suppress mask exploration and yield zero mask difference throughout the training process in (a).
indicator of convergence. In our method, the
GPT-3
84M
GPT-3 2B
LLaMA-2 7
Prior Type
Prior Mask
Learned Mask
Prior Mask
Learned Mask
Prior Mask
Learned Mask
Magnitude
1.1SE+04
16.07
6.02E+004
12.06
S4.71
6.77
SparseGPT
79.84
15.39
24.43
11.59
10.46
6.72
Wanda
51.37
16.39
27.08
12.18
11.29
6.80
No Prior
18.62
14,31
9.12
Table The effectiveness of transfer learning with prior masks. We report the Wikitext PPI. 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 crucial for effective learning. As mentioned Section 3.2 scaling factor x controls the randomness of sampling. To illustrate this, we visualize the learning process in Figures SarressSaaans [5] showing the mask difference between adjacent steps and the maximum probability of the learnable distribution, respectively. With large factor, such x=le5, 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 small scaling factor, such as k=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 appropriate scaling factor crucial, which should guarantee sufficient randomness and an acceptable convergence speed. In this work,
l=le2 and linearly increase it to Se2 for all experiments.
Finding 4. Maintaining a large magnitude of the remaining weights improves downstream tasks.
Maintaining large magnitude of the remaining weights In Equation we introduce regularizer in the form of -xE, I @ M1 This regularizer is crucial for both mask learning and transfer learning, as it directly influences the mag- nitude of gradients during training. For instance, if certain layers are pruned to 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 weight regularization under different sce- narios, such as mask training, L.LM fine-tuning after pruning, and transfer learning to downstream tasks. will be elaborated subsequent sections, proposed regularization helps the learning of lossless masks for downstream tasks. We provide more analysis in Section F of the Appendix.
Task (2B)
w/o Reg.
w/ Reg.
Mask-only Sub-domain Finetuning
11.59
11.42
7.61
7.39
10.21
9.96
Table 3: Weight Regularization on re- maining weights helps mask learning
effectiveness

---

Domain
Pascul
Story
Japanese
Chinese
OpenWeb
Averuge
GPT3-28 Dense
1.78
1.54
2.50
14.76
9.71
8,75
8.25
12.05
7.42
Magnitude
1.38E+03
1.72++03
L64E403
2.12E+05
6,950E+02
5.67E-022
7.22E+02
123E+05
4.27E+0M
SparseGPT
2.54
2,41
3,86
30.37
26.99
28.69
26.93
28,66
18.80
MearseGPT-Updaae
2.20
2.11
3,09
25 25.43
20.35
22.34
20.55
26.69
15.36
4.75
36.75
MaskLLM
1,76
1.54
1.94
15.58
9.61
7.96
6.92
13.84
7.39
Domain
CUDA
VHDL
Javascript
BigScience
Reddit-Plus
Book
Araiv
MedAbs
Average
LLaMAZ-7B Dense
1.74
1.86
2,01
6.28
11,05
7,02
3.49
4,95
4.80
Magnitude
2.52
72.95
32.17
SparseGPT
2,09
2.30
2.52
9.57
15.46
9.91
4.54
6.73
6.64
SparseGPT-Upple
1.91
2.32
9.63
14,52
9.78
4.21
6.14
6.32
Wanda
2.32
2.59
2,80
11.56
18.62
12.83
5,23
8.36
8.04
MaskLL.M
1.80
1,83
2.01
6,88
10.12
8.10
3.51
4.95
4,90
Table 4: Learning customized masks for downstream tasks with frozen LLM weights,
Mask Type
Avg. Task PPL
Storage Task (bits per param)
Model Size in Memory
7.42
Methods
Speed
General Mask 10.61
Scratch Mask
7.51
Finctuning Learned 2:4 masks
16
100%
10xx
Transfer Mask 7.39
0.65 (L 25x)
73%
1.4x
Table 5: Transfer learning is ef-
Table 6: Storage and inference of of llama-2 7B for downstream tasks
fective for downstream tasks.
4.4 Learning N:M Sparsity for Downstream Tasks
Finding 5. Learned masks can losslessly adapt frozen LLMs to downstream tasks, offering a 1.4x wall clock GPU speed up and 73% memory footprint.
Large language models achieve satisfactory quality across 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 translation, for which an LLM is over-paramettrizzed. This naturally introduces new problem: can we learn a mask for specific tasks to achieve lossless compression? To evaluate this, learn masks for 2,000 steps separately on different domains and tasks and report the task-wise PPL in Table 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 crepcoeptee general mask from Table or train an "expert" mask from scratch. However, both strategies show quality drop compared the dense model (PPI.=7.42) since they either allocate some capacity for other domains (PPI.=10.)) 1) or only limited data from target domains (PPL=7..51). Our work leverages the general mask prior and transfers it to the downstream tasks, which can produce lossless models (P31.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 single, shared copy of the original parameters. As shown in Table task-specific masks only need 0.65 bits per parameter for storage on disk using simple arithmetic coding" with static, uniform symbol distribution. For BS=1 inference an A6000 GPU, 2:4 sparsity brings 1.4x acceleration and 27% reduction the memory footprint (broader speedup results appear in Table 16 in the appendix).
5 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-to- end 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.
(log, (6)/4 bits for 6 mask candidates per group)

---

Acknowledgments and Disclosure of Funding
This work is in part supported by the Singapore Ministry of Education Academic Research Fund Tier (WBS: 1-0009440-01-00) We would like to thank Jorge Albericio Latorre for the fruitful discussion and feedback on the project.
References
[1] AI@Meta. Ilama 3 model card. 2024.
[2] Salch Ashkboos, Maximilian L Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows and columns. arXiv preprint arXiv:2401.15024, 2024.
[3] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv: 1308.3432, 2013.
[4] Vladimír Boža. Fast and optimal weight update for pruned large language models, arXiv preprint arXiv:2401.022 0222230,, 2024.
[5] Tom Brown, Benjamin Mann, Nick Ryder. Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, al. Language models are few-shot learners. Advances in neural information processing systems, 13:7197-1 2020.
[6] Wei-Lin Chiang, Lianmin Zheng. Ying Sheng, Anastasios Nikolas Angelopoulos, Tianle Li, Dacheng Li, Hao Zhang, Banghua Zhu, Michael Jordan, Joseph E. Gonzalez, and lon Stoica. Chatbot arena: An open platform for evaluating Iims by human preference, 2024.
[7] Together Computer. Redpajama: an open dataset for training large language models, October 2023.
[8] Rocktim Jyoti Das, Liqun Ma, and Zhiqiang Shen. Beyond size: How gradients shape pruning decisions in large language models. arXiv preprint arXiv:2311.04900 2023.
[9] Peijie Dong, Lujun Li, Zhenheng Tang, Xiang Liu, Xinglin Pan, Qiang Wang, and Xiaowen Chu. Pruner- zero: Evolving symbolic pruning metric from scratch for large language models. Forty-first International Conference on Machine Learning, 2024.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 19929, 2020.
[11] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv:1803.0306 2018.
[12] Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can accurately pruned one-shot. In International Conference on Machine Learning, pages 10323-10337. PMLR, 2023.
[13] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf. Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang. and Andy Zou. framework for few-shot language model evaluation, 2023.
[14] Emil Julius Gumbel. Statistical theory of extreme values and some practical applications: a series of lectures, volume 33. US Goverment Printing Office, 1954.
[15] Song Han. Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding, International Conference Learning Representations (ICLR), 2016.
[16] Song Han, Huizi Mao, and William Dally: Deep compression: Compressing deep neural networks with pruning, trained quantization and huffiman coding. International Conference on Learning Representations (ICLR), 2016.
[17] Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. In Advances in Neural Information Processing Systems, pages 2015.
10

---

[18] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. In Proceedings of the IEEE international conference on computer vision. pages 1389-1397, 2017.
[19] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song. and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint wXXi:2009.033000 2020.
[20] Ajay Jaiswal, Shiwei Liu, Tianlong Chen, and Zhangyang Wang. The emergence of essential sparsity in large pre-trained models: The weights that matter. arXiv preprint arXiv:2306.0388 2023.
[21] Eric Jang. Shixiang Gu. and Ben Poole. Categorical reparameterizztioo with gumbel-softmax. arXiv preprint (..:::6611.0014444 144, 2016.
[22] Sheng-Chun Chun Kao, Amir Yazdanbakhsh, Suvinay Subramanian, Shivani Agrawal, Utku Evci, and Tushar Krishna. Training recipe for n: M structured sparsity with decaying pruning mask. arXiv preprint arXiv:2209.076/1 076107 2022.
[23] Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
[24] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, al. Deja vu: Contextual sparsity for efficient ilms inference time. In International Conference on Machine Learning, pages 22137-22176. PMLR, 2023.
[25] Yucheng Lu, Shivani Agrawal, Suvinay Subramanian, Oleg Ryhakov, Christopher De Sa, and Amir Yazdanbakhsh. Step: Learning n: structured sparsity masks from scratch with precondition. In International Conference on Machine Learning, pages 22812-22824. PMLR, 2023.
[26] Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning large language models. In Advances in Neural Information Processing Systems, 2023.
[27] Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yanjie La, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint arXiv:2403.088. 00t.:24433.383,,, 2024.
[28] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[29] Asit Mishra, Jorge Albericio Latorre, Jeff' Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. Accelerating sparse deep neural networks. arXiv preprint arXiv:2104.08378, 2021.
[30] Pavlo Molchanov, Arun Mallya, Stephen Tyree. luri Frosio, and Jan Kautz. Importance estimation for neural network pruning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 12222 2019.
[31] Jupinder Parmar, Shrimai Prabhumoye, Joseph Jennings, Mostofa Patwary, Sandeep Subramanian, Dan Su, Chen Zhu, Deepak Narayanan, Aastha Ihunjhunwala, Ayush Datragupta, al. Nemotron-4 155 technical report, arXiv preprint arXiv:2402. 19, 2024.
[32] Jeff Pool, Abhishek Sawarkar, and Jay Rodge. Accelcrating inference with sparsity using the nvidia ampere architecture and nvidia tensort. NVIDIA Developer Technical Blog, htps://develppe. mvidia. 2021.
[33] Jeff Pool and Chong Yu. Channel permutations for n: M sparsity. Advances neural information processing systems, 14:13316-13327, 2021.
[34] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv e-prints. 2019.
[35] Dibyanshu Shekhar, Sree Harsha Nelaturu, Ashwath Shetty, and Ilia Sucholutsky. End-to-end learnable masks with differentiable indexing. 2023.
[36] Mohammad Shoeybi, Mostofa Patwary. Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv: I 0090000003... 2019.
[37] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 368:127063, 2024.
11

---

[38] Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprini arXiv:2306. 1695, 2023.
[39] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava. Shruti Bhosale, et al. Llama 2: Open foundation and line-tuned chat models. arXiv preprint 200,0:2200,002888. 2023.
[40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[41] Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedingy of the 2020 Conference on Empir- ical Methods in Natural Language Processing (EMNLP), pages 3151-6162, Online, November 2020. Association for Computational Linguistics.
[42] Haojun Xia, Zhen Zheng, Yuchao Li. Donglin Zhuang, Zhongzhu Zhou, Xiafei Qiu, Yong Li, Wei Lin, and Shuaiwen Leon Song. Flash-Ilm: Enabling cost-effective and highly-efficient large generative model inference with unstructured sparsity. arXiv preprint arXiv:2309. 10285, 2023.
[43] Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen, Sheared Ilama: Accelerating language model pre pre-training via structured pruning. arXiv preprint arXiv:231000669, 2023.
[44] Lu Yin, You Wu, Zhenyu Zhang, Cheng- Yu Hsich, Yaging Wang, Yiling Jia, Mykola Pechenizkiy, Yi Liang, Zhangyang Wang, and Shiwei Liu. Outlier weighed layerwise sparsity (owl): missing secret sauce for pruning lims to high sparsity, arXiv preprint arXiv:2310.05177 2023.
[45] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diah, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint 2rXiv:220.00000
2022.
[46] Yingtao Zhang, Haoli Bai, Haokun Lin, Jialin Zhao, Lu Hou, and Carlo Vittorio Cannistraci. Plug-and- play: An efficient post-training pruning method for large language models. In The Twelfth International Conference on Learning Representations, 2023.
[47] Yuxin Zhang, Mingbao Lin, Zhihang Lin, Yiting Luo, Ke Li, Fei Chao, Yongjian Wu, and Rongrong Ji. Learning best combination for efficient n: M sparsity, Advances in Neural Information Processing Systems, 35:941-955 2022.
[48] Lianmin Zheng, Wei-Lin Chiang. Ying Sheng. Siyuan Zhuang, Zhanghao WI, Yonghao Zhuang, 7i Lin, Zhuohan Li, Dacheng Li, Eric Xing, al. Judging Ilm-as-a-judee with mt-bench and chathot arena. Advances Neural Information Processing Systems, 36, 2024.
[49] Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hong- sheng Li. Learning n: m fine-grained structured sparse neural networks from scratch. arXiv preprint arXiv:2 102.00000 2021.
12

---

A Implementation Details
Here provide more details about the models, training data, training configurations and other resources used in our experiments.
LLaMA-2 For LLaMA-2, we collected a blended training set following the official paper 39 which consists of corpuses from 69 domains, covering CUDA, VHDL, Reddit etc. For training. we selected subset of 512k samples from the dataset and updated the learnable mask for 2,000 steps. with global batch size of 256. We used 64 A A100 GPUs during training with an 8-way tensor parallel configuration. The full training took 1,280 GPU hours for LLaMA-2 7B and 2,304 GPU hours for we also provide training results solely using the C4 dataset.
LLaMA-2 13B. In Table I1
Nemotron-4 For Nemotron-4, we collected a small training dataset covering three domains: CC- MAIN-2021-31, Open Web Math, and Gutenberg Fuzzy. We used subset of 512k samples and trained the model with 64 A100 GPUs using an 8-way tensor parallel configuration. The training process took 2,304 GPU hours.
GPT-3 (An Internal LLM). The GPT-3 multilingual models were pre-trained using the Megatron framework on a corpus of 1.1 trillion tokens. These models share similar network architecture with the official GPT [5], utilizing the standard transformer architecture [4) with layer normalization, SwiGLU activation function, and Rotary Positional Embeddings (ROPE) B7.. Both the 2B and 843M parameter models comprise 24 transformer layers with 16 attention heads. The hidden sizes are 2048 for the 2B model and 1024 for the 843M model. Furthermore, maximum sequence length for these models is 4096 tokens. For pre-training, multilingual dataset was collected, encompassing 110 domains such as HTML, C++, French, etc.
B Hyper-parameters
We summarize the hyper- parameters used in experiments in Table 7 The main results of hyper- parameter tuning are available in Table TO where we assessed different temperature, logit scaling factors and prior strength with GPT-3 843M.
Optimizer
Training Steps
Logits
Scaling
Temp.
Prior Streugth
Spurse Reg.
a.
AlamW((S-4, wd=0.1)
Nio, 0.01)
[e22 5e2]
[4, 0.05)
Spane(iP((o=3)
LLaMA-2 138
A5a4W(Se-4, wd-0.1)
2,000
A10.0.01
[1c. 5:2]
[4, 0.05
SpanetiTTa05
le-5
LLaMA-3 8B
AdamW(5e-4. wd=0.1)
2,000
A/0.0.01)
[1c2, Sc2]
[4, 0.05
(pmneGPT(=3)
le-5
Nemotron-4 4 14B
ASem wdw0.11
2,000
A(o,0001
[1c2, 5e2]
[4, 0,05
1c.5
GPT-3 2B
AdumWile 3, wd=0.1)
2,000
.((0,0001)
[1e2
[4, 0.055
SpaneGPT(om))
lc.5
GPT-3 843M
Ada- wd-0.1)
2,000
.(0000011)
[1e2, 562]
0.05)
Spane(iT((a-3)
1e.5
Table 7: Training details and hyper- for mask training
Temperature 7
+
0.05
+
0.05
4-+00
0.05
10
0.05
Wikitext PPI.
17.52
16.69
15.39
Table 8: Elyper-parameee tuning GPT-3
843M for Gumbel softmax temperature
Sealing Factor
1e2
5e2
1e3
Se3
1e5
Sc5
Wikitext PPL
5.97E+06
15.39
15.59
24.81
Table 9: Hyper-parameter tuning on GPT-3
843M for the logit scaling factor
Prior Strength
1
2
3
5
Wikitext PPL
18.62
23.65
15.59
15.39
15.48
Table 10: parmamer tuning on GPT-3 843M for the logit prior strength
13

---

C Mask Learning with the C4 Dataset
In table I we compare the learned masks on the C4 dataset (34] with those on the blended datasets discussed in Section [ Our blended dataset encompasses broader range of topies and domains compared to the C4 dataset, including coding, different languages, etc. Despite this, the result in table [I still indicates that MaskLLM is able learn accurate masks on the C4 dataset, with minor difference (APPL=0.07) compared to the result obtained on the blended dataset.
Method
Blended Data (See Sec.
Llama-2
5.12
5.12
SparseGPT [12
9.88
10.42
Wanda 38
11.25
10.29
Ma
6.72
6.79
Table 11: Wikitext-2 PPL of 2:4 LLaMA-2 7B pruned with different datasets
D 2:4 Results on Llama-3 8B
In table 12 we present additional pruning results for Llama-3 8B [1 adhering to the same training protocol as described Table For reproducibility, we utilize the C4 dataset for both calibration and mask learning.
Method
Weight Update
Wikitext-2 PPI.
Llama-3 8B Dense
5.76
Magnitude
2.61
SparseGPT 12
17.64
Wanda [38]
23.40
MaskLLM
8.50
Table 12: Wikitext-2 PPL of 2:4 LLaMA-3 8B, with the sequence length of 4096. We took the SparseGPT mask as the prior and learned the mask on the C4 dataset.
Comparison to More Pruning Methods for LLMs
In Table 13 we compare MaskLLM to several baseline methods that were not implemented using the Megatron framework. We report the official results Wikitext-2 PPL. and LLaMA-2 13B. Even compared to methods that incorporate weight updates, our method achieves superior perplexity results.
Method
Weight Update
Wikitext-2 PPI.
Llama-2 13B Dense
4.57
SparseGPT 12
8.32
Wanda 33
8.27
ADMM-Iter
7.78
GBLM 8
8.80
RIA 46
8.41
Pruner-Zero
7.41
MaskLLM
5.85
Table 13: Comparison to SOTA 2:4 pruning methods on LLaMA-2 13B, with all results collected from original papers or official implementations.
14

---

F
Sparse Weight Regularization
Weight Norm of Sparse LLMs. In Equation 8 we introduce an additional term to preserve sufficient gradients during training. As shown in Table 14 larger weight regularization facilitates large gradients during training, which beneficial for mask exploration. We also illustrate the weight magnitude of pruned LLMs, obtained using magnitude pruning, SparseGPT (Hessian) and MaskLLM in Figures [aa] and [bb] These figures show the relative L1 norm of pruned weights compared to the magnitude pruning baseline, which produces the largest weight norm after pruning. interesting observation is that, even when initialized with magnitude prior, learnable method may still select some smaller values during pruning, resulting in 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 fine-tuning for downstream tasks.
Regularization
Average Gradient Norm
0.219
le-5
0.542
le-4
0.559
Table 14: Average Gradient Norm over the First 500 Training Steps of GPT-3 2B, with Varying Levels of Sparse Weight Regularization. In this study, we use regularization strength of le-5 for mask learning, as offers a stable gradient while imposing minimal constraints on the search space.
1.0
122
(a) Relative norm of remaining weights (GPT-3 2B). (b) Relative norm of remaining weights (LLaMA2 7B).
Figure 6: The relative L1 norm pruned weights compared to magnitude pruning
G Layer Sensitivity
Sensitivity Analysis with Learnable Method. Figure we analyze the sensitivity 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 similar trend 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, can keep the last layer dense, achieving good trade-ofi off between efficiency and quality. In Table we report the pruning results when a few layers are kept dense.
Loarned
Strategy
Story
Open Web
Spanwtrtt Updaie
SparseGPT wio Uodste
Dense
14.76
12.05
533
Full Sparsity
15.58
13.84
Skip the last 1 layer
15.18
13.61
Skip first layer
15.41
13.88
Skip the last 4 layers
15.07
13.24
Skip the last 8 layers
14.95
12.92
Table 15:
Keeping sensitive layers dense can be way to trade-off quality and efficiency.
Figure 7: Layer Sensitivity of LLaMA-2 7B
15

---

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 batch size of I. Throughput is evaluated as the number of tokens processed per second. Over variety of input and output lengths, 2:4 sparsity achieves an overall acceleration of x.36x to 1.41 x compared to the dense model.
Throughput (Token/s/GPU)
Model
Input Len.
Output Len.
Dense
2:4
Speed Up
128
128
61.74
86.92
1.41 x
128
2048
59.18
82.11
1.39x
Llama-2 7B
2048
128
57.55
78.81
1.37x
2048
2048
55.40
75.10
1,36x
128
128
32.97
51.64
1.57
128
2048
31.81
49.00
1.54
Llama-2 13B
2048
128
31.14
47.19
1.55
2048
2048
30.09
45.06
1.50x
Table 16: Benchmarking LLaMA-2 7B and 13B on A6000 with TensorRT-LLM.
I Mask Difference
In Figure g we visualize the differences between learned masks and one-shot masks, using SparseGPT 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 masks compared these baselines, as shown in Figures Sa and 8b Additionally, we find that weight regularization is crucial for effective mask learning. Without weight regularization, the vanished gradient hinder mask learning, resulting in only 2.83% difference from the prior, as shown in Figure 8c
2%
5,29
2055
6407
2%
20.05
5,2%%
525
33%
0%
12.00
16.77%
23.0%
0%
22.0%%
2.0
(a) Mask Difference on LLaMA-2 7B with regularization
(b) Mask Difference on GPT-3 2B (c) Mask difference on LLaMA-2
with regularization
7B without regularization.
Figure 8: (a) (b) Notable mask difference exists between learned masks and one-shot ones. (c) without sparse weight regularization, the vanished gradient caused by pruning will hinder mask training, leading low difference between prior masks and learned masks.
J MaskLLM for Vision Transformers
To evaluate the generalizability of learnable semi-structured sparsity, we further extend the proposed method to Vision Transformers (ViTs) [00] Table 17 presents the top-1 accuracy on ImageNet- 1 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 M.s 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 single epoch of optimization, the learnable mask achieved an accuracy 76.23%, significantly outperforming SparseGPT's 71.52%. Morcover, fully optimized mask with 20 epochs of training
16

---

Method
Sparsity Pattern
Weight Update
Top-1 Acc. (%)
ViT-B/16
Densc
79.15
Magnitude
2:4
65.92
Wanda
2:4
63.28
SparseGPT
2:4
71.52
SparseGPT w/o Update
2:4
59.72
MasVLLM-4V (1 Epoch)
2:4
76.23
MaskLLM-4V (20 Epochs)
2:4
79.46
Table 17: MaskLLM for Vision Transformers
achieved lossless compression (AAcc +0.31%) 1) for the VIT-B/16 model. This observation also suggests the presence of a Lottery Ticket phenomenon [[]] in modern Vision Transformers, where a sparse sub-network can match the original model's performance without any weight update.
K Limitations.
In this work, we explore an end-to-end learning method for semi-structured pruning. Although our method yielded superior results, training LLMs with learnable masks inevitably consumes more resources compared to one-shot methods, which can produce masks efficiently. Improving the training efficiency of learnable masks is an important topic future works.
L Broader Impacts.
The technique proposed in this paper will not lead to negative societal impact. On the contrary, it offers significant benefits, including the reduction of energy costs and carbon emissions associated with the deployment of Large Language Models. By optimizing for Semi-structured (or 'N:M') Sparsity, our method reduces the computational resources required for inference, thereby contributing to more sustainable and environmentally friendly AI applications.
17

---

NeurIPS Paper Checklist
1. 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 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 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.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations 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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
18

---

Justification:
No theoretical results in this submission.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All hyperparameters and training settings are available in Table of the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
19

---

Answer: [No]
Justification:
This submission is not about data or models.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm
ooiiioionoly for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ ctpppp.ccppulli//guiii wwwiii/gguie///eeeeeeeei CoodeSubeiiieeeec ssi wwwinnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We discuss our datasets,
and tuning results in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: error bar is not available in this submission.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
20

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We report the number of GPUs and GPU hours in appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: sattps://not https///ccpccctttecccccceeeeeeee
Answer: [Yes]
Justification:
This submission follows the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
21

---

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 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.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This submission is not about data models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators 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?
Question:
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 dataset.
The authors should state which version of the asset is used and, if possible, include 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 (c.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, wwppppswpptseeee www.reeswthhhe....c. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
22

---

Answer:
[NA]
Justification: No new asset in this submission.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/coddemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurlPS 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.
23


Reviews:
[{"review_id": "OxYyQx0sxc", "reviewer": "Reviewer_AwSn", "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.", "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?\n- 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.\n- What's the calibration set size to produce Table 1.?\n- For Table.2 SparseGPT, did you also used the weight updates in Learned Mask setup?\n- 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": "all addressed.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Important and relevant problem setup.\n- The solution is novel.\n- Thorough evaluation across a range of model/dataset combinations.", "weaknesses": "- Unclear what the computational cost of this technique is, and how that compares with the alternative/more straightforward technique of sparse pretraining/finetuning.\n- Unclear why the authors do not report speedup for all models/datasets.\n- Using perplexity as a proxy for downstream coding task performance is sub-optimal. Using benchmarks such as HumanEval is preferred.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank reviewer AwSn for the valuable comments. \n\n> **Q1: Unclear computational cost. Comparison to sparse pretraining/finetuning.**\n\n**A1:** This submission discusses the computational cost of mask learning in Lines 433-435 of the appendix. According to the original tech report of Llama-2 [1], MaskLLM's training cost is 0.6% of pre-training (1,280 GPU hours / 184,320 GPU hours). On a single node with 8xA6000, MaskLLM (239 s/it) is slightly slower than the Straight Through Estimator (218 s/it) due to its additional sampling process. For a quantitative comparison, please refer to our response to Question 4.\n\n[1] Llama 2: Open Foundation and Fine-Tuned Chat Models\n\n> **Q2: Speedup for all models/datasets.**\n\n**A2:** We supplement more benchmark results for Llama-2 13B and Nemotron-15B here:\n|Model|Input Len.|Output Len.|Dense|2:4|Speed Up|\n|:-|:-:|:-:|:-:|:-:|:-:|\n|**Llama-2 13B**|\n||128|128|32.97|51.64|1.57×|\n||128|2048|31.81|49.00|1.54×|\n||2048|128|31.14|47.19|1.55×|\n||2048|2048|30.09|45.06|1.50×|\n|**Nemotron-4 15B**|\n||128|128|38.15|59.40|1.56×|\n||128|2048|37.51|57.93|1.54×|\n||2048|128|37.44|57.56|1.54×|\n||2048|2048|36.85|56.37|1.53×|\n\n*Table 1: Benchmarking Llama-2 13B and Nemotron-4 15B with A6000 and TensorRT-LLM*\n\nThe same 2:4 model will exhibit consistent speed-up across different datasets because the acceleration of the 2:4 pattern is determined by the inference engine (TRT-LLM in our case) and the hardware (A6000).\n\n> **Q3: PPL is sub-optimal for coding tasks. Using benchmarks such as HumanEval is preferred.**\n\n**A3:** Thank you so much for the insightful advice. We chose perplexity (PPL) in our paper as it's simple and general across different tasks or domains. We are still working on more specific evaluation metrics, such as HumanEval as suggested by the reviewer, and will provide additional results in the revised appendix.\n\n> **Q4: How this compare with more straightforward baselines of sparse finetuning/pretraining? What's the compute cost v.s. accuracy trade-off?**\n\n**A4:** Following the suggestion, we provide more comparison results for different strategies: (1) Mask-only learning with frozen weights, such as the MaskLLM and Wanda, (2) Sparse Fine-tuning of the remaining weights after pruning and (3) Learning both weight and sparsity, such as STE. \n|Method|Weight Update|LLM|Dense PPL↓|Sparse PPL↓|Δ PPL ↓|HellaSwag ↑|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Wanda + Sparse Finetuning [1]|✔️|Llama-1 7B|5.68|7.02|+1.33|n/a|\n|Wanda + LoRA [1]|✔️|Llama-1 7B|5.68|8.24|+2.56|n/a|\n|SR-STE [2]|✔️|Llama-2 7B|5.12|6.80|+1.68|51.26|\n|SparseGPT + SPP [3]|✔️|Llama-2 7B|5.12|n/a|n/a|51.33|\n|Wanda + SPP [3]|✔️|Llama-2 7B|5.12|n/a|n/a|50.61|\n|**MaskLLM + Sparse Finetuning**|✔️|**Llama-2 7B**|**5.12**|**5.83**|**+0.71**|**54.66**|\n|**MaskLLM**|-|**Llama-2 7B**|**5.12**|**6.72**|**+1.60**|**50.91**|\n\n*Table 2: Comparison to more finetuning-based baselines.*\n\nWe re-implement SR-STE and collect other results from the related works. Some missing results are marked with \"n/a\". This table indicates that MaskLLM even without any weight updates, can achieve competitive results (PPl=6.72) compared to SR-STE (PPL=6.80). This is expected as MaskLLM can explore all mask candidates while SR-STE mainly focuses on weights with large magnitudes. Furthermore, incorporating sparse fine-tuning into MaskLLM can significantly improve both perplexity (PPL) and zero-shot accuracy on HellaSwag.\n\n**The cost-accuracy trade-off:** A key message in this work is that incorporating more computing can effectively enhance accuracy. One-shot methods are efficient yet not sufficiently accurate. Techniques like LoRA improve accuracy with more training, but they may not fully explore various mask combinations. In contrast, MaskLLM thoroughly examines different masks through the proposed differentiable sampling, incurring a 9% higher training cost than STE (See Q1), but achieving the highest accuracy among these methods. Practically, one-shot methods are preferable if resources are limited. However, if sufficient resources and samples are available for end-to-end training, MaskLLM is a good choice for compressing LLMs.\n\n[1] A Simple and Effective Pruning Approach for Large Language Models.   \n[2] Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch.   \n[3] SPP: Sparsity-Preserved Parameter-Efficient Fine-Tuning for Large Language Models.  \n\n> **Q5: Finding 1. is vague. 1) just state the size of the dataset explicitly. 2) Just say your technique better leverages computational resources than prior art.**\n\n**A5:** Following the advice, we will replace \"large-scale dataset\" with \"larger calibration set\" and specify its exact size (512k samples with 2.097B tokens) in the paper. Regarding point 2, we agree with the reviewer that the phrase \"better leverages\" is more appropriate and will revise the submission accordingly.\n\n> **Q6: What's the calibration set size to produce Table 1.?**. \n\n**A6:** For the baseline methods, we used 128 samples for one-shot pruning following the official implementation. For MaskLLM, we utilized 512k samples.  \n\n> **Q7: For Table.2 SparseGPT, did you also used the weight updates in Learned Mask setup?**.   \n\n**A7:** When used as the prior, the weight update in SparseGPT is disabled.\n\n\n> **Q8: As for weight magnitude regularization (Finding 4), have you tried tuning the learning rate for downstream finetuning? This sounds like a learning rate issue.**\n\n**A8:** We indeed conducted tuning experiments and still observed this issue with different learning rates. Our hypothesis is that low magnitude might not be a good initialization for downstream training. The validation PPL at 4,000 training steps is shown bellow, where the regularization effectively alleviates this problem.\n\n|Lr|Val. PPL @ 4K Steps|\n|-|:-:|\n|1e-3|5.91|\n|5e-4|5.85|\n|5e-4 + Reg.|**5.62**|\n|1e-4|5.84|\n\n*Table 3: Validation PPL of downstream finetuning*", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fkVK6BYjmi", "reviewer": "Reviewer_2tsd", "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.", "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?\n2. In Table 11, when the authors mention the score from RIA, does that include the channel permutation method enabled in the paper?\n3. 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": "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).\n2. 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper formulates the selection of N:M masks as an optimization problem, which enables scaling the method to large datasets and models.\n2. The paper walks through the math required to understand their optimization formulation step-by-step, making it simple to understand.\n3. The method focuses on common problems introduced by pruning (such as vanishing gradients of pruned weights etc.) and proposes ways to resolve them through regularization.\n4. For each of the parts of the proposed method, the authors present ablations and results that validate their design choices.\n5. The authors also showcase how their method can scale for transfer learning to downstream tasks (both via fine-tuning or via mask-transfer, as proposed via their prior initialized method)\n6. The paper shows performance results from using their method (~1.3~1.4x faster) on A6000 GPUs.", "weaknesses": "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.\n2. 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?\n4. 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? \n    - 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for the invaluable comments and suggestions about baselines, datasets, and hyperparameters!\n\n> **Q1: 1) The comparison between MaskLLM with SparseGPT prior and SparseGPT seems to be an unfair. The method seems to be much closer in performance to the SparseGPT method without prior. 2) It will be good to see how the method performs for say the 13B model without using any priors.**\n\n**A1:** Thank you for the insightful comment. One-shot methods like SparseGPT excel at efficiently identifying a coarse mask. In contrast, MaskLLM, a learning-based method, identifies better masks but is also more resource-intensive. Without a prior, 2,000 steps might not be sufficient to learn all the masks. Moreover, using a trivially-computable Magnitude mask as the prior also yields a respectable PPL of 6.77, compared to the 6.72 achieved with a SparseGPT prior. This is likely due to the high similarities of the masks found by Magnitude and SparseGPT (demonstrated in Figure 8 of the appendix). From our perspective, mask prior is an important design in this work and it's always encouraged and feasible to start training from some prior, which can be easily identified with even magnitude pruning.\n\n**13B model without using any priors:** Regarding the requested 13B experiment, the training is still ongoing. We will keep working to provide new results; however, due to resource and time constraints during the rebuttal period, we cannot immediately present the training results for the 13B model.\n\n> **Q2: 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?**\n\n**A2:** For the one-shot method, we use the C4 dataset as was suggested in the original papers. For training with MaskLLM, we utilize a blended dataset collected from the Internet, which covers 69 domains such as programming languages as mentioned in Line 430 of the appendix. Following the advice, we provide more results to evaluate the effectiveness of SparseGPT and MaskLLM on different data sources. The following table shows results obtained by **only** using the blended dataset or the public C4 dataset. It can be observed that the results between datasets are comparable and MaskLLM is still better than the baselines. We will include these additional results in the revised version of the paper.\n|Calibration Data|Blended Data|C4|\n|-|:-:|:-:|\n|Llama-2 7B|5.12|5.12|\n|SparseGPT|9.88|10.42|\n|Wanda|11.25|10.29|\n|MaskLLM (2K steps)|**6.72**|**6.85**|\n\n*Table 1: Wikitext-2 PPL of 2:4 LLaMA-2 7B pruned with different datasets.*\n\n> **Q3: 1) There are many hyper-parameters to tune. Do these hyper-parameters scale optimally for all model scales? Or is more fine-grained optimization needed as models scale up? 2) Why 1e-5 was selected for regularization?**\n\n**A3:** **1)** We agree with the reviewer that our method can be customized by a number of hyperparameters. Fortunately, as demonstrated in Table 7, we applied **the same hyperparameters across all models** and obtained consistently positive results. This indicates that the selected hyperparameters are robust and generalizable to different models and datasets. **2)** Regarding regularization, we opted for the relatively smaller regularization of 1e-5 since it produces a slightly better validation loss (1.83 vs. 1.86) during training. An over-large regularization may limit the searching space of mask learning.\n\n> **Q4: 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?**\n\n**A4:** We do not have access to the official Llama-2 dataset as it is not publicly available. As discussed in Line 431 of the appendix, we collected a blended dataset following similar principles as stated in the Llama paper, encompassing 69 domains such as Programming (C++, C Sharp, Python) and general corpus. We will clarify this point in the revised manuscript. Additionally, we will provide the C4 results mentioned above, and make the code and our learned binary masks public to enhance the reproducibility of our experiments.\n\n> **Q5: In Table 11, when the authors mention the score from RIA, does that include the channel permutation method enabled in the paper?**\n\n**A5:** Thank you for the comment. The RIA result is collected from the official paper. Channel permutation is disabled for this result. If enabled, the PPL of RIA will be 7.77 compared to our 5.85. We will update Table 11 with this result.\n\n> **Q6: Can the authors verify their magnitude pruning results for downstream tasks for Llama2-7B (Table 1). The downstream results seem to be too high?**\n\n**A6:** We double-checked these results and found them to be accurate. The same phenomenon has also been reported in the Wanda paper, where the accuracy drop on zero-shot tasks is 8%.\n\n> **Q7: There are some inconsistencies in the links in the paper (please fix this).**\n\n**A7:** Thank you for the comment. We will fix these links according to the instructions and carefully review other links to ensure their correctness.\n\n> **Q8: 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.**\n\n**A8:** Tables 3 and 5 utilize GPT-3 2B for quick experiments, while Table 13 covers all three models: GPT-3B, 843M, and Llama-2 7B. We will clarify the models in the captions following your advice.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rFE3kedbsd", "reviewer": "Reviewer_HusL", "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.", "questions": "Please refer to \"Weaknesses\".", "limitations": "Limitations are adequately addressed.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The ideas of learnable masks and transfer learning are innovative for Sparse LLMs. The proposed method enables task-oriented compression for downstream applications without necessitating the re-training of LLM weights, making it practical for real-world applications.\n2. Results on several LLMs are positive. Table 1 demonstrates that the learnable mask method can achieve superior performance compared to state-of-the-art methods, while keeping the LLM weights frozen throughout the learning process. This indicates significant potential for further improvements in N:M sparsity within LLMs.\n3. The results in Table 4 are interesting. They show that the learnable mask does not require so many training steps and samples to outperform the oneshot baseline. And the proposed method is more scalable with large-scale datasets.\n4. The key ideas and findings in this paper are clear and well-organized. It is easy to understand the key messages in different experiments.", "weaknesses": "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.\n2. 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?\n3. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the invaluable comments and questions! \n\n> **Q1: 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.**\n\n**A1:** Thank you for the suggestions. We provide additional results for the public C4 dataset and Llama-3 7B. \n\n**1) The C4 dataset:** The following table demonstrates that our method achieves comparable results across different datasets. However, the blended dataset yielded slightly better results due to its inclusion of more diverse domains, such as programming.\n\n|Method|Blended Data|C4|\n|-|:-:|:-:|\n|Llama-2 7B|5.12|5.12|\n|SparseGPT|9.88|10.42|\n|Wanda|11.25|10.29|\n|MaskLLM|**6.72**|**6.85**|\n\n*Table 1: Wikitext-2 PPL of 2:4 LLaMA-2 7B pruned with different datasets.*\n\n\n**2) Llama-3 8B**: In addition, we also provide new results for Llama-3 8B only using the public C4 dataset, using the same training protocol and hyperparameters as Llama-2 7B for mask learning. Our method continues to achieve superior performance on the latest Llama-3 model. We will include all these new results in the appendix.\n| Method           | Weight Update | Wikitext-2 PPL |\n|----|:-:|:-:|\n| Llama-3 8B Dense | -             | 5.76           |\n| Magnitude        | -             | 2.61E+03       |\n| SparseGPT        | ✔️             | 17.64          |\n| Wanda            | -             | 23.40          |\n| MaskLLM          | -             | **8.50**       |\n\n*Table 2: Wikitext-2 PPL of 2:4 LLaMA-3 8B, with a sequence length of 4096. We took the SparseGPT mask as the prior and only used the public C4 dataset for this experiment.*\n\n> **Q2: 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?**\n\n**A2:** Regarding the Llama-3 model, the actual storage cost for the binary mask is 564MB, which is compressed using the lossless compression method -- ``np.savez_compressed``. In contrast, the original model requires 15GB for storage. For five downstream tasks, using SparseGPT would require 5x (15GB/2 (for 50% zeros) + 564MB) on disk, but MaskLLM would require only 15GB + 5x 564MB. The training is conducted with BF16 precision.\n\n> **Q3: 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?**\n\n**A3:** Thank you for the question. A mask prior can be a useful jump-start, helping to reduce the number of samples needed to achieve good quality. MaskLLM indeed inherits some masks from the prior and refines them to enhance performance. Figure 8 in the appendix illustrates the mask differences. A small difference is observed between magnitude pruning and other one-shot methods like SparseGPT and Wanda. However, after training, the mask difference between MaskLLM and one-shot methods can be larger than the differences among the one-shot methods themselves. Regarding the availability of different priors, we believe that the magnitude prior is always accessible for mask learning, showing comparable or even superior performance to other priors like SparseGPT or Wanda.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qxjC12t5ca", "reviewer": "Reviewer_qjsy", "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. \n\n\n----------------------\nThe 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.", "questions": "See above.", "limitations": "Yes", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I really enjoyed reading this paper. It was a very practical paper on many levels I have read on sparsity and applicability to LLMs. \n\nThe idea of differentiable masks is not new -- as pointed by the authors -- in CS (Savarese et al., 2019), STR (Kusupati et al., 2020) along with other methods mentioned in the paper. Same goes with N:M. However, bringing them together (again done earlier as mentioned by the paper) and combining them with pretrained LLMs makes is very useful in practice. I really like 2 other aspects apart from the base idea and practicality. THe transfer of one-shot pruning masks to N:M as scaffolds and then using learned N:M masks for downstream tasks was strong. \n\nThe results are pretty solid as well. I also like the downstream evals and not just having perplexity values. the findings of the paper highlighted are very useful in further research. \n\nI am strongly in support of the paper unless I missed something obvious. I appreciate the author on the comprehensive paper.", "weaknesses": "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. \n\nMy 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for the invaluable suggestions and the positive comments! We will polish our draft with the following new results and make our code and learned masks public for better reproducibility.\n\n> **Q1: 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.**\n\n**A1:** Thanks for the suggestion. We provide new results in the following table by re-implementing the SR-STE method [1]. Additionally, we compare MaskLLM to several fine-tuning methods [2,3] from existing works, including simple fine-tuning and PEFT. \nThis table indicates that MaskLLM, even without any weight updates, can achieve competitive results (PPL=6.72) compared to SR-STE (PPL=6.80). This result is expected, as MaskLLM can explore all mask candidates, whereas SR-STE primarily focuses on weights with large magnitudes. Additionally, incorporating sparse fine-tuning into MaskLLM can significantly improve both perplexity (PPL) and zero-shot accuracy on HellaSwag.\n\nWe will polish the related work section and the appendix to include those baselines, and will also make our code and learned masks public for better reproducibility.\n\n|Method|Weight Update|LLM|Dense PPL↓|Sparse PPL↓|Δ PPL ↓|HellaSwag ↑|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|\n|Wanda + Sparse Finetuning [1]|✔️|Llama-1 7B|5.68|7.02|+1.33|n/a|\n|Wanda + LoRA [1]|✔️|Llama-1 7B|5.68|8.24|+2.56|n/a|\n|SR-STE [2]|✔️|Llama-2 7B|5.12|6.80|+1.68|51.26|\n|SparseGPT + SPP [3]|✔️|Llama-2 7B|5.12|n/a|n/a|51.33|\n|Wanda + SPP [3]|✔️|Llama-2 7B|5.12|n/a|n/a|50.61|\n|**MaskLLM + Sparse Finetuning**|✔️|**Llama-2 7B**|**5.12**|**5.83**|**+0.71**|**54.66**|\n|**MaskLLM**|-|**Llama-2 7B**|**5.12**|**6.72**|**+1.60**|**50.91**|\n\n*Table 1: Comparison to more finetuning-based baselines.*\n\n\n[1] Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch    \n[2] A Simple and Effective Pruning Approach for Large Language Models    \n[3] SPP: Sparsity-Preserved Parameter-Efficient Fine-Tuning for Large Language Models", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all reviewers for their invaluable comments and suggestions. We will make every effort to provide additional results to support our response within this limited rebuttal period. To ensure better reproducibility, we also promise to release our code and learned masks in the future.

Meta-review: 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.**
Decision: Accept (spotlight)