MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs
Authors: Yan Sun
Qixin Zhang
Zhiyuan Yu
Xikun Zhang
Li Shen
Dacheng Tao
Venue: ICLR 2026 Poster (2026)
Keywords: Semi-structured sparsity
Policy gradient learning
Probabilistic Relaxation
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We propose a lightweight learnable method for N:M sparse masking.

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

Published as a conference paper at ICLR 2026
MASKPRO: LINEAR-SPACE PROBABILISTIC LEARNING FOR STRICT (N:M)-SPARSITY ON LLMS
Yan Sun"
Qixin Zhang'
School Computer Science
Generative Lab
Faculty of Engineering
College of Computing and Data Science Nanyang Technological University
The University of Sydney
ysun98998uni. syd.egyyynnn...yyy......mo sydney..
edu.
qixin. zhang@ntu. edu. sg
Zhiyuan Yu
Xikun Zhang
University of Science and Technology of China yuzhiyuan@mar usto. edu. cn
RMIT University
xikun. zhange edu au
Li Shen
Dacheng Tao
School of Cyber Science and Technology Shenzhen Campus Sun Yat-sen University ma mathshen1 lennnii@mmaa @mmai 1. com
Generative Al Lab
College Computing and Data Science Nanyang Technological University dacheng. taoêntu. edu.
ABSTRACT
The rapid scaling of large language models (LLMs) has made inference efficiency primary bottleneck the practical deployment. To address this, semi-structured sparsity offers promising solution by strategically retaining N elements out of every M weights, thereby enabling hardware-fries friendll acceleration and reduced memory. However, existing -NoM)-co approaches typically fall into two categories: rule-based layerwise greedy search, which suffers from considerable enrors, and gradient-driven combinatorial learning, which incurs prohibitive train- ing costs. To tackle these challenges, we propose a novel linear-space probabilistic framework named MaskPro, which aims learn prior categorical distribution for every M consecutive weights and subsequently leverages this distribution to generate the (N:M)-sparsin) throughout an N-way sampling without replacement. Furthermore, to mitigate the training instability induced by the high variance of policy gradients in the super large combinatorial space, we propose a novel update method by introducing moving average tracker of loss residuals instead of vanilla loss. Finally, we conduct comprehensive theoretical analysis and extensive experi- ments validate the superior performance of MaskPro, as well as its excellent scal- ability in memory efficiency and exceptional robustness to data samples. Our code is available at
INTRODUCTION
Recent studies have witnessed the rapid advancement of LLMs across various domains, establishing them as a highly promising solution for wide range of downstream tasks (Hendrycks al., 20200;
Brown et al. 2020 Achiam et al. 2023). However, the massive parameter size introduces significant overhead in both training and inference TTovrron al. 2023 Grattatiori et al. (2024), underscoring the pressing need for efficient approaches in real-worls world applications (Shen et al. 2023 Zhou et al.)
2024]. In response, stmi-structured sparsity has emerged as a technique with considcrable practical potential, its acceleration can be efficiently harnessed by accelerators (Mishra al.) 2021 Pool
et al. 2021 Fan et al. 2025b a). Specifically, it adopts a designated sparsity pattern, retaining only
"These authors contributed equally to this work.
1 "Corresponding author. Dr. Tao's research is partially supported by NTU RSR and Start Up Grants.

---

Published as a conference paper at ICLR 2026
M elements
N-way sampling w/o replacement
0
X1
X2
X3
X4
Logits
Probability p
Mask m
Parameters w
Loss
Refined Policy Gradient Estimator
Figure I: Implementation our proposed MaskPro for learning (2:4)-sparse masks.
N out of every M consecutive weights, scheme commonly referred to as (N:M)-sparstyy Owing to its effective support from parallel computing libraries, its inference performance is exceptionally efficient, offering viable path toward the practical and scalable local deployment of LLMs.
Although its procedural design is relatively straightforward, effectively implementing (N:M) -sparsity while preserving model performance still remains formidable challenge. One major obstacle lies in its enormous combinatorial scale, making it extremely difficult to identify the optimal mask. Existing methods can be broadly classified into two main branches. The first category encompasses rule- based approaches that bypass backpropagation by leveraging a calibration set to greedily minimize layerwise errors through the objective minm |Wx-- (m w xx| | (Frantar & Alistarh 2023). Based on this, a series of variants incorporating auxiliary information, e.g., 12-norm of input activations (Sun) al. 2023) and gradients (Das al. [023] Dong al., 2024) have been further applied, leading to certain improvements. However, such handcrafted metrics inherently suffer from considerable gaps with the end-to-end loss, ultimately capping the potential effectiveness of these methods. address this. Fang et al. (2024) propose learning-basee method MaskLLM. Specifically, it determines the optimal solution by directly optimizing the objective minm f(m @ w) in generation tasks on a large dataset. MaskLLM achieves remarkable results, but its training costs are prohibitively high, even exceeding the overhead of finetuning the LLM itself. For instance, training the (N:M)-sparsity on d-dimensional weights requires at least additional O ((M) d
memory to save the logits. As and M scale up, this memory overhead can even grow exponentially, yielding extremely poor scalability.
Our Motivation. Existing solutions either suffer from inherent biases or incur prohibitively high training costs, making them difficult to implement. This motivates us to further explore memory- efficient learning-based method for this problem. Naturally, probabilistic modeling combined with efficient policy gradient estimators (PGE) emerges as promising study. However, due to the vast combinatorial space and large model size, the variance of policy gradients can become so substantial that training nearly impossible. Moreover, the memory overhead required to store the logits remains excessively large. To enable effective training, these two challenges must be adequately addressed.
To tackle these challenges, introduce linear-space probabilistic framework termed MaskPro. Compared with the current state-of-the-ee MaskLLM (Fang et al., 2024), instead the probability distributions for all possible masks of . weights, our proposed MaskPro establishes categorical distribution for every M consecutive elements and then utilizes this distribution to generate the (N:M)- sparsity through -wway sampling without replacement. This implies that for any (N:Mstsparsity pattern, we only require (d)) memory to store the logits. Furthermore. we propose novel PGE update accelerate and stabilize the entire training process, which modifies the independent loss metric in vanilla PGE by the loss residuals with moving average tracker. We provide the rigorous theoretical analysis for our probabilistic modeling and prove the unbiasedness and variance reduction properties of the proposed PGE. To investigate its effectiveness, we conduct extensive experiments on several LLMs and report the performance across various downstream tasks. Experiments indicate that the proposed MaskPro can achieve significant performance improvements while maintaining memory usage comparable to rule- based methods, with substantially lower training overhead than Mas Moreover, the MaskPro method demonstrates remarkable robustness to data samples, which can achieve stable performance even with only I training sample.
We summarize the main contributions of this work follows:
2

---

Published as a conference paper at ICLR 2026
We propose linear-space probabilistic framework MaskPro, formulating the -spariiy as a process of N N-way samplings without replacement within a categorical distribution over
O(d) M consecutive elements, which reduces the memory for logits from O ((W) d M
We propose an enhanced policy gradient that substitutes the raw loss in standard policy gradients with per-minibatch loss residuals. To maintain stability, we further incorporate a moving-average baseline that adaptively tracks the residual dynamics during training.
We provide the comprehensive theoretical analysis to understand the memory effectiveness of MaskPro and the variance reduction properties of the proposed policy gradient update. Extensive experiments validate its significant performance. Moreover, exhibits outstanding robustness to data samples, maintaining stable results even with only I training sample.
RELATED WORK
Model Pruning. Model pruning an important compression technique that has been adopted in
& Carbin 2018 Liu et al., 2019 Xia et al. 2023b |Snn jtt al. 2023] Sreenivas et al. 2024] Luo et al. 2025). ). It also demonstrates strong practicality in
several domains (Han et al. 2015, Frankle FrankFrankke
real-world applications of LLMs. A series of structured learning and optimization methods pruning
and training have been proposed and widely applied, including the depth- and width-baser (Ko et al.
2023), , kernel-based (Xia et al., 2023a), LoRA-based (Chen al. 2023 Zhang al. [023] Zhao et al.
2024), row- and column-based [Ashkboos et al.
2024). channel based Gao et al. 2024b Dery et al.
2024), layer-based (Yin al. [0033] Men et al.
2024 Zhang al. 2024a), attention head-base (Ma
et al. 2023). MoE-based (Chen et al. [022] Xie et al.
2024). These methods leverage prune-train process to effectively reduce the number of effective parameters while maintaining efficient training,
bring a promising solution for the practical application and deployment of LLMs in the world scenarios. However, structured pruning typically considers a specific model structure as the minimal pruning unit, which can significantly impact the model's performance. The fundamental unit of model each individual weight, implying that unstructured pruning methods generally have higher potential the performance (Frantar & Alistarh| [0233 Jaiswal et al.) (023). Such methods can typically identify a fine-grained mask that closely approaches the performance dense models.
Semi-structure Pruning. Due to the inability of GPUs and parallel computing devices to perfectly support arbitrary element-wise sparse computations, the practical efficiency of sparse models remains significantly constrained.
sparsity offers promising pathway for practical appli- cations (Zhou et al., 202T /Zhang et al. 2022 Lu et al. 2023), which is also called (N:M)-sparsity.
A series of methods supporting semi- structurdd sparsity have been consistently applied, primarily
including base-based (Han et al.) 2013 Frantar & Alistarh 2023 Sun et al.) 2023 Das et al. 2023
Dong al., 2024: Zhang et al.] 2024b) and learning- based (Holmes et al.) 2021 Fang al. 2024
Huang et al.) 2025) approaches. Our work is the first to adopt policy gradients for learning semi-
structured masks on LLMs. Enormous variance of policy gradients caused by the vast combinatorial
space makes learning -parsss via PGE more challenging than those gradient- based methods.
PRELIMINARY
3.1
SEMI-STRUCTURED SPARSITY
The core idea of semi-structured sparsity aims to divide the entire weights E Rd into groups of M consecutive elements and then retain N effective weights for each group. More specifically, we can formulate the semi-structured sparsity the following combinatorial optimization problem:
m"
arg min m=(m, MSNNN
EADD [f @ w))),,
(I)
where f(-) denotes the corresponding loss function, the symbol @ stands for the element-wise multiplication, & ~ DD represents the minibatch sampled from the underlying distribution D and SN:M (m; E BIXM |mmlli NN} (IB is the Boolean set and || ||, denotes 1 norm).
Generally speaking, in order to find the optimal mask m for problem [ we are confronted with two significant challenges: i) Huge Search Space: In the context of I. L.LMs, the model parameter
3

---

Published as a conference paper at ICLR 2026
scale d can become extremely large, which will result in the search space for problem reaching
an astounding size of (N) d/M iii Non-Differentiabilit of Mask Selection: The discrete nature of problem | prevents us from utilizing the gradient-based methods such as SGD (Lan) 2020).
conditional gradient (Braun et al., 2022) and 70 Sun et al. (2025) to search for the optimal mask m
To address these aforementioned issues, we will introduce an innovative probabilistic framework termed as MaskPro for problem D in the subsequent sections. Prior to that, we first review the state-of-the-ar learning-based MaskLLM method (Fang et al. 2024).
3.2 RETHINKING THE GENERAL PROBABILISTIC MODELING AND ITS INEFFICIENCY
Recent advance provides a learning method to address Problem 1 named MaskLLM (Fang et al. 2024). Specifically, for each group of M consecutive weights, MaskLLM defines a categorical distribution with class probability Ph....... .PisM.rr] where P, 1, and each p; represents the probability of the corresponding element SN:M By random sampling, if certain mask performs better, it is reasonable to increase the probability of the sampled mask. Otherwise, the sampling probability should be decreased. Thus, Problem can be transformed as,
[
(m:)) arg min {p(m.]]
[f @ E)
where p(m;) is the categorical distribution of the i-th mask m; over NNWMM
To enable the end-to-end training, MaskLLM further introduces Gumbel-Max(umbbe 1954) as reparameterization to relax the discrete sampling into continuous form, making naturally differen- tiable. This reparameterized loss-driven mask learning method is highly effective on various LLMs, providing innovative perspective for addressing this problem.
However, the memory overhead in the MaskLLM training process is extremely large. Firstly, the backpropagation of gradients typically requires storing large number of intermediate activation values and a substantial amount of optimizer states must be maintained during updates. A more notable issue is the separate probability assigned to each possible selection of m; over SN-M, which may cause extreme memory explosion. Concretely, when learning (N:M)-sparsity for the weights
W E IRd, MaskLLM requires at least O
(.)
space to save the logits for learning probabilities, M
d
which approximately reaches O ( (2 at the worst case (N = M/2). This implies that the computational resources required by MaskI.LM can even increase exponentially as M becomes large, significantly limiting its scalability in practical scenarios, especially with extremely large model size.
4 METHODOLOGY
In this section, we present the details of our proposed MaskPro method. Specifically, Section 4.1 we introduce the novel linear-space probabilistic framework to tackle the memory drawback of the vanilla sampling process in MaskLLM (Fang et al. 2024). Then, in Section 4.2 we propose to adopt the backpropagation-liee
policy gradient for training. Moreover, we further refine the logits update via utilizing the loss residual with a smoothing tracker instead of vanilla loss metric, which enhances the effectiveness and stability of the learning process.
4.1 MASKPRO: A LINEAR-SPACE PROBABILISTIC RELAXATION
Before going into the details of our proposed MaskPro probabilistic framework, we first present a representation theory of the concerned N:M mask set N {m; BIxM M ||m,]] 1 N). In order to better illustrate our results, we need to introduce a new operation @ for the coordinate- wise probabilistic sum of two vectors. Formally, for any E RIxM and b E RIXM
define a e b = 1 (1 a) @ (1 b), where the symbol 1M denotes the M-dimennional vector whose all coordinates are is worth noting that this @ is a symmetric associative operator, namely, a@b = b a. Therefore, it also makes sense to apply the operation @ to set of vectors. Specifically, given multiple diddimenti vectors (ai...., aN. ). we can define that
N
N
a; = a
+ 22
① aN = las
(1M a:)
(3)
i=1
4

---

Published as a conference paper at ICLR 2026
With this operation then can derive sparse representation for the N:M mask set S N:M. ii...
Theorem (Representation of N:M Sparsity)
S.:M a, fei..... Vi [M] and a1 # as ... aN )
(4)
where each e, denotes the j-th basis vector of the space RIXM
From high-level viewpoint, Theorem I offers parameter-reduced representation of the mask space SN:M, Notably, representing N distinct M M-dimensional vectors (ai...., aN typically requires at most (M)) unknown parameters. contrast, the mask set SN:M often has enormous size of (N) Particularly when is comparable to M, the parameter scale N M of vectors {ai,.... aN can be significantly smaller than the space complexity (N) SN:M.
Motivated by the results of Theorem if we represent each mask m; E SNM in problem as Vjj [N] and ai.j: + a Vj + je. then we naturally can reformulate our concerned mask selection problem [ as binary optimization with variables {a Vi [M that to say,
probabilistic of {a,,I,... 1..... a,,v) where ai, .e,,..., ear
) s.t. aij. # anja Vii # j2 E [N],
min (eL
REND
+
as, Wink
(5)
where the symbol w, denotes the i-th group of the whole weight vector E Rd and E [
Notably, in Eq we only employ NM * M Nd unknown parameters, which is significantly parameters scale used by the MaskLLM method. However, this new parameter-reduced formulation equequtioon 5f problem still remains discrete combinatorial optimization problem such that we cannot directly utilize gradient information to search for the optimal mask. To overcome this hurdle, we further introduce novel probabilistic relaxation for 5 in the subsequent part of this section.
smaller than the ((W) d M
problem
Note that in Eq we restrict each group variables (a......., to be distinct basis vectors IR 1 that is, a,,j E (ei..... ew) ), Vj E [N] and a. * a Vji + je E [N]. other words, we hope to identify an effective N-size subset from the basis vectors {el, eM which closely resembles an N-way sampling-without-rees without-replacement process over (ei,..., e ). Inspired by this finding, design a novel cantinuuuu-rellxaatinn framework named MaskPro for Eq 5 i.c., Firstly, allocate a categorical distribution P. = (P......... Pt.M) for each group of variables aa........ a,,v }. Subsequently, employ every categorical distribution p: sequentially generate N different random basis vectors [ei.i...., e,,x) throughout an N-way where ei {eI...., em} and ei j eij, Vji * j2. Finally, we assign these sampled basis vectors to the variables (a.. a by setting aij == e.,, [N].
haapling-wwthuut-epplaaae
trial
Specifically, under the previously described probabilistic framework, the discrete problem equation can naturally be converted into a continuous optimization task focused on learning the optimal categorical distributions P, across the basis vectors (el, (ei...., e,r) that is.
[ A a O w J] (
((p) = -B <p.EN [
min
(6)
where (ar.)/ P, represents the N-step sampling- without--eplacemment process guided by the categorical distribution p,. Note that representing all different categorical distributions { d typically requires / * M = d unknown parameters. Thus, by introducing randomness, the parameter scale of problem T can be further reduced from the previous Nd of problem t linear d.
Next, we utilize the pr-aammtterization trick to eliminate the unit simplex constraint inherent the problem 6 namely, (p: E [0,1] ||p:]] 1}. This step is crucial it enables to avoid the
5

---

Published as a conference paper at ICLR 2026
computationally expensive projection operations. Specifically, reset p; softmax(T,)) where n; (Tt.1.... Ts.M) is the logits of softmax function. With this reformulation, we can transform
the problem 6 as an unconstrained optimization regarding the logits и = { 1 that is,
N
min (TT) == [n., ) -
+
a O w.,
To avoid repeatedly using the cumbersome notation b in the remainder of this paper, we define m; == θ= ==1 a.j for any i E [ and also use p(m; m)) to denote the probability of our MaskPro generating the mask m, under logits и.. Then, the previous problem 7 can be rewritten as:
w. {) = Be If (m w, w, p(m| (8) min ((T) = Bg m- (m, m) [f (m
where E Ixxd fxddp is the concatenation of all mask fm...... m and p(m| II p(m (T.).
4.2 POLICY GRADIENT ESTIMATOR AND REFINED (N:M)-SPARSITY LEARNING
Thanks to the probabilistic formulation of Eq. we thus can facilitate an efficient optimization via a policy gradient estimator. Specifically, we have the following equality:
V (Φ(T) (T) = RR m= im mmm, [x.) [f(m @ w, V log (p(m|
As for the proof of Eq 9 and the specific calculation of p(m) (T) in MaskPro, please refer to
Appendix C.2 and B Note that Eaan be computed purely with forward propagation. Therefore, we can update the logits variables via a mini-batch stochastic gradient descent, that is to say,
T++1 गि f (m, O w, V log (p(m, (me))
(10)
Although 10.(10 may perform well in elementary tasks, it faces one major challenge in the context of LLMs, which is caused by the inherent differences in loss values among different minibatch.
Ambiguity on Mask m. and Minibatch Ç. The policy gradient updates logits based on the loss metric, aiming to encourage the logits to select masks that result in lower loss values. However, when the loss variation caused by mask sampling is significantly smaller than the loss variation caused by changing the minibatch, the loss metric alone cannot effectively distinguish whether the current mask is beneficial detrimental. For example, denote flow the minibatch whose loss is inherently low and Eligh as the minibatch with high loss. Then we sample two masks and denote one that achieves lower loss by mggod and the other by mbad- There are typically scenarios during training.
(Mgood ① (low) (Maod ④ W, Slow) and f(Mgood w, Enigh)
W, Enigh)-
A bad case: [Mahad W, Slow) S f (Mgood ④ w, Stigh).
The first is likely hold in most cases, good mask can generally reduce the loss on most minibatches. But when the bad case occurs, Eq0 interprets that the lower- loss sample as the better one, yielding more erroncous learning on Mabd To better illustrate this phenomenon, we randomly select two minibatches during the training of LLaMA-2-7E and extract the logits the 500-th it- eration. We then sample 1000 masks and plot their loss distributions, as shown in Figure 2 that f(Inadd parities between minibatches are quite common, causing (0) frequently encounter conflicting information when learning solely based on loss value f(m O w,E). E).
random minibatch E
100
random minibatch &
80
Fplm f(mm @ w. Ej)
40
Mgdbo moad
Mgood mad
It is clearly observed
1 < [mmood C w, E2). Such dis-
2.944
2.948
2.952
2.956
Loss f(m w. 5i)
Figure 2: Loss-related misconceptions.
To address this issue, propose
the loss residual update the logits, which can distinguish the loss variations independently caused by mask changes. By rethinking the first case above, to accurately evaluate whether mask is better, should fix the impact of minibatch. Similarly,

---

Published as a conference paper at ICLR 2026
Algorithm 1 Learning (N:M)-Sparsity via MaskPro
Input: frozen weights w, initial logits To, initial mask mo. learning rate 77, smoothing coefficient a = 0.99, smoothing tracker = 0.
Output: learned logits TT
1: for = 0, 1, 2.
--11-0 do
2: sample a minibatch & for training
3: reshape иy into groups of M elements and calculate pt softmax(T.) for each group
4: perform N-way sampling without replacement by pe to generate the mask m,
5: perform inference and calculate the loss residual f(m w.E) fmm @ w.
6: update logits 7t+1 = मे 77 (f 3 w.E) - f(mo O w. 8) V log
7: update the smoothing tracker 5 = ad (1 ax) (f(m @ w, f(mo w,
8: end for
introduce f fi,, C w)
f(mo O w, instead of f(m, w. alone to evaluate whether the current sampled mask m, is better than the baseline of initial mo- Thus, the update is refined as:
1 ग. 7 (m, O -w f'mo w, () log (p(m)/==))
(11)
In experiments, the effectiveness of Eq. 11 is significantly better than that of 1Q. However, exhibits poor numerical stability. To further handle the potential numerical explosion during training, motivated by /Zhao et al. (201T), we introduce a moving average tracker to evaluate the averaged loss residual under the current logits. Specifically, we reformulate Eq. [T] as follows:
T++ = गे 77 (m, @ w. - of (mg @ w.E) {)- 5) V log (p(m)
(12)
== ad + a) (f(m, O w. f(mo ww ())
12. not only effectively distinguishes the loss variations caused by each sampled mask but also stabilizes its numerical distribution around zero through the ö term. This prevents aggressive logits updates caused by large loss variations, ensuring a more stable training process. We also provide a theoretical intuition and understanding for the ö term in Appendix C.3.3
We summarize the training procedure in Algorithm| At t-th iteration, we first reshape the logits Te into groups of M consecutive elements and then apply the softmax funetion to generate the corresponding probabilities p.. Based on Pt, we perform an N-way sampling without replacement for each group, resulting strict (N:M)-sparse mask. We then calculate the policy gradient update the current logits. By calculating the loss residual on the corresponding minibatch ç, we can obtain the independent impact of the loss value. With the assistance of smoothing tracker, we ensure that the distribution of loss residuals used for the policy gradient remains stable. Then we complete the policy gradient update of the logits. Finally, we update the smoothing tracker Regarding the final output, since the output consists of the logits Ty of all weights, in our experiments, we directly select the top-N positions with the highest logits within each group of M elements as the mask. Actually, more refined approach is to perform multiple N N-way sampling- sampling-withoute without- rellaettreelacement processes and then evaluate them on a small calibration set to select the optimal mask.
UNBIASEDNESS AND VARIANCE REDUCTION
In this section, we primarily demonstrate the unbiasedness and variance-reduee properties of our proposed PGE update. For clarity of exposition, we denote these three updates as:
9p f(m @ w, V log (p(m|T)),
= (f(m C w, - [mm @ w, V log (p(m)r)),
Gsr (f(m w f(mo C w. d) V log (p(m(T)),
where gp is the vanilla PGE, 9r is the update via loss residual and gur is the update via loss residual with smoothing tracker 5. Then the following theorem holds (proof is deferred to Appendix C33
7

---

Published as a conference paper at ICLR 2026
Table I: Zero-shot evaluations of (2:4)-sparsity. In the test, freeze weight updates and directly apply masks. The results corresponding to each model name reflects the evaluation of dense weights.
Wiki.
HellaS.
RACE
PIQA
WinoG.
ARC-E
ARC-C
OBQA
Memory
GEMMA-7B
112.39
60.54
40.19
79.71
73.09
81.65
49.91
32.80
MASKLLM
25.42
20.10
51.52
49.49
25.21
21.59
18.40
467. 14 G
MAGNITUDE
25.23
21.24
51.85
50.75
26.43
21.84
12.40
16.32
SPARSEGPT
26.07
22.39
55.11
50.36
30.64
18.43
14.80
34.94 G
WANDA
26.80
22.78
56.47
48.86
32.66
17.75
13.60
29.63 G
GBLM
26.81
22.49
54.52
51.07
32.38
17.66
14.00
39.38 G
PPRNEER-EERE
25.27
21.63
53.21
50.75
24.58
22.70
15.20
39.38
MaskPro
26.97
23.26
57.88
52.82
32.92
22.65
16.40
48.63 G
VICUNA-I.3-7B
11.86
56.32
41.91
77.37
69.46
74.28
42.41
34.60
MASKLLM
14.91
49.07
39.13
75.24
65.35
65.57
33.57
25.60
331,16
MAGNITUDE
389.92
40.19
28.61
67.03
57.62
54.59
28.75
19.40
12.82
SPARSEGPT
24.93
44.87
37.81
70.62
63.30
62.92
32.42
25.00
22.20 G
-WANDA
25,24
44.28
37.89
70.57
61.56
61.70
32.17
23.00
21.25 G
-BBMM
24.60
44.29
38.37
70.51
61.80
62.84
31,40
24.00
26.87 G
PRUNER--EER
24.02
44.77
37.42
71.22
62.75
62.33
32.76
24.00
26.87
MaskPro
21.10
46.81
38.76
71.60
64.25
64.23
33.19
24.80
35.90
LLAMA-2-77
8.71
57.15
39.62
78.07
68.90
76.35
43.34
31.40
MASKLLM
12.55
51.17
38.56
74.70
65.04
69.57
35.67
26.80
331.16 G
MAGNITUDE
307.39
45.43
31.48
70.08
60.93
61.87
30.20
21.80
12.82
SPARSEGPP
21.07
43.20
36.56
70.89
64.56
64.52
31.48
24.60
22.20
WWANDA
23.44
41.32
35.89
70.46
62.12
62.79
30,20
24.20
21.25 G
GBLM
21.64
41.79
34.61
70.57
62.75
63.17
29.86
23.20
26.87 G
PRUNER-ZERO
22.09
41.17
34.64
70.18
62.35
61.32
27.05
22.80
26.87 G
MaskPro
17.17
46.18
37.13
73.07
65.82
66.12
32.85
26.20
35.90
DEEPŠEEK-7H
9.70
56.94
39.62
79.27
70.40
75.25
43,60
32.60
MMASKLM
12.90
51.73
39.14
75.95
65.80
68.10
35.32
25.80
339.56 G
MAGNITUDE
285.06
40.97
28.52
69.75
60.06
54.92
27.56
20.80
13.13 G
SSPARSEGPT
19.12
45.58
37.80
73.94
65.43
66.37
32.94
24.80
22.50
-WANDA
19.68
45.38
35.12
73.56
63.14
65.49
32,00
22.80
21.55
-GBLM
19.55
45.34
36.17
73.99
62.98
65.82
32.85
23.60
27.98 G
-RUNER-ZERO
20.71
44,93
35.22
73.23
62.12
64.94
30.89
23.20
27.98 G
MaskPro
17.97
47.78
37.75
74.72
65.59
66.74
33.49
28.60
36.82
Theorem The proposed PGEs are all unbiased estimators of the policy gradient, i.e.,
E [gp1
[g.1 E [g.r] V
(13)
Furthermore, when the sampled mask satisfies f(m; C w) f. @ w, we have: Var [g.r] S Var [g.] Var [gp].
(14)
In Theorem 2 Eq showes that our proposed updates 9r and Jsr are both unbiased estimators the gradient V x) effectively supporting the training process. Furthermore, when f(m, O w.g) > f O w, 3) from Eq 14 of Theorem 2, we know that before the loss of the sampling mask m, decreases less than half of the initial one, using the update via loss residual with smoothing tracker can achieve more efficient training. Once the optimization process has sufficiently progressed such that the loss is less than half of the initial loss, new of continue efficient training. In practical experiments, this condition is almost easily satisfied, as the loss rarely drops below half of the initial value when training with initial mask with simple priors.
be selected replace my to
8

---

Published as a conference paper at ICLR 2026
0.025
30
Sparseupt
Prumer-Zero
f(m; O w. 5 ) f(mo w.
Perplexity (PPL)
28
Wanda
MaskPro
-0.025
GBLM
26-
-0.050
0.075
24
-0.100
22
-0.125
20-
-0.150
Loss Residual
-0.175
Vanilla Policy Gradient
Loss Residual with Smoothing Tracker
18
-0.200
16
2000 3000 4000
5000
6000
256
512
Ieertion
Dataset Size
(a) Training Effectiveness of Three PGE Updates. (b) Training Performance of Different Dataset Size.
Figure 3: (a) We show the different loss curves trained with the three PGEs. (b) We report the PPL on Wikitext of different methods trained with I, 16, 32, 64, 128, 256, and 512 data samples.
6 EXPERIMENTS
In this section, first introduce the baselines along with details of the dataset and models. Then present the main experiments. We also conduct sensitivity studies of a and C on Appendix A.9 and to provide proper guidance for the reproducibility and extensibility.
Baselines, We select the backpropagation-fee free methods including Magnitude (Han et al.
2015)
SparseGPT (Frantar & Alistarh (2023). Wanda (Sun et al.) 2023), GBLM-Pruner (Das et al. 2023)
and Pruner-Zero (Dong et al. (024)) as baselines. We also report the results of the backpropagation- based MaskLLM (Fang
methods perform sparsification by minimizing the layer-wise wise emrrr of the output activations caused by sparse weights, while MaskLLM updates the mask by optimizing masks through the loss function of the text generation task.
2024 The backpropagation-fie
Models & Dataset. We evaluate the performance on 4 LLMs, including Vicuna-7B (Chiang et al. 2023), LLaMA-2-7B (Touvron et al. 2023). Deepseek- (DeepSeek-A) Al 2024). Gemma- (Team et al. 2024). To ensure a fair comparison, we use the C4 dataset (Raffel et al., (2220) as unified calibration or training dataset for each method and adopt the LM-evaluation soailaattoo-hareesss harness framework (Gao
et al. 2024a) for zero-shot evaluations. Due to the page limitation, more details of the hyperparameters and experimental setups for reproducibility can be found Appendix A.
Performance. In Table we report the zero-shot evaluation on several downstream tasks for the (2:4)-sparsity, We conduct extensive experiments on several 7B models to validate the effectiveness of our proposed method. MaskPro generally outperforms existing non-backpropaggation methods, achieving an average performance improvement of over 2% over the top-2 accuracy. On certain models and datasets, it achieves performance nearly comparable PPL test, the MaskPro method also shows a consistent improvement, about 3 on LLaMA-2-7B and over 3 the others. 'he weights of the Gemma-7B model are not sufficiently sparse, resulting in suboptimal performance of its corresponding sparse model and unstable PPL results. We show more evaluations in the Appendix A.33 More experiments of (4:8)/ / (8:1rstS are stated in Appendix A.5 and A.7 We also evaluate MaskPro on 13B and 30B models in Appendix A.8
MaskLLM. On the Wikitext
Optimizers. In Figure (a), we evaluate the training performance of vanilla PGE, loss residual and loss residual with the smoothing tracker. The metric on the y-axis represents how much the loss value of the current minibatch is reduced by the mask sampled from the current logits compared to the initial mask. It can be observed that the vanilla policy gradient update is almost ineffective, with the loss oscillating around zero without effectively learning any useful information. After applying the loss residual update, significant improvement is observed the logits receive effective guidance to sample better masks. However, its effect is not sufficiently stable after achieving a certain level of improvement, large oscillations occur, preventing further learning progress. The update of loss residual with the smoothing tracker can efficiently and stably train this task, leading to better results.
Size of Training Set. Our proposed MaskPro requires significantly less data samples compared to other learning- based methods. As shown in Figure (b), we evaluate the PPL of the Wikitext dataset on LLMMA2-BB after training 10k iterations with training set sizes of 1, 16, 32, 64, 128, 256, 512.

---

Published as a conference paper at ICLR 2026
According to the experimental results reported by Fang al. (2024). MaskLLM requires at least 1280 training samples to achieve the results of SparseGPT. and 520k samples for convergence. In contrast, proposed MaskPro can be trained with minimal number of training samples while maintaining nearly stable performance even with 1 data sample. We also provide results in Appendix [A.12 comparing runs initialized from different masks with 1 sample versus 128 samples. Our experiments show that training with single sample remains stable, with only slight loss in performance.
Training Efficiency. We evaluate efficiency primarily by comparing memory usage, training time, and the size of the training dataset. Traditional rule-based methods learn masks by evaluating specific metrics on a small validation set. For example, in the (2:4sssparsity on LLaMA-2-77, the Pruner-Zers requires 26.87 GB of memory and 128 C4-en data samples. And for the learning- based MaskLLM, it requires 330 GB of memory across 8x A100 GPUs and 520k training samples, taking over 1200 GPU hours. A significant advantage of our proposed MaskPro method is its low computational and memory overhead during training. More details of training time profile reports on different patterns and corresponding acceleration techniques of sampling are shown in Appendix A.13 and A.14
SUMMARY
In this paper, we propose novel memory-efficiet framework named MaskPro, which leverages policy gradient updates to learn semi-structured sparsity. By reformulating the (N:M)-sparsty as linear-space probability relaxation, our approach reduces the memory for logits storage from vanilla O (( d to O(d). Furthermore, we propose novel PGE that replaces the vanilla loss metric M
with loss residuals, refined by a moving average tracker, effectively accelerating training and reducing variance. Lastly, comprehensive theoretical analysis and extensive experiments demonstrates the effectiveness of our MaskPro in achieving substantial performance gains with minimal training costs.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.086 087744 2023.
Saleh Ashkboos, Maximilian Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows and columns. arXiv preprint arXiv:2401.15022 2024.
Gábor Braun, Alejandro Carderera, Cyrille W Combettes, Hamed Hassani, Amin Karbasi, Aryan Mokhtari, and Sebastian Pokutta. Conditional gradient methods. arXiv preprint arXiv:2211. 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Tianyi Chen, Tianyu Ding, Badal Yadav, Ilya Zharkov, and Luming Liang. Lorashear: Efficient large language model structured pruning and knowledge recovery. arXiv preprint arXiv:2310.18356. 2023.
Tianyu Chen, Shaohan Huang, Yuan Xie, Binxing Jiao, Daxin Jiang, Haoyi Zhou, Jianxin Li, and Furu Wei. Task-specific expert pruning for sparse mixture-of-experts. arXiv preprint 0rXiv:2206.00277. 2022.
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing. Vicuma: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023, httpps
mayms http://s.oo//booo2/33/3//3/ 2023- -
Rocktim Jyoti Das, Mingjie Sun, Liqun Ma, and Zhiqiang Shen. Beyond size: How gradients shape pruning decisions in large language models. arXiv preprint 2Xiv:2311.00400 2023.
10

---

Published as a conference paper at ICLR 2026
DeepSeek-AL. Deepseek Ilm: Scaling open- language models with longtermism. arXiv preprint LIM
arXiv:2401.02954, 2024. URLhtttps://giiiae wtt.:/tthh..occeeeeeeeeeeeeeeeee wtwp//peeeeaaaa-eeeaeeee
Lucio Dery, Steven Kolawole, Jean-François Kagy, Virginia Smith, Graham Neubig, and Ameet Talwalkar. Everybody prune now: Structured pruning of Ilms with only forward passes. arXiv preprint Sttiv:2402.0554000 2024.
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. In Proceedings of the 41st International Conference on Machine Learning. PMLR, 2024. URL [arXiv: 2406.02924).
wwg/abs///aeoooooootogogo... https::/ org/ abs/ 2406. 02924
Ziqing Fan, Shengchao Hu, Yuhang Zhou, Li Shen, Ya Zhang. Yanfeng Wang, and Dacheng Tao. Task-aware harmony multi-task decision transformer for offline reinforcement learning, arXiv preprint rrXiv:241..01144 146, 2024.
Ziqing Fan, Siyuan Du, Shengchao Hu, Pingjie Wang, Li Shen, Ya Zhang, Dacheng Tao, and Yanfeng Wang. Combatting dimensional collapse in Ilm pre-training data via submodular file selection. In The Thirteenth International Conference on Learning Representations, 2025a.
Ziqing Fan, Cheng Liang, Chaoyi Wu, Ya Zhang. Yanfeng Wang, and Weidi Xie. Chestx-reasonee Advancing radiology foundation models with reasoning through step-by-step verification. arXiv preprint vrXiv:2504.2093: 2025b.
Ziqing Fan, Yuqiao Xian, Yan Sun, and Li Shen. Joint selection for large-scale pre-training data via policy gradient-based mask learning. arXiv preprimt arXiv:2512.24265. 2025c.
Gongfan Fang, Hongxu Yin, Saurav Muralidharan, Greg Heinrich, Jeff Pool, Jan Kautz, Pavlo Molchanov, and Xinchao Wang. Maskllm: Learnable semi-structured sparsity for large language models. Advances in Neural Information Processing Systems, 37:7736-7758, 2024.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: /803.0333 2018.
Elias Frantar and Dan Alistarh. SparseGIT: Massive language models can be accurately pruned in one-shot. arXiv preprint arXiv:2301.007: 00774, 2023.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'l 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. A framework for few-shot language model evaluation, 07 2024a. URL htt.hhttps/eeeee /tt::// arg./reccodd//26680e
Yuan Gao, Zujing Liu, Weizhong Zhang, Bo Du, and Gui-Song Xia. Bypass back-propagatioon: Optimization-baseed structural pruning for large language models via policy gradient. arXiv preprint arXiv:2406. 10576, 2024b.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et The llama herd of models. arXiv preprint arXiv:2407.2178. 2024.
Emil Julius Gumbel, Statistical theory of extreme values and some practical applications: a series of lectures, volume 33. US Govermment Printing Office, 1954.
Song Han, Huizi Mao, and William Dally, Deep compression: Compressing deep neural networks with pruning, trained quantization and huffiman coding, arXiv preprint arXiv:1510.00149 2015.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.033 03300, 2020.
Connor Holmes, Minjia Zhang, Yuxiong He, and Bo Wu. Numtransformer: Stmi-stre sparsifi- cation for natural language understanding via admm. Advances in neural information processing systems, 18411818-888 1830, 2021.
11

---

Published as a conference paper at ICLR 2026
Weiyu Huang, Yuezhou Hu, Guohao Jian, Jun Zhu, and Jianfei Chen. Pruning large language models with semi-structural adaptive sparse training. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 24167-24172 2025.
Ajay Jaiswal, Shiwei Liu, Tianlong Chen, Zhangyang Wang, al. The emergence of essential sparsity in large pre-trained models: The weights that matter. Advances in Neural Information Processing Systems, 36:38887-38901, 2023.
Jongwoo Ko, Seungjoon Park. Yujin Kim, Sumyeong Ahn, Du-Seong Chang, Euijai Ahn, and Se- Yun. Nash: A simple unified framework of structured pruning for accelerating encoder- decoder language models. arXiv preprint arXiv:2310. 10054, 2023.
Guanghui Lan. First-order and stochastic optimization methods for machine learning, volume 1. Springer, 2020.
Zechun Liu, Haoyuan Mu, Xiangyu Zhang, Zichao Guo, Xin Yang, Kwang-Ting Cheng, and Jian Sun. Metapruning: Meta learning for automatic neural network channel pruning. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 3296-3305, 2019.
Yucheng Lu, Shivani Agrawal, Suvinay Subramanian, Oleg Rybakov, Christopher De Sa, and Amir Yazdanbakhsh. Step: learning n: M structured sparsity masks from scratch with precondition. In International Conference on Machine Learning, pp. 22812-22824. PMLR, 2023.
Haotian Luo, Li Shen, Haiying He, Yibo Wang, Shiwei Liu, Wei Li, Naiqiang Tan, Xiaochun Cao, and Dacheng Tao. Ol-pruner: cengthhharonniiigg fine-tuning for ol-like reasoning pruning. arXiv preprint arXiv:2501. 12500 2025.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models. Advances in neural information processing systems, 36:21702 36:21702-217200 2023.
Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are redundant than you expect. arXiv preprint arXiv:2403.03853, 2024.
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.
Jeff Pool, Abhishek Sawarkar, and Jay Rodge. Accelerating inference with sparsity using the nvidia ampere architecture and nvidia tensorrt. NVIDIA Developer Technical Blog, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
Li Shen, Yan Sun, Zhiyuan Yu, Liang Ding, Xinmei Tian, and Dacheng Tao. On efficient training of large-scale deep learning models: A literature review. arXiv preprint arXiv:2304.03589 2023.
Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi, Marcin Chochowski, Ameya Sunil Mahahaleshwarkar Gerald Shen, Jiaqi Zeng, Zijia Chen, Yoshi Suhara, Shizhe Diao, et al. Llm pruning and distillation in practice: The minitron approach. arXiv preprint arXiv:2408. 1796, 2024.
Mingjie Sun, Zhuang Liu, Anna Bair, and J Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306.1166) 2023.
Yan Sun, Tiansheng Huang, Liang Ding, Li Shen, and Dacheng Tao. Tezo: Empowering the low- rankness on the temporal dimension in the zeroth-order optimization for fine-tuning 1lms, arXiv preprint arXiv:2501. 19057, 2025.
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, et al. Gemma: Open models based gemini research and technology. arXiv preprint arXiv:2403.08295, 2024.
12

---

Published as a conference paper at ICLR 2026
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229-256, 1992.
Haojun Xia, Zhen Zheng, Yuchao Li, Donglin Zhuang, Zhongzhu Zhou, Xiafei Qiu, Yong Li, Wei Lin, and Shuaiwen Leon Song. Flash-llm: Enabling cost-effective and highly- large generative model inference with unstructured sparsity. arXiv preprint arXiv:2309. 10285, 2023a.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared llama: Accelerating language model pre- are-training via structured pruning. arXiv preprint arXiv:2310.06694 2023b.
Yanyue Xie, Zhi Zhang, Ding Zhou, Cong Xie, Ziang Song, Xin Liu, Yanzhi Wang, Xue Lin, and An Xu. Moe-pruner: Pruning mixture-of-ext experesss large language model using the hints from its router. arXiv preprint 2400: 2024.
Lu Yin, You Wu, Zhenyu Zhang, Cheng-Yu Yu Hsieh, Yaqing Wang, Yiling Jia, Gen Li, Ajay Jaiswal, Mykola Pechenizkiy, Yi Liang, et al. Outlier weighed layerwise sparsity (owl): A missing secret sauce for pruning Ilms to high sparsity. arXiv preprint arXiv:2310.05115. 2023.
Mingyang Zhang, Hao Chen, Chunhua Shen, Zhen Yang, Linlin Ou, Xinyi Yu, and Bohan Zhuang.
Loraprune: Structured pruning meets low-rank parameter-efficien fine-tuning. arXiv preprint arXiv:2305. 18403, 2023.
Nan Zhang, Yanchi Liu, Xujiang Zhao, Wei Cheng, Runxue Bao, Rui Zhang, Prasenjit Mitra, and Haifeng Chen. Pruning as a domain-specific Ilm extractor. arXiv preprint arXiv:2405.06275 2024a.
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, 2024b.
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-999 993,99333222 2022.
Bowen Zhao, Hannaneh Hajishirzi, and Qingqing Cao. Apt: Adaptive pruning and tuning pretrained language models for efficient training and inference. arXiv preprint arXiv:2401. 12200, 2024.
Tingting Zhao, Hirotaka Hachiya, Gang Niu, and Masashi Sugiyama. Analysis and improvement of
policy gradient estimation. Advances in Neural Information Processing Systems, 24, 2011.
Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. Learning n: m fine-grained structured sparse neural networks from scratch. arXiv preprint arXiv:2102.0400 2021.
Zixuan Zhou, Xuefei Ning, Ke Hong, Tianyu Fu, Jiaming Xu, Shiyao Li, Yuming Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, et al. A survey efficient inference for large language models. arXiv preprint arXiv:2404. 20.::2444.12294, 2024.
13

---

Published as a conference paper at ICLR 2026
The Use of Large Language Models. In this work, only evaluate the performance on LLMs in our experiments and employ LLMs to refine the writing and presentation our manuscript. Other aspects of the work are unrelated to LLMs.
This paper presents a memory-efficient training framework for learning semi-structured sparse masks based on policy gradient, achieving comprehensive improve- ments in performance and efficiency through substantial upgrades in both the probabilistic modeling and optimizers. A limitation of this paper is that when training large-scale models, the primary time consumption lies in simulating the mask sampling process. Utilizing more efficient sampling simulations can further enhance training efficiency. The core contributions of this paper mainly include linear-space probabilistic modeling and optimizer enhancements. These two aspects can be widely applied various model pruning tasks, not just the specific task addressed in this work.
Limitation and Broader Impact.
A EXPERIMENTS
A.1 EXPERIMENTAL DETAILS AND REPRODUCIBILITY
In this paper, we reproduce the baselines using their official open-source codes provided in each paper. For fairness, we use the C4-en dataset the calibration/rraining dataset. For the MaskLLM, we follow Fang et (2024) to adopt 520k C4-en samples for training 2k iterations with batchsize 256. For other methods, we follow their setups adopt 128 C4-en samples calibration dataset.
Hyperparameters. For the MaskPro, we evaluate wide range of dataset sizes, ranging from 1 to 320k. We select the learning rate from [25, 50, 100, 200] for each model and 50/ proves to be a relatively effective choice. In the training, we use batchsize as 32 and training for ~ 100k iterations. Using batchsize larger than 32 is also encouraged, as larger batches generally lead to stable training. In all experiments, we adopt the smoothing coefficient = 0.99 to stably follow the loss residual. We summarize the selection of certain hyperparameters in Table
Table Hyperparameters selections.
Model
Learning rate
Logits Magnitude
Smoothing coefficient a
Initial Mask
Gemma-7B
50/ / 100
10.0
0.99
Top-N / Sparsegpt
Vicuna-V1.3-4
50
10.0
0.99
Top-N / Sparsegpt
ILaMA-2-7B
50
10.0
0.99
Top-N Sparsegpt
DeepSeek-7B
50 100
10.0
0.99
Top- Sparsegpt
Initialization. The initialization of logits in MaskPro is crucial. Standard random initialization or zero initialization are ineffective. This is because the logits determine the sampling scale. For instance, zero initialization implies that each position is sampled with equal probability, leading to very large number of negative samples during the initial training stage. Consequently, it becomes exceedingly difficult to identify effective positive samples for learning. In our experiments, we initialize the logits based on logits magnitude. A larger C' indicates that the mask changes less compared to the initial mask mp. effectively maintaining balance between positive and negative samples in the early training stages. The design of mo is flexible. In practice, training can also start with a randomly generated mask; however, this approach typically requires longer training period. We recommend directly using the results from the Sparsegpt method or selecting the Top-N positions over M elements per group.
mo C, where me is a pre-defined mask and C' is the initial
Training Environment. We train our proposed MaskPro on single H100 / A100 GPU device. Other details are stated in T
Table 3: Training Environment.
GPU
CPU
CUDA
Driver
Pytorch
1x 11100/ / A100
128x AMD EPYC 9354 32-C
12.4
535.230.02
2.5.1
14

---

Published as a conference paper at ICLR 2026
Evaluations. For fair comparisons, all evaluations conducted the public benchmark framework A.Mevaauation-harneee framework (Gao et al. (2024))
harness.gu git). Please refer to the relevant reproduction guidelines.
A.2 THE IMPORTANCE OF C IN LOGITS INITIALIZATION
We have previously discussed the selection of C' in the experiments. Here, will visualize some practical scenarios encountered during the experiments and illustrate why C' must be sufficiently large to effectively drive the training process. We analyze the distribution of loss values of training LLAMA-2-7B within 100 steps with minibatch of 32 samples under different C' initialization settings, as shown in Figure
3.0
14.0
2.99
f(mo O ww)
12.0
2.98
f(m w
10.0
2.97
2.96
8.0
2.95
6.0
2.94
4.0
4
9 10
Vaule of
Figure 4: The distribution of loss within 100 steps under different C. used for logits initialization.
We first explain which variables are affected by C. Since we use the softmax function generate the probabilities for the corresponding positions, the logits values determine whether the initial probability of being sampled a specific position is sufficiently large. In other words, when sampling mask, it ensures how many positions with high probabilities remain unchanged. This point particularly important because the sampling space extremely large. Without constraining the sampling space, there is high probability of sampling poor masks. Extremely poor masks incapable of capturing useful information effectively. Therefore, randomly initializing the C' value or directly setting it space, i.e., whether the distribution of positive and negative samples the sampling space is balanced.
is completely ineffective,
cannot ensure the stability of the sampling
Next, explain the meaning of Figure 4 We show the distribution of loss values 100 training steps using minibatch under different C' initialization settings on the LLaMA-2-7B model. In the subplot, the red line corresponds the loss of the initialized mask mo- When C' small, it is evident that the training fails the loss surges from the initial 2.95 to over 10. A large number of negative samples flood into the training process, leading chaotic learning. As increases gradually improves. However, it is still insufficient. As shown in the subplot, even when C = 6, more than 90% of the sampled masks still exhibit extremely poor performance. Until increases to 9 and 10, it can be observed that the distribution of positive and negative sampled masks during training gradually maintains a 1:1 ratio. By this, the training can proceed effectively.
the stability
Here, we provide additional example to explain and guide the selection of C for different network parameters. As mentioned earlier, one probabilistic interpretation C is to determine, on average, how many positions
sampled differently from the initialized mask. We can succinetly express this probability in a mathematical form. Suppose the initialized mask my is [0, 1, 1, 0]. then its initial logits
[o, C, C, 0] and the corresponding softmax probability is +1)+ 2( 217+11 Ze
Thus we have:
2C
p m = C,C,C, (eC +1)
In fact, the size of the sampling space where positive and negative samples are evenly distributed is difficult estimate for different model parameter sizes. However, we can reasonably speculate that
15

---

Published as a conference paper at ICLR 2026
the total number of parameters is generally proportional to the above probability value. For larger models, using a larger C' can further maintain the effectiveness the training space.
A.3 MORE EXPERIMENTS ON DIFFERENT TASKS
In addition to the primary comparisons presented in the main text, we extend our evaluation to encompass over dozen additional tasks to provide more comprehensive demonstration of the effectiveness our proposed method. These extended tests are carefully selected to cover diverse data distributions and task complexities, allowing to assess the robustness and generalizability of our approach. The results from these comprehensive experiments consistently highlight the superior performance of method across various scenarios, further reinforcing its effectiveness. The detailed outcomes of these evaluations are presented as follows.
Table 4: Zero-shot evaluations of (2:4)-sp -::assppasstty other more tasks.
LLAMA-2-7B
DeepSeek-7B
Dense
Sparsegpt
Pruner-Z
MaskPro
Dense
Sparsegpt
Pruner-Z
MaskPro
WMDP
39.29
26.61
26.52
26.95
41.00
27.15
27.07
28.22
TMLU
29.58
25.03
25.13
25.38
37.17
25.99
24.36
25.37
Prost
23.60
24.26
24.03
24.41
28.19
28.22
27.62
29.57
Exams
21.04
23.65
23.65
23.65
23.65
23.65
23.65
23.65
AClue
27.47
25.33
25.31
26.24
32.34
27.17
26.88
27.31
ANLI-1
36.40
33.20
33.60
34.40
34.10
31.10
31.19
32.20
ANLI-2
37.20
34.10
33.90
34.10
36.60
33.70
33.20
33.50
ANLI-3
37.58
33.08
33.00
35.67
37.75
33.33
33.04
33.85
SCIQ
94.00
91.10
91.10
91.10
94.10
92.30
90.20
90.90
MathQA
28.24
23.72
23.55
23.95
29.48
25.93
25.12
26.76
Haerae
22.27
18.88
18.91
18.79
29.70
25.57
18.26
22.18
BoolQ
77.68
71.10
69.13
71.12
72.81
66.91
66.36
67.77
ComQA
32.92
20.80
20.08
22.03
36.69
23.10
22.95
23.18
LogiQA
25.65
21.66
21.78
22.89
25.04
21.73
21.35
22.58
COPA
87.00
81.00
79.00
79.00
84.00
86.00
84.00
87.00
WIC
49.84
47.81
47.22
49.84
51.10
48.00
48.81
49.06
WSC
36.54
36.54
36.54
36.54
64.42
36.54
36.54
36.54
CB
42.86
41.07
39.29
57.14
55.36
42.86
43.44
48.21
MultiRC
56.97
57.20
56.37
56.93
57.22
57.20
57.20
57.20
RTE
62.82
58.48
59.12
61.37
67.87
63.43
63.15
66.32
Mutual 70.84
68.01
67.44
68.53
71.30
67.43
67.24
68.33
WebQS
0.0586
0.0541
0.0544
0.0566
0.0876
0.0468
0.0226
0.0494
In this experiment, we evaluate the performance MaskPro across diverse set of tasks to com- prehensively assess its effectiveness on LLaMA-2-7B and DeepSeck-7B. The experimental design includes a variety of downstream tasks. MaskPro consistently demonstrates superior performance over competing methods, such SparseGPT and Pruner-Zero, in the majority of datasets. The method effectively balances accuracy and computational efficiency, achieving more favorable out- comes without compromising memory constraints. This consistent performance across multiple tasks highlights the robustness and generalizability MaskPro handling different scenarios. On smaller datasets, the performance gains of MaskPro relatively moderate, the evaluation is constrained by limited sample diversity. However, when tested on larger datasets with extensive testing samples, MaskPro consistently demonstrates substantial improvements over baseline methods.
16

---

Published as a conference paper at ICLR 2026
A.4 TRAINING WITH DIFFERENT DATASET SIZE
In this section, we report the training results using different numbers of samples. In Figure 5 we present the loss residuals of training on LLaMA-2-7B model with 1, 32, 128, and 320k samples, respectively. We set batchsize as 32 for all others expect for as 1. All are trained for 10k iterations.
0.02
0.0
Train with 1 Sample
Train with Samples
7mm O w. ( - -mmo w.
0.00
-0.1
00.0
f(m, fmmo
-0.2
0..4
-0.06
-0.3
-0.08
-0.4
-0.10
-0.12
-0.14
4000 Iteration
8000
2000
4000
6000
8000
10000
Iteration
(a) Training set =
(b) Training set = 32.
Am; O w. E)- f(mo O w, 5
0.02
Train with 128 Samples
Amm O w. {) fmmo O w, E)
Train with 320k Samples
0.000
-0.02
-0.04
-0.050
-0.06
-0.075
-0.08
-0,100
-0.125
-0.12
-0.150
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Iteration
Iteration
(c) Training set = 128.
(d) Training set = 320k.
Figure 5: Loss residual curves of training on LLaMA-2-7B model with 1, 32, 128, and 320k samples.
It can be observed that MaskPro does not require large number of training samples. Even with just sample (in a single minibatch), can complete training and achieve stable performance. The loss on single training sample steadily decrease, but this does not necessarily imply continually decreased loss on the test dataset. In fact, despite the persistent reduction in training loss, the test set performance may have already stabilized. In Figure (b) of the main text, we report the testing results of the learned mask on the Wikitext dataset. Next, we evaluate the zero-shot accuracy on series of downstream tasks, shown in Table
Table 5: Zero-shot evaluations of masks trained with different dataset size on LLaMA-2--BB
HellaS.
RACE
PIQA
WinoG.
ARC-E
ARC-C
OBQA
Avg.
320k samples
46.18
37.13
73.07
65.82
66.12
32.85
26.20
49,62
128 samples
46.10
37.03
72.47
65.62
65.49
32.25
25,80
49.25
samples
46.32
36.89
72.80
65.27
65.95
32.66
25.80
49.38
sample
46.39
37.61
72.96
64.64
65.70
32.59
24.40
49.18
It can be observed that although the performance slightly declines, overall, training with just 1 sample can still maintain satisfactory results, and in some datasets, the performance is even slightly higher.
17

---

Published as a conference paper at ICLR 2026
A.5 PERFORMANCE OF (4:8)-SPARSITT
In this section, we report the results for (4:8)-sparsity in Table 6 and corresponding training loss curves in 6 The training hyperparameters consistent with those reported in Table
Table 6: Zero-shot evaluations of (4:8)-sparsity. The MaskLLM method suffers from severe memory explosion and exceeds the memory limitation of 8x A100 GPUs (> 640 G).
Wiki.
HellaS.
RACE
PIQA
WinoG.
ARC-E
ARC-C
OBQA
LLAMA-2-7B
8.71
57.15
39.62
78.07
68.90
76.35
43.34
31,40
MASKLLM
MAGNITUDE
61.99
46.05
35.31
72.20
62.27
64.81
34.07
25.80
SPARSEGPT
14.99
48.19
38.55
73.78
67.72
68.15
36.01
27.80
WANDA
15.28
47.04
38.18
74.14
66.77
67.00
34.56
26.40
GBLM
15.21
47.32
37.51
74.16
67.56
67.13
34.56
27.20
PRUNER-ZERO
15.10
47.82
38.13
74.07
67.23
68.18
34.97
27.20
MaskPro
13.73
49.51
39.33
74.65
68.43
68.64
35.92
28.20
DEEPSEEK-77
9.70
56,94
39.62
79.27
70.40
75.25
43.60
32.60
MASKLLM
MAGNITUDE
109.37
45.32
32,06
72.42
61.64
56.31
32.68
23.40
SPARSEGPT
14.67
48.36
38.09
75.24
65.82
70.20
36.69
29.20
WANDA
14.76
49.09
38.47
75.46
64.88
68.48
34.22
27.20
GBLM
14.74
49.03
38.76
75.73
65.11
68.18
34.13
27.00
SAUNER-ZEROO
14.85
48.22
38.32
75.12
65.66
69.23
35.50
27.80
MMskPPr
13.89
50.97
39.25
75.87
66.27
69.51
36.89
29.80
imm w.W. E) f((m O w. 5
Train with 512 Samples
f(m. O W. ( - f(mm O w. EE)
Train with 512 Samples
-0.02
-0.04
-0.06
.0.04
-0.08
2000
4000
6000
8000
10000
2000
4000
6000
8000
10000
Iteration
Iteration
(a) 2LaMA-2-7B.
(b) DeepSeek-7B.
Figure 6: Loss residual curves of training for the (4:8)- -pprrsity.
A.6 MEMORY SCALABILITY
In this section, we report the memory scalability in Table
Table Memory (GB) reugired for training on DeepSeek-78
MaskLLM
MaskPro
(1:4)-Sparsity 266.35
36.82
(2:4)-Sparsity
339.56
36.82
(4:8)-Sparsie)
>640.00
36.95
The MaskPro method, due to its linear probability modeling, almost does not cause memory growth the (N:M) ratio scales. When training the (4:8)-sparsity DeepSeek-7B model, MaskI.LM has
18

---

Published as a conference paper at ICLR 2026
encountered OOM (Out of Memory) on 8x A100 (>x40G). In contrast, MaskPro achieve the expansion with almost no additional memory overhead.
A.7 PERFORMANCE OF
Moreover, we provide the (8: 1spaarstyy pattern to evaluate the performance of our proposed MaskPro method. This setting involves significantly larger combinatorial spaces which can greatly support the efficiency of MaskPro.
Table 8: Zero-shot evaluations of -sparisy on LLaMA2-7B.
HellaS.
RACE
PIQA
WinoG.
ARC-E
ARC-C
OBQA
Avg.
LLAMA-7B
57.15
39.62
78.07
68.90
76.35
43.34
31.40
56.40
MAGNITUDE
52.27
35.02
72.74
64.48
67.68
37.03
27.20
50.92
SPARSEGPT
50.19
39.04
74.43
66.22
70.45
36.43
28.80
$2.22
WANDA
49.77
39.14
75.30
66.61
70.62
36.18
28.80
52.35
GBLMM
49.51
39.90
75.68
66.38
69.91
36.43
27.60
52,20
PRUNER-ZERO
50.12
38.68
75.22
66.13
69.93
35.48
27.80
51.91
MaskPro
53.15
39.23
76.15
66.56
72.87
40.13
29.60
53.96
LLAMA-13B
60.05
40.48
79.11
72.22
79.42
48.46
35.20
59.28
MAGNITUDE
55.43
37.51
74.48
66.06
68.94
38.05
27.60
52.58
SPARSEGPT
54.24
40.38
77.15
70.19
75.08
41.31
31.00
55.62
WANDA
54.50
39.62
77.09
70.09
73.19
40.36
30.80
55.09
GBLM
54.45
39.18
76.35
69.92
73.75
40.07
29.60
54.76
PRUNER- ZERO
54.11
38.64
76.28
70.41
72.92
40.55
30.00
54.70
MaskPro
57.35
39.92
77.83
70.68
76.45
43.26
30.60
56.58
Under this sparsity pattern, the memory requirement of MaskLLM becomes extremely large, even exceeding the resource demands commonly used in the community to train models with hundreds of billions of parameters. Moreover, our MaskPro approach introduce minor training cost, while achieving better results than rule-based methods.
A.8 PERFORMANCE ON LARGER SCALE MODELS
In this section, we present the results of applying MaskPro larger models, specifically the 13B and 30B variants. We retain the same hyperparameter settings used for the 7B model, with the only adjustment being a slight tuning of the initialization logits magnitude.
Table 9: Zero-shot evaluations of (2:4)-sparsity 13B/30B models.
HellaS.
RACE
PIQA
WinoG.
ARC-E
ARC-C
OBQA
Avg.
LLAMA-13B
60.05
40.48
79.11
72.22
79.42
48.46
35.20
59.28
MAGNITUDE
50.10
36.84
71.76
61,88
62.29
31.74
23.40
48.29
SPARSEGPT
47.73
38.95
73.61
69.22
69.95
36.35
27.40
51.89
WANDA
46.24
38.47
73.94
67.32
68.73
34.13
24.20
50.43
GBLM
46.65
37.97
73.46
69.04
69.33
34.75
25.80
51.00
PRUNER- -ERO
46.15
38.85
73.13
67.24
67.52
33.89
25.20
50.28
MaskPro
49.24
38.91
75.12
70.33
71.85
38.26
27.40
53.02
L.LAMA-3R
63.36
39.14
80.63
75.85
80.64
51.45
36.40
61.07
MAGNITUDE
49.57
35.69
70.24
65.59
57.32
31.66
27.80
48.27
SPARSEGPT
55.25
37.77
77.45
73.68
75.25
43.27
31.80
$6.35
WANDA
54.18
40.00
77.69
73.24
74.24
42.15
31.60
56.16
GBLM
54.68
37.35
75.24
73.12
74.68
42.32
30.80
55.46
PRUNER- -ZERO
53.69
37.13
75.86
73.04
74.23
41.25
31.20
55.20
MaskPro
59.76
37.28
78.24
73.32
76.83
45.65
33.20
57.75
19

---

Published as a conference paper at ICLR 2026
Notably, MaskPro remains highly effective when applied to models at the 30B scale. This demonstrates the robustness and scalability of the proposed probabilistic formulation. Furthermore, due to the linear probability modeling and the use of selicy-gradient-ber basedd optimization, MaskPro achieves this performance with significantly reduced computational overhead. In particular, the train- ing process requires far fewer resources compared methods that rely dense mask representations or exhaustive combinatorial search. These properties highlight the practical advantages of MaskPro, especially in large-scale scenarios where both memory efficiency and training stability are critical.
A.9 SENSITIVITY OF TRACKER COEFFICIENT
In this part, we demonstrate the sensitivity studies of the tracker coefficient a. In our PG update, the parameter a is used to track stable estimate of the current baseline and prevent from being overly influenced by the stochastic variance of sampled losses. Conceptually, this plays the same role as 31 or 32 in the Adam optimizer. To examine its sensitivity, we conducted the following set of experiments:
Table 10: Sensitivity studies of tracker coefficient a.
0.7
= 0.9
0.95
0.99
0.995
LLAMA-7B
34,25
48.28
49.37
49.62
49.21
LLAMA- 13B
38.68
51.23
52.78
53.02
52.74
We find that using a = 0.99 consistently across all tasks provides the most stable and reliable performance. Therefore, only report the selection of 0.99 for reproduction in the main text. This hyperparameter requires almost no additional tuning.
A.10 SENSITIVITY OF LOGITS MAGNITUDE
In this part, demonstrate the sensitivity of the logits magnitude C. In the initialization, the parameter C' is used for stable sampling space. A detailed explanation provided in Appendix A.2. If is too small, single sampling step has high probability of producing a poor mask, which can lead to severe imbalance between positive and negative samples during training, ultimately hindering the learning process of combinatorial optimization. Therefore, choosing sufficiently large C' during initialization allows the training to remain stable. We evaluated different values and the results are as follows:
Table 11: Sensitivity studies of initial logits magnitude C.
C =
9
10
C'
11
12
LLAMA-7B
49.17
49.62
49.59
49.55
LLAMA-13B
52.45
52.94
53.02
52.99
A.11 ABLATION STUDIES OF INITIAL MASK mo
In this part, we evaluate how the diffferent initialization mask ma affects the results. Unlike gradient- based methods, RL methods typically converge more slowly, so good initialization can significantly shorten the training process.
Table 12: Ablation studies of initial mask mg.
Random
Top-K
Wanda
GBLM
Sparsegpt
LLAMA-7B
30.27
45,97
46.71
46.56
47.97
MASKPRO
36.35
48.35
49.33
49.45
49.62
IMPROVEMENT
+6.08
+2.38
+2.62
+2.89
+1.65
20

---

Published as a conference paper at ICLR 2026
In practice, when using SparseGPT for initialization, the model converges in roughly 10000 steps. With TopK initialization, extending the training over 20000 steps yields a relatively smooth result. We can see that random initialization can also train the mask, but the training process is slow. We further conduct longer training experiment specifically for random initialization, and the results are as follows:
Table 13: Long-term training on the random initialization.
T 20000
T 30000
T= 50000
T=70000
ACCURACY
36.35
38.43
40.74
42.37
This training process is quite lengthy, and we estimate that completing the full experiment would require at least 300000 steps. Such behavior is consistent with the theoretical convergence rate of RL-based methods, which why we do not encourage training from random initialization. We hope that these two experiments address the reviewer's concerns: it is not that RL-based methods cannot be trained from random initialization, but rather that unnecessary, as simple priors such as top-K can significantly shorten the training cycle.
A.12 TRAINING WITH DATA SAMPLE FROM DIFFERENT INITIAL MASK me
In this part, additionally evaluate the stability of the training process of "with data sample" from different initialization.
Table 14: Ablation studies of dataset size on different initial mask mo.
Random
Top-K
Wanda
GBLM
Sparsegpt
WITH 128 DATA SAMPLES
36.35
49.35
49.33
49.45
49.62
WITH DATA SAMPLES
35.97
49.12
49.04
49.21
49.18
We would like to clarify that MaskPro is indeed not very sensitive to the number of samples. The essence of RL-based methods lies in accurately estimating and constructing the reward, rather than relying on large data volumes. While we do not deny that using larger dataset may yield further improvements, the performance obtained with only a few hundred samples already very close.
A.13 TRAINING TIME PROFILES ON MASKPRO
In this part, we mainly show the training time profiles of our MaskPro method under different patterns and models. We torch. mnch.multinomial function for (N:M)-sparsity sampling, which simulates the sampling process through a lookup-based mechanism and provides high accuracy. The forward pass implemented through standard wrapper function. Specifically, we wrap the linear layer with an additional mask parameter and integrate the mask computation directly inside the linear operation. This design avoids modifying Py PyTorch's 's computation graph and enables efficient inference. The logits updates are computed entirely through matrix calculation, and PyTorch's h built-in libraries already provide the necessary parallelization. To further illustrate the implementation details, we report the per-step training time follows:
Table 15: Averaged time required in each step on (2:4)-Sparsity.
Mask Sampling
Forward
PG Update
TIME
RATIO
TIME
RATIO
TIME
RATIO
LLAMA-7B
2.328s
85.94%
0.062s
2.29%
0.319s
11.77%
LLAMA-133 -33 4.739s
85.83%
0. 0.139s
2.52%
0.644s
11.65%
21

---

Published as a conference paper at ICLR 2026
Similar to the general RL method (Fan et al. 2024), the main souree of time consumption comes from the sampling process. We also evaluated the sampling performance under different sparsity patterns, shown in the table below.
Table 16: Mask sampling time each step on different (N.M)-Sparsity.
(2:4)-Sparsity
(4:8)-Sparsity
(8:16)-parsity
LLAMA-7B
2.328s
1.334s
0.794s
LLAMA-13B
4.739s
2.6928
1.574s
We can observe that doubling the model size roughly doubles the sampling time. Another interesting observation is that the sampling time of (N:M)-Sparsity depends on M. With the same model size, larger M leads shorter sampling time, due to parallel optimizations in the sampling process. For d-dimensional model, there are total of d sampling groups. Although increasing N and M makes M each group more expensive to sample, the total number of parallel groups decreases proportionally. This reduction in the number of groups results in a more favorable computation pattern for hardware. Consequently, for more complex (N:M)- the time required for a single sampling step can actually be lower. Large M is GPU-friendly selection.
A.14 ALTERNATIVE STRATEGIES FOR ACCELERATING SAMPLING
In this part, we additionally explore two alternative accelerated sampling strategies along with their corresponding results. Sampling is the primary computational bottleneck of RL-based methods. Therefore, we explored several alternative acceleration strategies speed up the training process. and their effects are summarized below.
Table 17: Acceleration of mask sampling and their corresponding performance on (2:4)-Sparsity.
torch.multinoiial
Naive Gumbel-Topk
Eaussion-Topl
TIME
ACC.
TIME
Acc.
TIME
ACC.
LLAMA-7B
2.328s
49.62
1.821s (1.27x)
49.24 (-0.38)
1.496s (1.58x)
48.84 (-0.78)
LLAMA-13B
4.739s
53.02
3.645s (1.30x)
52.59 (-0.43)
3.061s (1.55x)
52.22 (-0.80)
Within an acceptable error range, the training time can be further reduced. However, we still recommend using higher-precision sampling methods, as the current training time requirement of MaskPro is already quite reasonable. On the impact of randomness on experiments, RL methods rely on sampling, so they are generally less sensitive to random seeds compared with gradient-based methods, and tend to exhibit stronger robustness across settings.
B DETAILED DESCRIPTION OF WITHOUT-EELLCEMEE PROBABILITY p(m[T)
This section aims present a specific form of p(m/x) and its related gradient log (p(m) ()). Note that in Eq 8 we define p(m .) I (m;));))))))))))))))(() where m, == + az, for any E / and p( ((m./mr mr) denotes the probability our MaskPro generating the mask m; under logits m.. Therefore, before presenting the details of p(m /T), we firstly investigate the probability (m)..
B. DETAILED DESCRIPTION OF (T.)
It is worth noting that the mask vector m, NNM such that can assume m.; EiEINi Cb where ej denotes the j-th basis vector of the space IR IRIM 1xxM k; E [M], Vi e [N] and k * 12 + ... KN- In other words, (ki,..., ,...., is N-size subset of [M] (1...., M).
22

---

Published as a conference paper at ICLR 2026
From the definition of m,, know that m, +;= adj- Furthermore, according to Eq
in Section C.I we also can know that, in order to ensure that ① j== Bijj EiE[N] ex; we typically require a one-to-one assignment of the previously defined N distinct basis vectors (ek....., eky } to (ar. J. In general, there are N! different ways to perform this matching.
To better illustrate our results, we introduce the concept of permutation from group theory to represent these N! one-to-one assignment. More specifically, for any one-to-one assignment from fokr (Ck)....,.0 } to {a..1...., N we represent (kI...., ky ). Here, each bijection a means that we match each basis vector er(j), Vj E [N] to the j- th sampled vector a in the sampling-wite without-rrplacemmen process, namely, a,, = ex()) Vj [N]. Moreover, we denote all such bijections as B (m.), that is to say,
bijective function o =fl,.... .......A. .N.
Bx (m:) == {o is a bijection from [N] to (kI,...,.w) kN}}.
With the notions of a and x (m,). we next present the specific form of p(m, At first, like RLAM
Section 4.1, we assume मं (T..1..... TT,,r) and define vi(T;; Σ e5s
" of as
softmax function. Then, for specific assignment a E Bx (w(m),, we have that
Pr (far.) = eooi) II
(15)
1 Σe= (TTT)))o(a)
where the symbol 'Pr' denotes the probability and [y:)m denotes its j-th component. Moreover, Eq 15 when 1, we define the summation Σ [ [v((s))o((a and simultaneously specify
It is important to note that in Eq 15 the value 1->: (1) stands for the j-step sampling- without-replaacmmet probability. Finally, from the result of Eq
have that
p(m(/T))= = Σ
Pr ((ar) = ( (s) (m)
rEBx(mi)
(16)
(T(T)))))))
II
Ex= (TTH))
(m.)
B.2 DETAILED DESCRIPTION OF p(m(FT)
Due to that p(m|T)) = II. 37 pmm. T.)) and Eq
16
we then can show that
( (T)))o Σi= (MT)))occo)
)
p(m|#) pmm
Σ
II j=1
(17)
BEBN (m.
B.3 COMPUTE THE GRADIENT V log (p(m)T))
Note that in Eq 10 in order to update the logits via mini-batch stochastic policy gradient descent, we need to frequently compute the gradient V. log (p(m[r)). Thus, in this subsection, we give the detailed form of this v, log (p(m( [)).
At first, due that p(m( = 11 * such that
p(m mT), can know log (p(m(z)) = Σ =i log (p(m,/m))
V, log (0 = V log (mm) Vny log
log (p(m T
Therefore, in the subsequent part of this subsection, show the specific form of Vz, log (p(m. (T)) for any i d1 Like Section B.I assume that m, EsE[N] Ck2 where k; E [M] [M] and kn # k2 # ... # kx.
23

---

Published as a conference paper at ICLR 2026
Then, when nT (Ta,,) 1......... Te,M), for any lo [M] [f.....,M we have that
(log log (p(mi(mT))
o(v (m)i
1
dTi
p(mm (T))
① EoeBy((..)
[1 1
(T.))(Ts))
OTick
① II j=1 E' (M(sso())
(18)
[Ta
05igh
or Bay (m.)
① (II (((T)))())) θT,
N
Σ
II
1
Ex (1b( (1b(m)))((a)
(mB)
j=1
a II
1- Ex i
p(m: TT) (m.)
17 1-52 ((
a (wr))a and
Next, we compute the definition of (TT), we show
m(a)
in Eq 18
At first, from the
d
1ms.1
a [1)(m)) (+(T))) [bb(T))] ,Kk + E [M];
OTik
(19)
0 ( [+)(al)l (1 (1-(((T))l)))
As result, have that
a (I==110(Ts))oos
[IMm))00)) ) (1IKE (kili=n) N (MmT)ll) ) (20)
θTi
where is the indicator function.
(
((()))
As for
we have that
[IN
j=1 Σ1= (o
0 ( (
0
(21)
1=1 -i-1 (TTTllooa))
EE= [e) (TTT))l((a)
[+((T)) (Ii o- o1() - E (TT))o(a)))
1 - (
1 1- 1ー [1((T)))
where II is the indicator function and o~ denotes the inverse mapping of a. Especially when E [ki...., k. e.g., & = ke where c [N]. we set o -(k) = c. As for & # [ki..... kN }. we
set -kk) such that Tíj > o -k)) 0 for any j E [N].
24

---

Published as a conference paper at ICLR 2026
Merging Eq 21 and 2q into a log (p(m.
18
18
finally have that
N
[1 (ki) N [T((T))R))
E
II
Σ (TT(Ta)l((a)
p(m;/T=;)
BEBN (m))
j=1
N (xb(TTs) (o(s)
N [+TTallle (Ili (k)] - ( (TT
Σ dEBN
Ex (1b (TTT)))l(aa)
p(milTT)) - I (
/=1
where p(m:/ (T)) Eme
and m EiE(Ny ek &NN.M
C
PROOFS
In this Section, we provide the detailed proofs of the main theorems.
C.I PROOF THEOREM I
This subsection aims to present rigorous proof for the representation Theorem 1 Before going in the details, we first assume that, Eq[4 a; ek;, Vi E [N] where k; E [M], Vi [N] and 1 # 12 # ... # kN. With this assumption, then can show that,
(1M -aa) = 1M
0j
(22)
(K
We verify this Eq 22 by induction. Firstly, when N = 1, Eq. 22 naturally holds. Subsequently, we
assume that when N M, Eq 22 is right. As result, we can show that, when N m+1 < M
m++1
m
(1M a.)
(1M ai) = O (1M a:)
(1M am+1)
1M
Σ
Bj
3 (1 @kmn1)
1M
Σ
Cj
Ckm+1
Σ (Okm-1 @ e3) -IM-
Σ
Cj
jE (ks.....mm
jE(Ki
where the final equality follows from that @km+1 holds for any S M.
ej = 0, when kmm+1 As a result, the Eq.
22
According to the result of Eq.
22
we can easily have that
N
+ a
1
a.
Σ ej.
(23)
i=1
JE(MI....Kxx)
Therefore, from 23
23 we can infer that, when E (ei...., GM Vi E [N] and + a2 .... + aN. i= a; BIXM and @ azlli N such that @ a; NN:M where B denotes the Boolean set. Furthermore, for any binary vector b E MNNM
redefine b EiE(N es, where 8; E [M], Vi [N] and 81 o 82 +... ... + 8N. Then, we set a = es, for {l,..., ni,, acoording
to the result of Eq 23 we can have
+ a Σ A
C., b.
iE[N]
As result, we can establish that
N
SN:MM =
a; a, E (c,,..., CM a J Vi E [N] and aj + a2 #... # aN
25

---

Published as a conference paper at ICLR 2026
C.2 PROOF OF POLICY GRADIENT ESTIMATOR EQUATION
9
From the notations in Eq 8 we have that
(T)
(m [f (m w. [))
E8g [f (m ww (24)
It is worth noting that in right-hand side(RHS) of 24 only the component "p(m|r)" contains the unknown logits variable w. As a result, we have that
V V Bg if (m (..))p((m/))dm
IE [f (m O w, V.p(ml
-f |EE l/ (m w,
[xp(m]=) p(m (m//)ddm
p(m|
/ Ee lf (m ww (v log (p(m) r)) ( ((m|r)dm
EgwD.m (mm.mm)-pmmm...)) [f(m @ w. log (p(m|=))).
where the forth equality comes from the relationship that (log([(z))) a ((1oR(((0)) (x)
dx
[(
C.3 PROOF OF THEOREM 2
We first investigate the properties of the policy gradient update method applied this paper. As shown in Eq. 99. the general policy gradient satisfies the following equation:
( Bp(ml=) [f'm @ w)V log (p(m|r))),
where Eg [f O w. f(m c w).
In the training, due to the limitations of data samples, instead of computing the full loss [(m 3 w). we typically use small mini-batch stochastic gradient, that is,
9p f(m C w, V log (p(m|rr),
C.3.1 Loss RESIDUAL AND SMOOTHING TRACKER ARE UNBIASED ESTIMATORS OF V
We denote 9. as update via loss residual:
9. (f(m @ w, - w. ()) log (p(m) [)]
and for update via loss residual with smoothing tracker:
Is = (f(m @ )) - ((mo @ w, - 8) V log (p(m|=)),
8 8 ad + a) (f(m O w, [mo w)
It is worth noting that these two introduced additional terms s(mo @ w. and 5 are independent of the logits variable π such that we can know that
Vp(m|T) p p(m]
Ep(mlr) [Cf(mo w. + 8) loo (p(m[=))) (f @ w) + +)) p(m|T))
dm
=(f(mo OW p(m I (f O {) + 01 =.
Therefore, our proposed update using the loss residual with smoothing tracker remains an unbiased estimator of the standard policy gradient. Similarly, letting å 0, it degrades to the update with only loss residual, which is also a unbiased estimator of the standard policy gradient. In fact, our proposed enhanced version of the policy gradient update be viewed as auxiliary training method that introduces a baseline term, similar to the approach in reinforcement learning (Williams 1992). Results also consistent with existing similar analysis (Fan al. 2025c
26

---

Published as a conference paper at ICLR 2026
C.3.2 EFFICIENCY OF On
We first investigate the properties of updating via loss residual f(m @ w, f'mo w. w) We have the variance of the standard policy gradient Ip as:
Var [spl E f(m @ w, w. (V (p -)) - Bff(m e. (p(m) m))]
E
(V log ( (+TT)
Similarly, since E [91]] - (T(T), the variance of achieves:
Var [or] = [f(m @. f(mo @ w, (Vlog (p(ml (
Thus have:
Var [gr]
Var [sp]
[x(m)=)
V log (p(ml=))) Eg
(f(m w, f(mo o - f(m x)
v log (p(m) (m))² Rgg s(mo o w.
w. 2 3 w) (f(mp (f
Their relative magnitudes are determined by f(mo O w. 2m w. term and we have:
Var [9r] 2 Var [sp] when f(mo @ w,E) 2 2f( O w,E), Var [9r] Var [sp] when f(mo O w, < 2f ④ w,
Therefore, updating via loss residual can always achieve a lower variance when f'm @ w.E) E) > [mo @ g. This implies that the variance in the initial training stage significantly lower than that of the vanilla PGE 9p. enabling substantial accelcration. We also validate this in our experiments, where the vanilla policy gradient converges extremely slowly and barely learns effective information, while can achieve a rapid reduction in loss within only hundreds of iterations.
C.3.3 EFFICIENCY OF ISk
Similarly, since E [9sc] ( the variance of far achieves:
Var [asr] [ w, fmmo @ w. 8)² (V log (p(m)=)))) V2(T)
And we have:
Var [g<r]
Var [gp]
log (p(m)) ()))² Bg (J(m w. - {(mo @ w, 8)2 - s(m @ w, s>²
(p(m)r)
[p(m)))
log (p(m) m))² Bg [f o w. {
Ep(ml>) (V (V (p(m|=))) Eg [6²]
-2Ep(mas)
(V log (p(m) m)) m)))² Eg [f (m @ w, 6)0]
+ 2E 2(p(m)))
(V log (p(m) (T)) Eg [mo o w, 0)
2EE (m)
v log (p(m) T)))² Eg [m @ wwgg) C ww
[p(m)=)
(v log (p m TT) )
82 + 2E pm (Vlog ()(m))))) (f(mo @ ww) f(m @ w))
denoted by A220
denoted by
log 2 Eg [f(m, o w. (m @ w. - 2 w. ww(
(m/ss)
Varjai] Var[g,]
27

---

Published as a conference paper at ICLR 2026
Clearly, when 8 0. Var [ Var [ge]- Next, discuss the case where 8 * 0. The above expression can be viewed as a quadratic function w.r.t 6, i.e.,
Var [gw] Var [gp) V(6) A52 Bö (Var [++] Var (sp])>
According to the definition of 6, it is the moving average of the f(m C w. - fmm O w. term.
By considering f'm @ w, 2 2/ {mp w, w,g), we can intuitively examine the corresponding magnitude relationships through the function plots. As shown in Figure 7 when [6] < l we always have Var [9xe] < Var [9t]- Furthermore, if å maximum. Therefore we have:
B the extent of variance reduction will reach its
24
(V log (p(mm)))) (J(m @ w) [(mo @ w)) (V log (p(mlz)))²
Bp(m(n)
8*
2AA
(pp())))
(V log (p(m) (p(m
(f(m w) s(mo @ w))
E (p(m(r)
(V log (p( (m /
[f(m o w) f(mo @ w)],
where p(m| = (V p [=)
Sycodss) [ (m))))
Clearly, d* can be interpreted as the expectation of f(m @ w) w O w) under the optimal distribution ((m/F), or equivalently, as the weighted average over all possible cases. It feasible to accurately measure this distribution. When the original distribution p(m|TT is known, the optimal distribution can be derived; however, the corresponding computational overhead calculate is prohibitively high. Therefore, we track all stochastic sampling in the training process and calculate the moving average of each f (m, @ w, - f (mg O w. {) as a compromise. After a long iteration and enough samplings, 8 can achieve significant and stable performance.
Therefore, we have Var [gw] S Var [yr] < Var [gp].
Varlg..I-Varlg,1
Varlg.I-Varrgal
8 0
<
8
0
if
B
0
6
6
Varlg.I-Varlg.1
Varlg.I-Var(g.J
(a) When B < 0.
(b) When B > 0.
Figure 7: Schematic illustration of the optimal solution.
28

---

Published as a conference paper at ICLR 2026
And the theoretically maximal variance reduction can be expressed as:
max (Var [gpl Var [9vr]]
-pp(m)mm
V log (o(m/T)) Eg [f O w. (f -w - 2 ow
v log (p(ml [))) (f(m w) - f(mo o w)))) ?
[p(mlss)
pp(mlm)
(V log (p(m [
(V log log (p(m|T))) (f(m @ w) f(mo 2 w)))) 2
(pp(m))) m)
Var [9p]
Var [gr] +
E (m)r)
(V log (p(m|z)))
29


Reviews:
[{"review_id": "ACb518xVqz", "reviewer": "Reviewer_NXWQ", "summary": "The paper introduces a new probabilistic framework for learning hardware-friendly semi-structured sparsity in LLMs. Unlike prior methods, MaskPro reduces both memory and computation overhead by modeling sparsity as an N-way sampling without replacement from categorical distributions over M consecutive weights. It leverages a linear-space parameterization that reduces memory complexity to linear. and optimizes via a refined PGE that replaces the raw loss with a loss residual tracked by a moving average, improving stability and variance reduction. Extensive experiments on multiple 7B LLMs show that MaskPro outperforms baselines while approaching the their accuracy at 10x less memory and training cost.", "questions": "How sensitive is MaskPro's performance to the choice of $\\alpha$?\n\nDoes the sampling-without-replacement process still scale linearly in memory and remain computationally tractable for large M? This needs further discussion.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "Sound theoretical explanation, linear memory efficiency, and comprehensive validation experiments.", "weaknesses": "Experiments focus mainly on 2:4 sparsity and 7B models; scaling to higher sparsity ratios or larger models, e.g., 70B, is not demonstrated. (might due to hardware constraints)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "hoCblFzPPQ", "reviewer": "Reviewer_4Q5R", "summary": "This paper presents MaskPro, a linear-space probabilistic framework for learning strict N:M sparsity in LLMs. The work addresses the limitations of existing approaches, rule-based methods that introduce bias and learning-based ones that demand prohibitive computational cost, by proposing a lightweight probabilistic formulation. It models each group of weights with a simple categorical distribution and learns sparsity patterns through a refined policy gradient update that incorporates loss residuals and a moving average tracker, improving both stability and convergence. Empirically, MaskPro achieves strong results across multiple 7B-scale models, reaching good accuracy while reducing memory usage by an order of magnitude and requiring only a handful of training samples. The approach is conceptually clear, theoretically supported, and practically appealing.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The reformulation of the (N:M) sparsity problem into a linear-space probabilistic model is both interesting and convincing. It reduces the memory requirement from exponential to linear scale and lowers the data demand compared with prior learning-based methods.\n2. MaskPro demonstrates consistently strong performance across diverse benchmarks and model backbones, achieving comparable or superior accuracy to established baselines while maintaining high efficiency.\n3. The approach is remarkably data-efficient, as evidenced in Figure 3, where stable results are obtained even with minimal training samples, making the method highly practical for large-scale applications.", "weaknesses": "1. While the results on 2:4 sparsity are strong, it remains to be seen how well MaskPro generalizes to other configurations, such as 8:16. These settings involve significantly larger combinatorial spaces, roughly 12,870 combinations per group in MaskLLM, and would serve as a valuable test of the proposed probabilistic formulation’s scalability.\n2. It would be interesting to explore whether MaskPro can be extended to jointly optimize both the sparsity pattern and the LLM parameters. Such a joint optimization could potentially yield further improvements in performance and adaptability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "fmUiA16TSh", "reviewer": "Reviewer_rhLs", "summary": "This paper introduces MaskPro, a memory-efficient probabilistic framework for learning strict (N:M)-structured sparsity in large language models (LLMs). The method models each group of M weights using a categorical distribution and generates sparse masks via N-way sampling without replacement, reducing the memory complexity from exponential to linear. It also proposes a new policy gradient estimator (PGE) that replaces the vanilla loss metric with loss residuals and employs a smoothing tracker to stabilize optimization. Experiments on several 7B-scale models show moderate improvements in performance and significant reductions in memory usage compared to prior approaches such as MaskLLM.", "questions": "- Can the authors provide wall-clock runtime and throughput comparisons with MaskLLM and rule-based baselines?\n- How sensitive is MaskPro to the smoothing parameter $α$ and the initialization constant $C$?\n- What happens when masks are initialized randomly instead of using SparseGPT priors?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "- The proposed update rule based on loss residuals, combined with a moving average tracker, is an interesting and practical way to stabilize the inherently noisy gradients in policy-gradient-based mask learning. The authors demonstrate theoretically that this estimator remains unbiased while reducing variance.\n- MaskPro achieves substantial memory and computational efficiency, requiring roughly 36 GB of memory compared to over 300 GB for MaskLLM, while running on a single GPU. \n- The method performs well even when trained with very small datasets (sometimes as few as one sample), showcasing robustness to data scarcity.", "weaknesses": "- The experimental evaluation focuses solely on mid-sized (7B) models and only on (2:4) sparsity configurations, with minimal exploration of other ratios or architectures. There is no validation on larger or smaller models, no finetuning experiments, and no runtime or latency measurements to verify practical benefits. As a result, the empirical evidence supporting MaskPro’s claimed generality and scalability remains limited.\n- The method's success is critically dependent on a complex initialization strategy. Standard random or zero initializations are \"ineffective\", and the training fails if the initial logits magnitude $C$ is not large enough.\n- The comparisons with baselines such as MaskLLM, SparseGPT, and Wanda are not entirely fair, as these methods differ in their reliance on fine-tuning versus frozen weights, training data sizes, and initialization strategies. Moreover, MaskPro benefits from initialization using precomputed SparseGPT masks, which gives it an advantage that is not acknowledged or ablated. The paper would be stronger if these factors were controlled more carefully.\n- Although the authors emphasize MaskPro’s efficiency, the claimed linear-space scaling overlooks constant factors associated with sampling and softmax operations, and the actual runtime improvements are not measured. The claim that the method can train effectively with one sample seems implausible without relying on strong priors, undermining the claim of data robustness. More empirical validation is needed to substantiate these statements.\n- The algorithmic description omits critical implementation details such as how N-way sampling without replacement is realized efficiently, how randomness affects reproducibility, and how logits are updated in parallel. The notation (e.g., the $\\oplus$ operator) is nonstandard and lacks intuition, while figures are schematic and do not convey architectural structure or ablation results. Overall, the presentation is mathematically heavy and could be made clearer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "vElXH4ENgk", "reviewer": "Reviewer_HgWG", "summary": "The paper proposes **MaskPro**, a probabilistic framework to learn strict (N:M)-sparsity masks for LLM weights with **linear memory** in the number of parameters. It models each group of M weights with a categorical distribution and generates N active positions via sampling without replacement. It trains logits using a policy-gradient estimator on loss residuals with a moving average baseline to cut variance. A representation theorem rewrites N:M masks as a probabilistic sum of basis vectors; proofs claim unbiased gradients and reduced variance. Experiments on several 7B models show competitive zero-shot accuracy vs rule-based pruning and proximity to MaskLLM at much lower memory and data cost.", "questions": "1. How is $p(m\\mid \\pi)$ computed efficiently for N-way sampling without replacement? Provide explicit formulas and complexity, not just appendix references. Can you replace sampling with a differentiable top-N estimator while keeping linear memory?  \n2. How sensitive is performance to the initial mask $m_0$ and the logit magnitude $C$? Report ablations that start from random $m_0$, weaker heuristics, and different $C$.  \n3. Give wall-clock profiles: percent time in sampling, forward passes, and PG updates. Can cached logits or Gumbel-Top-k variants reduce sampling cost further?  \n4. Show end-to-end throughput and latency on A100/H100 with vendor N:M kernels before and after MaskPro, at fixed accuracy. Accuracy-only tables are insufficient to claim deployment readiness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- **Memory scalability.** Reduces logits storage from combinatorial $O \\left(\\binom{M}{N}^{d/M}\\right)$ to linear $O(d)$. This is the key lever for practicality.  \n- **Clean probabilistic formulation.** N-way sampling without replacement per group with explicit $p(m\\mid \\pi)$ and REINFORCE-style training.  \n- **Variance control.** Loss-residual update with a moving average tracker yields an unbiased estimator with lower variance than vanilla PG under stated conditions. Empirical curves support stability gains.  \n- **Robustness to tiny data.** Reasonable performance even with very few calibration samples; plots include the 1-sample case.  \n- **Empirical breadth.** Evaluates on multiple 7B models and diverse zero-shot tasks; MaskPro beats rule-based baselines by ~2 points on average and narrows the gap to MaskLLM.  \n- **Training efficiency.** Orders-of-magnitude lower memory and dataset size than MaskLLM; single-GPU feasibility reported.", "weaknesses": "- **Novelty vs known tools.** The representation of N:M masks via basis-vector sums and a categorical policy with REINFORCE is conceptually straightforward; the main novelty is the linear-space parameterization plus the residual baseline. The paper should better position against prior PG-based pruning and Gumbel/relaxation methods.  \n- **Sampling cost.** The approach still simulates N-way sampling per group at train time. The paper notes this as a time bottleneck but gives limited profiling or asymptotic-to-wall-clock mapping.  \n- **Baseline fairness and scope.** Results focus on (2:4) with (4:8) pushed to appendix. No strong hardware-level throughput benchmarks for inference with real N:M kernels. Comparisons use C4 for all methods, but details like per-layer masks and calibration choices could bias outcomes; MaskLLM numbers are partly taken from prior work.  \n- **Initialization dependence.** Logit init relies on a good starting mask (Top-N or SparseGPT). The sensitivity to poorer inits and to the magnitude hyperparameter is only partly explored.  \n- **Theory conditions.** The variance-reduction claim depends on a condition $f(m_t \\odot w,\\xi) > \\tfrac12 f(m_0\\odot w,\\xi)$. Practical validity across models and tasks is not deeply tested; switching $m_0$ mid-training is suggested but not ablated.  \n- **Model scale.** Only 7B models are reported. No 13B–70B results, where memory and sampling costs and accuracy regressions are more consequential.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)