BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
Authors: Chenyi yang
Wenjie Nie
Yuxin Zhang
Yuhang Wu
Xiawu Zheng
GUANNAN JIANG
Rongrong Ji
Venue: ICML 2025 poster (2025)
Keywords: Convolution Neural Networks
N:M Sparsity
Network Pruning
Primary area: applications->computer_vision
TL;DR: We introduce BAME, a novel N:M sparse training method that achieves state-of-the-art performance with far less training FLOPs.

Abstract: N:M sparsity stands as a progressively important tool for DNN compression, achieving practical speedups by stipulating at most N non-zero components within M sequential weights. Unfortunately, most existing works identify the N:M sparse mask through dense backward propagation to update all weights, which incurs exorbitant training costs. In this paper, we introduce BAME, a method that maintains consistent sparsity throughout the N:M sparse training process. BAME perpetually keeps both sparse forward and backward propagation, while iteratively performing weight pruning-and-regrowing within designated weight blocks to tailor the N:M mask. These blocks are selected through a joint assessment based on accumulated mask oscillation frequency and expected loss reduction of mask adaptation, thereby ensuring stable and efficient identification of the optimal N:M mask. Our empirical results substantiate the effectiveness of BAME, illustrating it performs comparably to or better than previous works that fully maintaining dense backward propagation during training. For instance, BAME attains a 72.0% top-1 accuracy while training a 1:16 sparse ResNet-50 on ImageNet, eclipsing SR-STE by 0.5%, despite achieving 2.37 training FLOPs reduction. Code is released at \url{https://github.com/BAME-xmu/BAME}

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
Chenyi Yang' Wenjie Nie' Yuxin Zhang'
Yuhang Wu
Xiawu Zheng'
Guannan Jiang? Rongrong Ji
Abstract
et al., 2016), object detection (He et al., 2017a), and se- mantic segmentation (Girshick al., 2014), etc. These progressions are chiefly driven by an augmented parameter burden and an increasingly onerous computational cost. Re- grettably, this tendency presents significant impediments for the deployment of DNNs on resource-constraing edge de- vices like smartphones and various Internet of Things (loT) apparatuses. Consequently, there has been proliferation of interest in model compression research (Hubara et al., 2016; Howard et al., 2017: Lin et al., 2020), with the explicit objective of reducing the model's computation and parame- ter complexity whilst preserving comparable performance to the original model, thereby alleviating the deployment tribulations experienced with DNNs.
N:M sparsity stands as a progressively impor- tant tool for DNN compression, achieving prac- tical speedups by stipulating at most N non-zero components within M sequential weights. Unfor- tunately, most existing works identify the N:M sparse mask through dense backward propaga- tion to update all weights, which incurs exorbi- tant training costs. In this paper, we introduce BAME, a method that maintains consistent spar- sity throughout the N:M sparse training process. BAME perpetually keeps both sparse forward and backward propagation, while iteratively perform- ing weight pruning and-regrowing within des- ignated weight blocks to tailor the N:M mask. These blocks are selected through a joint assess- ment based on accumulated mask oscillation fre- quency and expected loss reduction of mask adap- tation, thereby ensuring stable and efficient iden- tification of the optimal N:M mask. Our em- pirical results substantiate the effectiveness of BAME, illustrating it performs comparably to better than previous works that fully maintain- ing dense backward propagation during training. For instance, BAME attains a 72.0% top-1 ac- curacy while training a 1:16 sparse ResNet-50 ImageNet, eclipsing SR-STE by 0.5%, de- spite achieving x training FLOPs reduction. Code is released at https: httss::/ggtttbbccco com/ BAME-Xma/BR XAMEBAME.
Among these techniques, network sparsity has proven many successes (Han et al., 2015; LeCun et al., 1989; Luo et al., 2017) by zeroizing weights to yield lightweight, sparse net- works at different granularity levels, from fine to coarse. Fine-grained sparsity (unstructured sparsity) (LeCun et al., 1989; Ding et al., 2019) removes individual weights and is demonstrated to well retain performance even at high sparsity rates. Regrettably, the deployment of such fine- grained sparse networks onto mainstream hardware systems becomes exceptionally challenging, given the irregular ma- trix patterns created by sparse weights. In contrast, coarse- grained sparsity, otherwise known as structured sparsity, (He et al., 2017b; Lin et al., 2020) procures substantial accel- eration, purging whole convolution filters in the process (Liu et al., 2019; Lin et al., 2020). Nevertheless, structured spar- sity can experience severe performance degradation, espe- cially under high sparsity conditions. Recent developments indicate N:M sparsity as an auspicious avenue towards ef- fectively balancing the dual requirements of acceleration and performance retention (Zhou et al., 2021; Pool & Yu, 2021). By imposing a restriction of, at most, N non-zero elements within M sequential weights throughout the input channel dimension, N:M sparsity can substantially enhance the performance of structured sparsity, while concurrently assuring swift inference, ably facilitated by the N:M sparse tensor core (Nvidia, 2020).
1. Introduction
In recent years, the vision community has precipitously bol- stered the performance of Deep Neural Networks (DNNs) across various tasks, including image classification (He
'Equal contribution 'Key Laboratory of Multimedia Trusted Perception and Efficient Computing, Ministry of Education of China, Xiamen University *Contemporary Amperex Technology Co., Limited (CATL) 'Institute of Artificial Intelligence, Xiamen University. Correspondence to: Rongrong Ji <rrji wwt..imm.uuu.........m.
The crux of maintaining the performance of N:M sparse networks lies in identifying the optimal N:M sparsity mask. To achieve this, prevalent methodologies involve updating all weights during training to determine the most effective
Proceedings of the 4nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
Origin Mask
New Mask
Mask Update
Mask
Group Index
OBF
Weight
Weight
LMA
Prune
SAD Score
Revive
Sparse Training
Figure I: Framework of BAME. It iteratively performs weight
through Loss-Aware Mask adaption (LMA) and Oscillation-aware Block Freezing (OBF), which leads to stable and efficient location for the optimal N:M mask.
N:M mask, adopting a straight-through estimator to approx- imate the gradients of the pruned weights (Zhou et al., 2021; Zhang et al., 2023b) or learning the importance criteria for all weights (Zhang et al., 2022). Despite their efficacy, the computation of dense gradients invariably imposes sub- stantial training overhead. Notably, the reduction of training costs has been a focal research point within the sparsity comunity in recent years (Liu et al., 2021; Evei et al., 2020; Dettmers& & Zettlemoyer, 2019). With the ever-grove weer-growing size of cutting-edge models, the significant computational de- mands and energy consumption of training sparse networks escalating critical environmental, ethical, and financial concerns. Consequently, the development of efficient and scalable N:M sparse training methods is paramount, poten- tially even more urgent, to support the widespread accessi- bility and democratization of DNNs.
free efficient manner.
We conduct extensive experiments on validating the effec- tiveness and efficacy of BAME for N:M sparse training. The results show that BAME is able get state-of-the-art per- formance when training N:M sparse networks across wide range of sparse pattern, datasets, and prevailing DNNs, even with much fewer training FLOPs compared with existing work. Illustratively, BAME attains a 72.0% top-1 accu- racy while training a 1:16 sparse ResNet-50 on ImageNet, eclipsing SR-STE (Zhou al., 2021) by 0.5%, while using far less training FLOPs. Our work provides fresh insights in N:M sparse training without dense weight updates and we anticipate that BAME will not only equip practitioners with a robust training tool but also lay the groundwork for subsequent explorations into the training efficiency of N:M sparsity.
In this paper, we present BAME as a way of maintaining consistent sparsity in both forward and backward propaga- tion throughout the N:M sparse training process. As shown in Figure 1, BAME escapes from dense weight's update through block-aware N:M mask evolution. It specifically executes weight aandro--nnd- within each consec- utive M weights in order to adapt the sparse mask. Such mask evolution oceurs solely when the detrimental effects on loss caused by pruning certain weight is outweighed by the gain in loss from restoring another already pruned selectively adapt the mask of N:M
2. Related Work
2.1. Network Sparsity
By removing redundant weights to eliminate the parame- ter and FLOPs burden, network sparsity has emerged as fervent area of research over the last decade (LeCun et al., 1989; Han et al., 2015; Louizos et al., 2017). Traditional ap- proaches can broadly be classified into two categories based on their pruning granularity: unstructured and structured sparsity. The former involves the elimination of individ- ual weights at any location within the network, achieving sparsity at fine-grained level (Han et al., 2015; Lee et al., 2019; Ding et al., 2019). In essence, unstructured sparsity can rival the performance of their dense counterparts even at exceedingly high sparsity ratios, such as 90% (Mostafa & Wang, 2019). Nonetheless, the generated sparse weight tensors generally precludes acceleration on standard hard- ware platforms unless the sparsity ratio reaches or exceeds
weight. Concurrently, blocks, as some blocks are experimentally observed to ex- hibit frequent oscillations on their masks during training, leading to unstables loss landscape. To this end, we employ exponential moving averaging (EMA) to accumulate the in- cidence of mask fluctuations for each block, choosing those with fewer fluctuations for mask evolution ensure stable optimization for the N:M sparse network during training. In this manner, BAME can stably optimize the N:M mask while conducting N:M sparse training in backward-
2

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
95% (Wang). Conversely, structured sparsity achieves no- table acceleration by extensively removing entire weight rows or convolution filters (Luo & Wu, 2020; Lin et al., 2020). Regrettably, structured sparsity often leads to sub- stantial performance degradation at sparsity levels exceed- ing 50%, attributed to the constraints imposed on sparsity flexibility. Diverging from conventional sparsity granulari- ties, this paper delves into N:M sparsity that removes weight in mid-level granularity and has garnered significant re- search interest in recent years (Zhou et al., 2021; Sun et al., 2021; Pool & Yu, 2021).
network sparsity (Hoefler al., 2021; Evci et al., 2020; Han et al., 2015; Liu et al., 2021). The most represen- tative method RigL (Evci et al., 2020) prunes weights of smaller magnitudes during inference and subsequently re- grows the same quantity of weights based on their gradient values throughout backward propagation. Sparse Momen- tum (Dettmers & Zettlemoyer, 2019) employs the mean momentum magnitude of each layer as a benchmark for redistributing parameters. (Kusupati al., 2020) proffer layer-wise learnable thresholds strategizing the reallocation of parameters across layers. Moreover, (Liu et al., 2021) proposed to gradually increase the sparsity level during training to further enhance the performance of sparse net- works. While these approaches predominantly concentrate on boosting unstructured sparsity, our endeavor in this paper differs by targeting the training of N:M sparse networks, innovatively designing block-aware selection mechanism for pruning and reviving N:M sparse weights.
2.2. N:M Sparsity
The recent development of N:M sparsity upholds the conser- vation of N-out-of-M -M consecutive weights in DNNs (Nvidia, 2020; Pool & Yu, 2021; Sun et al., 2021: Zhou et al., 2021; Chmiel et al., 2021; Hubara et al., 2016; Zhang et al., 2022). Supported by the NVIDIA Ampere Core (Ronny Krashin- sky, 2020), N:M sparsity fosters superior storage and com- putational efficiency, establishing an immaculate harmony between model efficiency and precision, outdoing both un- structured and structured sparsity. To illustrate, 2:4 spar- sity can realize 2x speedups on an NVIDIA A100 GPU, while unstructured sparsity might further decelerate the inference speed at identical levels of sparsity. As trail- blazing work, ASP (Nvidia, 2020) employs a traditional tri-phase workflow encompassing model pre-training, high- magnitude weight extraction (Han al., 2015), and net- work fine-tuning. (Zhou et al., 2021) subsequently proposed to train N:M sparse network from scratch by introducing the Sparse-refined Frriigh-TThroug Estimator (SR-STE). More specifically, N-out-of-M weights of higher magni- tudes are selected in each forward pass, whileall weights are updated during the backward phase, utilizing the STE estimator, paired with a uniquely designed sparse penalty term. LBC (Zhang et al., 2022) further recasts N:M sparsity as a combinatorial problem, learning the optimal mask for each N:M block. MaxQ (Xiang et al., 2024) utilizes a multi- axis query to generate soft N:M masks during training to further improve the performance. Despite their effectiveness in preserving the performance of sparse networks, most ex- isting works require dense backward propagation update all weights to discover the optimal N:M mask, leading to massive training burden and memory cost. Our proposed BAME in this paper diverges from existing N:M methods as it performs both sparse forward and backward propagation during the entire training process, substantially alleviating the training cost.
3. Methodology
3.1. Background
We first recap basic preliminaries of N:M sparsity. For simplicity, we take the weights from a specific layer within DNNs for illustration. N:M sparsity forces at most N out of M consecutive weights in the weight row to have non-zero values. The weights can be therefore grouped into K blocks where each block contains M consecutive weights, denoted as W E RK x M And then, N:M sparsity can be formulated as multiplying W with a binary mask B E KR with the following objectives:
min L(W O OB; B) S.L. |Ba = N, W.B
(I)
where k 1,2, ..., K, @ is the point- clement-wise multiplication, 2()) denotes training loss function and 'D represents the observed training dataset, respectively. The zero elements in B indicate the removal of corresponding weights in the network, and vice versa.
Challenge of N:M sparse training. The crux of optimiz- ing Equation (1) falls into locating high-quality masks that correctly preserve important weights. As a pioneer work, ASP (Nvidia, 2020) chooses to mask out weights that have lower magnitudes, intuitively reducing the output derivation between dense pre-trained weights and N:M sparse weights. Nevertheless, the pre-training phase unavoidably carries huge training burden. In the literature, a more popular way to obtain the sparse mask is performing training-time weight selection by updating all weights (Zhou et al., 2021; Zhang et al., 2022; Fang et al., 2022; Zhang et al., 2023b). Partic- (STE) (Bengio et al.,
2.3. Sparse Training
ularly, the trriigh-tthrough-estt 2013) is leveraged to calculate the gradient of all weights, since the currently removed weights always receive no gradi- ent as their corresponding multiplied masks are Os. Formally,
Sparse training, which dynamically adjusts the sparse masks throughout the training process, has recently emerged as promising solution to enhance the training efficiency of
3

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
the gradients of W are derived
the pruned weights have the smallest magnitude within the N:M block.
OL
OC
d AWOB O
B
(2)
OW
B
To address this challenge, we introduce loss-aware mask adaption (LMA) that ensures weight saaniin--nddreviiiin always lead to loss decrease during training. LMA performs static sparse training in both forward and backward prop- agation, while only calculating dense gradient to perform mask adaption every AT' iteration. Here we use a specific N:M block W E R to illustrate the mask adaption proce- dure. Considering a currently preserved weight W where B 1, the loss change, denoted as AC(W upon its removal can be approximately derived using first-order Taylor expansion (Molchanov et al., 2017) as:
In this vein, all weights can be updated during the training process. By dynamically selecting weights with higher mag- nitude, such N:M sparse training can effectively boost the model performance, even without reliance on pre-trained weights. Despite recent efforts to further enhance N:M sparse training through additional norm constraints on pruned weights (Zhou et al., 2021) or gradual sparsity (Fang et al., 2022), one significant concern remains that dense proppproon and weight updates continue to incur sub- stantial resource consumption, posing challenges to scenar- ios with limited resources.
AL(W..)) k.i)
In this paper, we address the above hindrance of training inefficiency by proposing Block-Aware Mask Evolution (BAME), a method that ensures consistent sparsity through- out the forward and backward propagation phases of the N:M sparse training process. The unique contribution of BAME encompasses loss-aware mask adaption (LMA) that prune-and-revive weights to effectively decrease the train- ing loss, and oscillation-awarr block selection (OBS), limit- ing mask modifications within blocks demonstrating high- frequency mask oscillations, thus stabilizing the N:M train- ing proces. We introduce these two components as follows.
-VC( O B: D, B 0) L(W о B; D, B 1)
dc
<|C(W O B: D, B =1 a(W @ B)
(W 0)
Rí (BL, L(W @ B; D, B = 1)1-
(3)
If we ignore the first-order remainder Rı (Br,i = 0), then:
dC
AC(We.)) R a(W C B) W..l
(4)
Similarly, if we consider reviving currently removed weight Wa. k. back, the loss change AC(W,.)) = 0 can
3.2. I Loss-aware Mask Adaption
be derived as:
Owing to the great benefit of training cost reduction, adapt- ing the sparse mask during training while escaping from dense gradient calculation has been a hot topic within tra- ditional unstructured sparsity literature (Evei et al., 2020; Dettmers & Zettlemoyer, 2019; Liu et al., 2021; Jayakumar et al., 2020). The central philosophy of these methods in- volves performing global pruning and revival based on instantaneous gradient information every few training itera- tions. Specifically, several of the weights with the highest gradients among all pruned weights are restored and the same number of weights with the lowest magnitude among all retained weights pruned, therefore reducing the loss to the fastest extent.
AL(W -C[( O B; D. B =1) - - D, Baj = 0)
<CC(W O B: D. B 0)
OC
ac
0 0- - a(W @ B B)&
a( O B B)K.j
+ R1 (Bx, = 1) -L(W + B; D, B = 0
OC 2
207
O(W G B)RJ
(5)
where 77 is the current learning rate. Based on the preced- ing derivation, we can articulate the following conclusions. On one hand, Eq. (4) tells that for the preserved weights, pruning those with comparably minor A.....)) ensures the loss does not undergo notable alterations. This per- spective concurs with traditional network sparsity knowl- edge (Molchanov et al., 2017; Zhang et al., 2023a). Con- versely, considering the presently pruned weights, their re- vival will invariably benefit the minimization of loss as observed in the derivation of Eq. (5). Simultancously, it bears mentioning that restoring weights with significantly larger AC(W will induce the most substantial degree of loss mitigation. Therefore, at each mask adaption cycle,
Regrettably, prior methodologies for globally altering the sparse topology are unsuitable within the context of N:M sparsity. Following a fixed sparsity budget for each N:M of weights can only be car- ried out in each independent N:M block. This presents substantial risks for the mask adaptation: The gradients of the weights in the same block are likely to have minor differences due to the continuous input received, as is the magnitude of the weights. Hence, directly applying tradi- tional sparse methods can have high possibility of resulting in the recovery of weights yielding less loss benefit com- pared to the disruption caused by weight pruning, even if
block,
4

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
2:4 Sparsity
1:4 Sparsity
1:16 Sparsity
234
340
Black Numbbr
200
Block Number
400
Black Number
523
100
150-
300
62-
200-
100
30B
130
200
290
280
130
200
300
400
SAD
SAD
SAD
Figure 2: Sparse Architecture Divergence (SAD) of N:M blocks during training N:M sparse ResNet-50 on ImageNet. The majority of blocks remain under minimal mask variation, yet minority experience frequent mask oscillations.
we first calculate the loss-aware metric AL(WA.) of all weights in an N:M block using Eq. (4) and Eq. (5). Then, we adapt the mask of weights follows:
Algorithm 1: BAME for N:M Sparse Training.
Require: Weights W; Initial and final iterations for mask adaption t, and ty; Update interval AT.
Output : Sparse weights W
Bk.m= = {
0, if (F(CWW)) ()) < Top(AC(W A.M.. M- M),,,
for E [t....., do
1, otherwise,
if % AT== 0 then
(6) where m 1,2, .... M and B is the updated mask. Such mask adaptation perceptively prune-and-revive weights by looking at the effects imparted on the loss, conducting an inclusive ranking within each N:M block. Paradoxically, preceding arts that mandates the pruning of lowest magni- tude weights while refurbishing those with highest gradi- ents (Evei et al., 2020; Zhang al., 2023a), although jus- tified when executed across the entire weight matrix, may potentially be harmful for N:M sparsity with limited amount of weights in each block. To explain, the increment to the loss prompted by restored weights could indeed be consider- ably less than the disturbance to the loss distribution induced by pruned weights. Hence, our proposed LMA effectively realizes loss-aware optimization of sparse typologies.
Calculate AL(W) via Eq. (4) and Eq. (5) Obtain the adapted mask B via Eq. (6)
// pass--wware mask adaption Get the restricted mask B via Eq. (9)
s
// Oscillation-are freezing
6
end & N=W&B//
Apply N:M mask
8
Sparse Forward and backward propagation
end
culate mask fluctuation at e-th and + 1-th LMA cycle as follows:
SAD(B BB
B
B
(7)
It is also noteworthy that LMA necessitates the computation of dense gradients only intermittently, every AT iterations, while primarily conducting truly sparse training and weight updates at other times. This starkly contrasts previous N:M sparse training methods (Zhou et al., 2021; Fang et al., 2022; Zhang et al., 2023b), which obligate the calculation and stor- age of full gradients at each stage of optimization, culminat- ing in significantly greater training impedance relative to LMA.
where B denote the k-th block's mask at c-th LMA cy- cle. In Figure 2, we show the accumulated SAD score of different N:M blocks during N:M sparse training. The ob- servation reveals significantly higher frequency of mask alterations occurring in a certain number of blocks com- pared to others. On reflecting upon the primary intent of LMA, the apex aim during the LMA process constitutes the pruning of weights of lesser magnitude, making way for the revival of a more significant one, thereby pinpointing an enhanced position while ensuring consistent training there- after. Nevertheless, some blocks endure recurrent deviations, alternately zeroing the weights, unquestionably inducing os- cillations in loss, and thereby impeding network training. Consequently, we harness the capabilities of Exponential Moving Average (EMA) to accumulate the episodes of mask perturbations across each block, electing those exhibiting lesser fluctuations for mask evolution, thereby ensuring sta-
3.3. Oscillation-aware Block Freezing
Other than LMA which enables efficient mask adaption in- ner each N:M block, we further stabilize the N:M sparse training process by oscillation-aware block freezing (OBF).
The impetus behind OBF stems from our observations of the
mask fluctuation for each block across different LMA cycles. In particular, we employ the Sparse Architecture Divergence (SAD) (Zhou et al., 2021) cal-
5

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
Table 1: Results for sparsifying ResNet-32 and MobileNet-V2 on CIFAR-I0.
Model
Method
N:M Pattern
Top-1
Epochs (Train)
FLOPs
FLOPs
Accuracy (%)
(Train)
(Test)
ResNet-32
Baseline
94.52
300
1x(322e16)
(x(133e))
ResNet-32
ASP
2:4
94.68
600
1.5x
0.51
ResNet-32
SR-STE
2:4
94.52
300
0.83x
0.51
ResNet-32
LBC
2:4
94.81
300
0.72x
0.51
ResNet-32
BAME(ours)
2:4
94.99
300
0.63x
0.51
ResNet-32
SR-STE
1:4
94.52
300
0.74x
0.26x
ResNet-32
Bi-Mask
1:4
94.43
300
0.49x
0.26x
ResNet-32
BAME(ours)
1:4
94.71
300
0.39x
0.26x
ResNet-32
SR-STE
1:16
92.92
300
0.67
0.11x
ResNet-32
Bi-Mask
1:16
92.77
300
0.37x
0.11x
ResNet-32
BAME(ours)
1:16
93.15
300
0.29
0.11x
MobileNet-V2
Baseline
94.55
300
1x(1.4e17)
Ix(4.8c7)
MobileNet-V2
SR-STE
1:16
93.14
300
0.67x
0.11x
MobileNet-V. V2
Bi-Mask
1:16
92.48
300
0.37x
0.11x
MobileNet-V V2
BAME(ours)
1:16
93.32
300
0.29x
0.11x
bilized optimization for the N:M sparse network during the course of training. Concretely, we devise vector O E gK equivalent in magnitude the count of blocks, devised for logging the frequency of mask alterations, as
non-zeros parameters budget linearly decreased from M to the targeted N of each block in the early T; iterations. Then, we perform BAME to find the best N:M mask from T; to Ty. At last, we set the masks all freeze and conduct static training for the N:M sparse network within the remained iterations. The workflow for performing BAME for N:M sparse training outlined in Alg. It should be noted that although BAME effectively reduces the overhead of N:M sparse training, the sparse weights during backpropa- gation may not maintain N:M sparsity due to the transpose operation, which introduces specific hardware implemen- tation challenges. To address this issue, we observe that BAME can effectively integrated with bidirectional N:M masks (Zhang al., 2023b) to ensure N:M sparsity through- out backpropagation; i.e., we impose N:M sparsity on the backward masks of BAME based on the magnitudes of the corresponding weights. We experimentally elaborate this in the subsequent sections.
0 %% 7) SAD(BE- BB),
(8)
where y the momentum of EMA updating. Then, we restrict 3 proportion of N:M blocks with the highest oscil- lation frequency from being updated by LMA as:
B if > Top 1 [8-A(),
B
B E otherwise.
Furthermore, within the mask adaptation selection of LMA, the occurrence of gradients sporadic, implying the prospect of particularly extraordinary gradient for speci- fied weight given stage. This could conceivably prompt an incorrect pruning of substantial weight, thereby precip- itating noteworthy effect network performance. Conse- quently, we confine LMA to transpire solely within weight blocks of lesser magnitudes to circumvent such inadvertent erroneous mask adaptations.
4. Experiment
4.1. Experimental Settings
Datasets and Networks. We validate the effectiveness of BAME by using to train N:M sparse networks on image classification tasks the CIFAR-10 (Krizhevsky et al., 2009) and ImageNet- datasets (Deng et al., 2009). For the networks, we sparsify ResNet-32 (He et al., 2016), MobileNet-V2 on CIFAR-10 dataset, and ResNer-18 (He et al., 2016), ResNet-50 (He et al., 2016), DeiT-small on ImageNet- 1K dataset.
1 kmm
if Of > Top(O%, [B.k]) and ||lll Top( W " kJ),
BC
otherwise,
(10) K. For the implemen-
where We tation of BAME, we follow (Jayakumar et al., 2020) to perform a three-step sparse training pipline, with T, and
Tr evenly divides the training schedule. In particular, we first employ gradual pruning (Zhu & Gupta, 2017) to set the
Implementation Details. We train N:M sparse networks
6

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
Table 2: Results for sparsifying ResNet-50 and DeiT-small on ImageNet. et al., 2023b) for backward propagation of BAME.
means we apply the bi-directional masks (Zhang
Model
Method
N:M
Top-1 Accuracy (%)
Epochs (Train)
FLOPs (Train)
FLOPs
Pattern
(Test)
ResNet-50
Baseline
77.1
120
(x(2.2e18)
I (x(822e9)
ResNet-50
ASP
2:4
76.8
200
1.24x
0.51
ResNet-50
SR-STE
2:4
77.0
120
0.83x
0.51
ResNet-50
LBC
2:4
77.2
120
0.72x
0.51
ResNet-32
Bi-Mask
2:4
77.4
120
0,66x
0,51x
ResNet-32
MaxQ
2:4
77.6
120
0.91x
0.51x
ResNet-50
BAME(ours)
2:4
77.4
120
0.59
0,51x
ResNet-50
BAME(ours)
2:4
77.4
120
0.63x
0.51x
ResNet-50
SR-STE
1:4
75.3
120
0.74x
0.26x
ResNet-50
Bi-Mask
1:4
75.6
120
0.49x
0.26x
ResNet-50
BAME(ours)
1:4
76.1
120
0.39
0.26x
ResNet-50
SR-STE
1:16
71.5
120
0.69
0.11
ResNet-50
Bi-Mask
1:16
71.5
120
0.37x
0.11x
ResNet-50
BAME(ours)
1:16
72.0
120
0.29
0.11x
DeiT-small
Baseline
79.8
300
(8.9e118)
Ix(9.2e9)
DeiT-small
SR-STE
2:4
79.6
300
0.83x
0.11x
DeiT-small
Bi-Mask
2:4
79.4
300
0.72x
0.11x
DeiT-small
BAME(ours)
2:4
79.7
300
0.63x
0.11x
from scratch via the Stochastic Gradient Descent (SGD) optimizer, paired with momentum of 0.9 and a batch size of 256. The initial learning rate set to 0. and gradually decayed based on the cosine annealing scheduler. Following previous works, train all networks for 300 epochs on CIFAR-10, 10, with weight decay of 0.005. On ImageNet, 120 epochs given for ResNet and 300 epochs for DeiT- small. For the implementation of BAME, we set the LMA update interval AT 100 and 0.5 for both a and 3 in OBF ALL experiments are implemented based on PyTorch and executed on NVIDIA Tesla A100 GPUs.
narios, even utilizing far fewer training FLOPs and param- eters compared with other methods. For instance, BAME achieves 94.71% top-1 accuracy when training 1:4 sparse ResNet-32, surpassing the recent baseline Bi-Mask that also pursues efficient backward propagation by 0.28%. More- over, compared with SR-STE which conducts dense gradient calculation training, BAME still achieves better performance retention for all N:M patterns even with sparse backward propagation. For example, when training 1:16 sparse MobileNet- V2, BAME yields 93.32 top-1 accuracy, surpassing SR-STE by 0.18% while only using 0.29% train- ing FLOPs (0.67% for SR-STE).
We juxtapose BAME N:M sparsity methods, includ- ing ASP (Nvidia, 2020), SR-STE (Zhou et al., 2021), LBC (Zhang et al., 2022), Bi-Mask (Zhang et al., 2023b). We experiment with a wide range N:M patterns for com- parison, including 2:4, 1:4, and 1:16. We report the Top-1 ac- curacy, the training/inference and parameter burden of N:M sparse networks.
Performance Metries and Baselines.
with several
ImageNet. For the large-scale ImageNet-IK 1K dataset that contains over 1.2 million images for training and 50,000 im- ages for validation in 1,000 categories, we first present the quantitative results for training sparse ResNet with depths of 18 and 50, along with DeiT-small in Table 2. Again, BAME substantially enlarges the performance of existing methods, with the minimum training FLOPs by efficient weight prun- ing and growing. For instance, it surpasses SR-STE by 0.5% Top-1 accuracy when training 1:4 sparse ResNet-50 (76. 16 for BAME and 75.3% for SR-STE), while consumes far fewer training FLOPs (0.39x for BAME and 0.74x for SR-STE). Although BAME's performance slightly infe- rior to MaxQ by 0.2 Top-1 accuracy, it achieves a more than 1.54-fold reduction training overhead. Moreover, when introducing sparsity constraints into backpropagation, BAME; still maintains notable performance advantages
point operations (FLOPs)
4.2. Image Classification
CIFAR-10. We first evaluate the efficacy of BAME for training sparse ResNet-32 and MobileNet- V2 on the CIFAR- 10 dataset, which includes 50,000 training images and 10,000 validation images within 10 classes. Tab I show- cases the performance comparison under different N:M pat- terns. BAME achieves state-of-the-art accuracy at all sce-

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
Table 3: Ablation studies on the training schedule.
93.2
Top-1 Accuracy (%)
t
t
Top-1 Accuracy (%)
FLOPs (Train)
92.9
0
100
91.87
0.07x
92.6
200
92.57
0.07x
300
92.22
0.07x
92.3
100
300
92.81
0.29x
92.0
200
300
93.05
0.41x
a=0.5, B=0.5
0=0.5, B=
100
200
93.15
0.29x
a-0.1,
8-0.5
a' 0-1.1,
0-1
91.4
10
102
10a
Table 4: Ablation study on LMA and OBF.
Update Interval (AT)
Figure 3: Ablation studies the update schedule.
Method
Top-1 Accuracy (%)
Stastic
90.03
compared with Bi-Mask, while readily leveraging N:M sparse tensor cores during training. It is also worth mention- ing that BAME holds its advantages when training sparse DeiT-small compared with other methods, demonstrating its scalability for other types of model structures beyond convolution neural networks.
RigL
92.01
LMA
92.98
LMA+OBF
93.15
contrast, prolonging BAME until the termination of train- ing, that is, designating ty to the final iteration, results in even more precipitous performance degradation. To ex- plain, the freshly grown weights, initialized to zeros as per Alg. 1, mandate substantial training following restoration to enhance the performance of the sparse network.
4.3. Performance Analysis
In this section, we provide the performance analysis of BAME, with all experiments conducted on training 1:16 ResNet-32 on CIFAR-10.
Mask Adaption. At last, we investigate the effectiveness of our mask adaption methods including LMA and OBF, We set static training the baseline, which means the bi- nary masks are randomly initialized and kept frozen during the entire sparse training procedure. In addition, we run RigL (Evci et al., 2020), a representative method for tradi- tion network sparse training that take weight magnitude and gradient for pruning and reviving, respectively. As shown in Tab. 4, both LMA and OBF contributes the overall or sparse training performance.
Hlyper-parameters We first investigate the influence of hyper-parameters within BAME, including the two restric- tion factors a and 8, and the updating interval AT. As shown in F, the best performance is obtained with AT = 100, a = 0.5, 3 = 0.5. To analyze, smaller a and 3. larger e all lead to an insufficient procedure for mask ex- ploration during the training schedule. Setting these hyper- parameters in the contrast direction, also resulted in poor performance, which fits into our claim that high frequency of mask oscillations can unavoidably harm the training sta- bility and lead to sub-optimal results. Nevertheless, it serves as a promising directions to automatically perform N:M sparse training without tyrrr-parrmettr choosen.
5. Conclusion
N:M sparsity has become an increasingly crucial DNN com- pression tool, delivering functional speed ups by imposing a maximum of N non-zero constituents within M consecutive weights. We introduce BAME, method that enhances the efficiency of the contemporary N:M sparsity methods while preserving the model's performance. BAME's s fundamental principle involves carrying out loss-aware mask adaptation to prune and revitalize weights within specific N:M blocks, whilst maintaining the stability of frequently-oscillatinng blocks. BAME surpasses existing methods in sparsifying mainstream networks across various vision tasks, all while greatly reducing the training FLOPs and the parameter strain by keeping both sparse forward and backward propagation through training. Hopefully, BAME will not only provide practitioners with a robust N:M sparse training instrument, but also set the groundwork for further investigations into efficient N:M sparsity.
Training Schedule. Further, we analyze the training sched- ule of BAME, Le, and f for stooping the gradual pruning and performing mask adaption. Tab. 3 delineates the quanti- tative results. Intuitively, establishing larger tę indicates an increase in training iterations for pre-training with grad- ual attainment of the desired sparsity level. Though this consequently induces a significant training cost, neglecting gradual pruning simultaneously results a considerable performance reduction. To explain, the randomly-inttiaizen weights require certain degree of pre-training to initiate an effective importance selection, which is validated in tra- ditional sparsity work (Liu et al., 2021; Jayakumar et al., 2020). Regarding the mask adaptation schedule, prema- turely halting BAME leads to a performance downturn due to inadequate identification of the optimal mask. stark
8

---

BAME: Block-Aware Mask Evolution for Efficient N:M Sparse Training
Acknowledgements
mantic segmentation. In IEEE International Conference Computer Vision (ICCV), pp. 580-587, 2014.
This work was supported by the National Science Fund for Distinguished Young Scholars NNo.62025603, the Na- tional Natural Science Foundation of China (No. 624B2119,
Han, S., Pool, J., Tran, J., and Dally, W. Learning both weights and connections for efficient neural network. In Advances in Neural Information Processing Systems (NeurIPS), pp. 1135-1143, 2015.
No. U21B2037, No. U22B2051, No. U23A20383, No. 62176222, No. 62176223, No. 62176226, No. 62072386,
No. 62072387, No. 62072389, No. 62002305 and No. 62272401), and the Natural Science Foundation of Fujian
He. K., Zhang, X., Ren, S., and Sun. J. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pp. 770-778,200 2016.
Province of China (No. 2021J06003, No.2022J06001..
Impact Statement
He, K., Gkioxari, G., Dollar, P., and Girshick, R. Mask r-cmn. In IEEE International Conference on Computer Vision 2017a.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
2017a.
He, Y., Zhang, X., and Sun, J. Channel pruning for acceler- ating very deep neural networks. In IEEE International Conference on Computer Vision (ICCV), pp. 1389-1397, 2017b.
References
Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for con- ditional computation. arXiv preprint arXiv: 2013.
Hoefler, T., Alistarh, D., Ben-Nun, T., Dryden, N., and Peste, A. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 12:1-124, 2021.
Chmiel, B., Hubara, I., Banner, R., and Soudry, D. Opti- mal fine-grained N:M sparsity for activations and neural gradients. In International Conference on Learning Rep- resentations (ICLR), 2021.
Howard, A. G., Zhu, M., Chen, B., Kalenichenko, D., Wang, W., Weyand, T., Andreetto, M., and Adam, Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: 1xXix1100.048800 2017.
Deng, J., Dong, W., Socher, R., L..JJ. Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 248-255, 2009.
Hubara, I., Courbariaux, M., Soudry, D., Yaniv, R., and
Bengio, Y. Binarized neural networks. Advances in Neural Information Processing Systems (NeurIPS), 29, 2016.
Dettmers, T. and Zettlemoyer, L. Sparse networks from scratch: Faster training without losing performance. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Jayakumar, S., Pascanu, R., Rae, J., Osindero, S., and Elsen, E. Top-kast: Top-k always sparse training. In Advances in Neural Information Processing Systems (NeurIPS), pp. 20744 20754, 2020.
Ding, X., Zhou, X., Guo, Y., Han, J., Liu, J., et al. Global sparse momentum sgd for pruning very deep neural net- works. In Advances in Neural Information Processing Systems (NeurIPS), pp. 6382-6394, 2019.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Kusupati, A., Ramanujan, V., Somani, R., Wortsman, M., Jain, P., Kakade, S., and Farhadi, A. Soft threshold weight reparameterization for learnable sparsity. In Inter- national Conference Machine Learning (ICML), pp. 5544-5555, 2020.
Evci, U. Gale, T., Menick, J., Castro, P. S., and Elsen, E. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning (ICML), pp. 2943-2222222 2943-2952, 2020.
Fang, C., Zhou, A., and Wang, Z. An algorithm-hardware co-optimized framework for accelerating N:M sparse transformers. IEEE Transactions Very Large Scale Integration (VLSI) Systems, 30(1 1): 157331115888 2022.
LeCun, Y., Denker, J., and Solla, S. Optimal brain damage. In Advances in Neural Information Processing Systems (NeurIPS), pp. 598 605, 1989.
Lee, N., Ajanthan, T., and Torr, P. Snip: Single-shot network pruning based on connection sensitivity. In International Conference on Learning Representations (ICLR), 2019.
Girshick, R., Donahue, J.. Darrell, T., and Malik, J. Rich feature hierarchies for accurate object detection and se-

---

BAME: Blork-Aware Mask Evolution for Efficient N:M Sparse Training
Lin, M., Ji, R., Wang. Y., Zhang. Y., Zhang, B., Tian, Y., and Shao, L., Hrank: Filter pruning using high-rank feature map. In IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 1529-15388 2020.
Sun, W., Zhou, A., Stuijk, S., Wijnhoven, R., Nelson, A. O., Corporaal, H., et al. Dominosearch: Find layer-wise fine- grained N:M sparse schemes from dense neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2021. 2021.
Liu, S., Chen, T., Chen, X., Atashgahi, Z., Yin, L., Kou, H., Shen, L., Pechenizkiy, M., Wang, Z., and Mocanu, D. C. Sparse training via boosting pruning plasticity with neuroregeneration. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Wang, Z. Sparsert: Accelerating unstructured sparsity on gpus for deep learning inference. In Proceedings of the ACM International Conference on Parallel Architectures and Compilation Techniques (ICPACT), pp. 31-42.
Xiang, J., Li, S., Chen, J., Chen, Z., Huang, T., Peng, L., and Liu, Y. Maxq: Multi-axis query for n. M sparsity network. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 15845- 15854, 2024.
Liu, Z., Mu, H., Zhang, X., Guo. Z., Yang, X., Cheng, T. K.- T., and Sun, Metapruning: Meta learning for automatic neural network channel pruning. In IEEE International Conference on Computer Vision (ICCV), pp. 3296-3305, 2019.
Zhang, Y., Lin, M., Lin, Z.,, Luo, Li, K., Chao, F., Wu, Y., and Ji, R. Learning best combination for efficient N:M sparsity. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Louizos, C., Welling, M., and Kingma, D. P. Learning sparse neural networks through L. regularization. In International Conference on Learning Representations (ICLR), 2017.
Zhang, Y., Lin, M., Zhong, Y., Chao, F., and Ji, R. Lottery jackpots exist in pre-trained models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2023a.
Luo, J. and Wu, J. Autopruner: An end-to-end trainable filter pruning method for efficient deep model inference. Puttern Recognition (PR), pp. 107461, 2020.
Zhang, Y., Luo, Y., Lin, M., Zhong, Y., Xie, J.. Chao, F., and Ji, R. Bi-directional masks for efficient M sparse train- ing. In International Conference Machine Leurning, pp. 41488-4 4497.. PMLR, 2023b.
Luo, J.-H., Wu, J., and Lin, W. Thinet: A filter level pruning method for deep neural network compression. In JEEE International Conference on Computer Vision (ICCV), pp. 5058-5066, 2017.
Zhou, A., Ma. Y., Zhu, J., Liu, J., Zhang, Z., Yuan, K., Sun, W., and Li, H. Learning N:M fine-grained structured sparse neural networks from scratch. In International Conference on Learning Representations (ICLR), 2021.
Molchanov, P., Tyree, S., Karras, T., Aila, T., and Kautz, Pruning convolutional neural networks for resource effi- cient inference. International Conference on Learning Representations (ICLR), 2017.
Zhu, M. and Gupta, S. To prune, or to prune: explor- ing the efficacy of pruning for model compression. In International Conference on Learning Representations Workshop (ICLRW), 2017.
Mostafa, H. and Wang. X. Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization. In International Conference on Ma- chine Learning (ICML), pp. 4646-4655, 2019.
Nvidia. Nvidia a100 tensor core gpu architec-
ture. https: www.wwwwwwww..wwwwwwwww..www..wp www///oo//ottoocooooooooooocccee
dam/en-zz/So http:///uuutiiss/aaaaaaaa-aa
nvidia ampere ampere wwchhirecture-witteppee pdf,2020
Pool, J. and Yu, C. Channel permutations for N:M sparsity.
In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Ronny Krashinsky, O. G. c. Nvidia ampere sparse tensor avidia. com /blg//
core.
nvidia-annd wridiia-apeeee-raee wwwppeeeeaaccitttttreeeeeeeeeeee /epth/, 2020.
10


Reviews:
[{"review_id": "rwAdFT7wPq", "reviewer": "Reviewer_2Nc7", "summary": "The paper introduces a novel approach called BAME (Block-Aware Mask Evolution) for training N:M sparse networks in an efficient manner. The authors argue that prior works often rely on dense gradient updates, which leads to considerable overhead. Instead, BAME keeps the network consistently N:M sparse throughout both the forward and backward passes. The core idea comprises two components: Loss-Aware Mask Adaption (LMA): and Oscillation-aware Block Freezing (OBF). Experimental results on CIFAR-10 and ImageNet show that BAME achieves accuracy comparable to or surpassing prior N:M sparsity methods, while consuming significantly fewer training FLOPs", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your positive and motivating comments. Please kindly see our response to your concerns below.\n\n**Q1**:  Actual training or inference speed measurement of BAME.\n\n**A1**:  We appreciate this constructive feedback. We first clarify that **reporting N:M sparsity patterns and theoretical FLOPs (rather than empirical acceleration ratios) is a standard practice in the N:M sparsity community** [1–3]. This is because N:M sparsity is inherently a hardware-software co-design problem: algorithmic works in this domain typically focus on ensuring that weight matrices satisfy N:M sparsity constraints during forward/backward propagation, while retaining the performance. Critically, once N:M sparsity is achieved in either phase (forward or backward), the acceleration effect becomes deterministic given certain hardware support[4,5]. For instance, [5] reports a 2.17× speedup for 2:8 sparse training when both forward and backward passes are sparse (equivalent to our BAME, T-mask[2], Bi-mask[3]), while methods with dense backward passes (e.g., SR-STE[1]) achieve only 1.33× speedup. Currently, the only open-source hardware framework is NVIDIA Ampere Sparse Tensor Core[4] that support 2:4 inference. Below are our test inference latency results on an NVIDIA A100 GPU (batch size=512):\n\n| Method | N:M | Latency per Batch (ms) | Top-1 Acc.|\n| - | - | -| - |\n| ResNet-50 |- | 59.88|77.4|\n| BAME | 2:4 | 37.19|77.4|\n\nThis demonstrates practical inference acceleration (1.61×) for BAME. Regarding training latency, existing N:M training frameworks are neither open-source nor easily deployable without specialized hardware expertise. Due to limited resources, **we prioritized algorithmic innovation (aligning with the community norms) to report only theortical FLOPs**. While FLOPs remain a standard metric for cross-method comparison, we will add the result of inference latency and clarify this point in our final version to avoid any misunderstanding.\n\n**Q2**: Theoretical proofs beyond first-order Taylor approximations for analyzing loss impact.\n\n**A2**: We fully agree that more rigorous theoretical derivations exist—such as second-order Hessian-based analysis. However, we employ first-order approximation due to its efficiency, as the derived metric w*g_w is readily available during training. This ensures training efficiency, and first-order approximations have also been widely validated as effective in reflecting loss impact[8], though exploring beyond them remains a promising future direction. Thanks very much for this insight.\n\n**Q3**: Hyperparameters might require domain knowledge.\n\n**A3**: We acknowledge that BAME does involve several hyperparameters, yet most follow established sparse training recipes[8] and require minimal tuning. Despite this, we recognize the potential of automatically performing N:M training without choosing hyperparameters, and we earmark this aspect for future work.\n\n**Q4**:  Clarify the initialization and final freeze phases in even more detail.\n\n**A4**: We re-initialize the pruned weights to their historical values rather than set them to zero. This is directly motivated by Equation 6 of our manuscript, which evaluates the impact of restoring a pruned weight on the loss by leveraging its historical state. For the storage of historical states, we directly store and freeze the values of pruned weight. We ensure that the above clarification will be added to our final version.\n\n**Q5**: The relationship between frequent updates and block importance.  \n\n**A5**: Thank you for sharing this insightful comment. Indeed, mask oscillation has been a persistent research point in the community. Zhou et al. [1] showed mask fluctuations correlate with loss reduction, where excessive variations yield suboptimal results. Their SR-STE solution increases weight decay on pruned weights to stabilize masks. The point you raised presents an intriguing research direction, as both our OBF and SR-STE seem to be heuristic to avoid mask oscillation. Perhaps these sensitive blocks may require specialized evaluation metrics to assess weight significance that avoids fluctuation without compromising performance. Addressing this challenge demands substantial algorithmic innovation, which we reserve for future work.\n\nWe sincerely appreciate the time and effort you have dedicated to reviewing our paper. Should you have any further inquiries, please let us know and we would be more than delighted to engage in further discussion with you.\n\n[1] Learning N: M fine-grained structured sparse neural networks from scratch. In ICLR, 2021\n\n[2] A Provable and Efficient Method to Find N:M Transposable Masks. In NeurIPs, 2021\n\n[3] Bi-directional Masks for Efficient N:M Sparse Training. In ICML, 2022.\n\n[4] Nvidia a100 tensor core gpu architecture, 2020\n\n[5] Efficient N:M Sparse DNN Training Using Algorithm, Architecture, and Dataflow Co-Design. In IEEE TCAD, 2023. \n\n[6] Rigging the lottery: Making all tickets winners. In ICML, 2020.", "claims_and_evidence": "The paper’s main claim is that BAME’s consistent N:M sparsity (applied throughout forward and backward passes) can significantly reduce training overhead without sacrificing accuracy. The authors present empirical results on CIFAR-10 and ImageNet using models such as ResNet, MobileNet, and DeiT, demonstrating that BAME achieves on-par or better accuracy compared to prior N:M sparsity methods at substantially lower training FLOPs.\n\nOne limitation, however, is that these claims rely on theoretical FLOPs counts. The authors do not provide actual training or inference speed measurements, such as latency, throughput, or wall-clock time, which would give more direct evidence of real-world performance gains.", "methods_and_evaluation_criteria": "The proposed method applies periodic “pruning-and-regrowth” within each N:M block, guided by LMA and OBF that freezes blocks prone to mask instability.\n\n- Top-1 Accuracy \n\n- Training/Test FLOPs \n\n- NM spasrity\n\nWhile these metrics are typical and demonstrate clear benefits (especially in theoretical FLOPs), the paper omits direct measurements of actual training/inference speed (e.g., wall-clock time, throughput, latency on real hardware). Including such real-world measurements would strengthen the argument that the reduced FLOPs indeed translate to practical speedups and resource savings.", "theoretical_claims": "The submission does not present extensive theoretical proofs beyond first-order Taylor approximations for analyzing loss impact.", "experimental_designs_or_analyses": "- The experiments are well-structured: multiple N:M patterns (2:4, 1:4, 1:16) are tested on CIFAR-10 and ImageNet with widely used architectures. The ablation studies around hyperparameters (α,β,ΔT) and scheduling (when to start and stop mask adaption) support the validity of the method.\n\n- The results consistently show superior or on-par accuracy with lower training FLOPs. The authors also compare with relevant baselines (SR-STE, ASP, LBC, Bi-Mask, MaxQ).\n\n- The experiments seem sound overall, but additional real-world speedup measurements (wall-clock time) on hardware supporting N:M sparsity (e.g., A100 GPU) might further demonstrate the actual training speed benefits.", "relation_to_broader_scientific_literature": "- BAME aligns with the growing literature on sparse training (RigL, SET, SNIP) and N:M sparsity (ASP, SR-STE, LBC, Bi-Mask).\n\n- The core novelty is combining a loss-aware local pruning/regrowth scheme with an oscillation detection approach specifically tailored for N:M blocks.", "essential_references_not_discussed": "I did not identify any critical missing references that would drastically affect the paper’s context.", "other_strengths_and_weaknesses": "Strengths:\n\n- Presents a well-motivated approach for block-level mask evolution.\n\n- Simple idea but demonstrates strong empirical results across multiple networks/datasets.\n\n- Reduces training FLOPs while retaining or improving accuracy compared to prior methods.\n\nWeaknesses:\n\n\n- More real-world training speed tests (beyond theoretical FLOPs) would be helpful to confirm the practical efficiency on hardware that supports N:M acceleration.\n\n- The choice and tuning of certain hyperparameters (α,β,ΔT) might require domain knowledge, although the paper provides some guidance through ablation.", "other_comments_or_suggestions": "I would recommend clarifying the initialization and final freeze phases in even more detail: how exactly are the pruned weights re-initialized? Are they set to zero or restored from historical values, etc.?", "questions_for_authors": "- For newly “revived” weights, are they always re-initialized to zero or do they retain their old (pre-pruned) value [1]? If always zero-initialized, does that hamper potential recovery for large updates? If not, how do you track historical states?\n\n-  Have you measured actual training speed on an NVIDIA A100 or similar GPU to confirm that the 0.29× or 0.39× training FLOPs factor translates into a similar wall-clock reduction? If so, please provide more context.\n\n- You mention that OBF (Oscillation-aware Block Freezing) identifies blocks with high-frequency mask changes and freezes them to improve training stability. However, one might wonder if these frequent updates are actually a signal that those blocks are “sensitive” or “important.” In other words, rather than straightforwardly freezing such blocks, could there be a more nuanced approach—one that better leverages this apparent sensitivity while still mitigating the risk of oscillation?\n\n\n[1] CHEX: CHannel EXploration for CNN Model Compression. CVPR 2022", "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": "9afdI1f2F4", "reviewer": "Reviewer_gJ4C", "summary": "This paper presents a novel approach for preserving sparsity in DNNs during training, with a focus on N:M sparsity. The authors introduce BAME (Block-Aware Mask Evolution), a technique that ensures both forward and backward propagation remain sparse while iteratively pruning and regrowing weights within predefined blocks. Unlike traditional methods that rely on dense backward propagation—often computationally costly—BAME offers a more efficient alternative. Experimental results on CIFAR and ImageNet show that BAME achieves performance comparable to or better than state-of-the-art methods while significantly reducing training FLOPs.\n\n## update after rebuttal\n\nMy concerns are addressed. I will keep my rating.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your positive and motivating comments. Please kindly see our response to your comment below.\n\n**Q1**: While BAME demonstrates lower training overhead compared to MaxQ, its performance appears to be slightly inferior to MaxQ in certain cases.\n\n**A1**:  We appreciate this insightful comment. We acknowledge that our method may exhibit marginally inferior performance to MaxQ in some scenarios. However, we would like to highlight that our approach achieves a significant reduction in training FLOPs—from 0.91× to 0.59× for training 2:4 sparse ResNet-50 compared to dense training. Thus, while maintaining comparable performance to MaxQ, our method offers substantial advantages in training efficiency.\n\n**Q2**: Another minor limitation of the proposed method is its reliance on hyperparameters. \n\n**A2**:  We appreciate this constructive feedback and acknowledge that BAME does involve several hyperparameters. However, we would like to clarify that most of these hyperparameters follow established conventions from prior sparse training methods [1,2]. As such, they are considered well-established defaults in the field and typically do not require extensive tuning—for instance, the update interval ΔT is fixed at 100, following standard practice. Despite this, we recognize the potential effectiveness of automatically performing N:M sparse training without choosing hyperparameters, and we earmark this aspect for future work.\n\n[1] Rigging the lottery: Making all tickets winners. In ICML, 2020\n\n[2] Sparse Training via Boosting Pruning Plasticity with Neuroregeneration. In NeurIPs, 2021.\n\nWe sincerely appreciate the time and diligence you’ve taken to participate in the review of our paper. If you have further questions, we are more than glad to discuss with you.", "claims_and_evidence": "The claims are well-supported by both theoretical analysis and empirical results. The authors provide a detailed theoretical proof for the efficacy of LMA, along with extensive experiment to demonstrate the performance of BAME compared with other methods.", "methods_and_evaluation_criteria": "Yes, the proposed BAME method is well-suited for the problem of N:M sparse training in DNNs. The use of block-level mask evolution is innovative and addresses the efficiency limitations of traditional methods. The evaluation criteria, including benchmark datasets like CIFAR and ImageNet, are appropriate and widely accepted in the field, ensuring the results are meaningful and comparable to prior work.", "theoretical_claims": "Yes, the theoretical claims regarding the LMA method appear to be correct. The proofs are well-structured and logically sound. No issues were identified in the theoretical analysis.", "experimental_designs_or_analyses": "The experimental designs are sound and well-executed. The authors conduct extensive experiments across multiple datasets (CIFAR, ImageNet) and network architectures to validate the effectiveness of BAME.", "relation_to_broader_scientific_literature": "The key contributions of this paper build on prior work in N:M sparse training. Specifically, the authors address the efficiency limitations of existing methods that rely on dense backward propagation, which is computationally expensive. By introducing block-level mask evolution, BAME offers a more efficient alternative that aligns with recent trends in reducing the computational cost of training DNNs. The paper effectively situates itself within the broader literature by comparing BAME to state-of-the-art methods and demonstrating its advantages.", "essential_references_not_discussed": "No, the paper adequately covers the relevant literature and provides sufficient comparisons with other methods.", "other_strengths_and_weaknesses": "*Strengths:**\n\n1. The paper is well-structured and clearly written, making it easy to follow.  \n2. **Innovative Methodology**: The paper introduces a fresh perspective on N:M sparse training by focusing on block-level sparsity evolution. This approach effectively balances efficiency and performance, showcasing significant potential for reducing the computational costs associated with dense training.  \n3. **Theoretical Rigor**: The theoretical proofs are thorough and provide a solid foundation for understanding the design principles behind the LMA metric. Additionally, the visualization of mask oscillations effectively validates the design rationale of OBF.  \n4. **Empirical Evidence**: The experimental results are extensive and demonstrate the efficacy of BAME across various network architectures and datasets. The results show that BAME can achieve substantial reductions in training FLOPs without compromising model accuracy.  \n\n**Weaknesses:**\n\n1. While BAME demonstrates lower training overhead compared to MaxQ, its performance appears to be slightly inferior to MaxQ in certain cases.  \n2. Another minor limitation of the proposed method is its reliance on hyperparameters. Future work could explore ways to automate N:M sparse training to reduce the need for such specific hyperparameter tuning.", "other_comments_or_suggestions": "No", "questions_for_authors": "I keep up with the literature in this area.", "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": "RUcmUY80AN", "reviewer": "Reviewer_hLEc", "summary": "This paper introduces a novel method, BAME, which maintains consistent sparsity throughout the N:M sparse training process. Specifically, BAME ensures sparse forward and backward propagation while iteratively performing Loss-aware Mask Adaptation (LMA) and Oscillation-aware Block Freezing (OBF) to adapt the mask. This approach not only optimizes the N:M mask but also ensures training efficiency. The authors have conducted extensive experiments to validate the effectiveness of BAME.\n\n## update after rebuttal\n\nI think my concerns are addressed and I am happy to keep my initial positive rating.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your careful review and constructive comments. Please kindly see our responses to your questions below.\n\n**Q1**: The proposed metric in Loss-aware Mask Adaptation is somewhat common in traditional sparse methods, which limits the novelty of the paper to some extent. \n\n**A1**: Thanks for this very insightful comment and we would like to share some explanations here.  While we acknowledge that pruning-and-reviving is indeed a well-established technique in sparse training literature, BAME introduces two key innovations that differentiate it from prior work: (1) block-wise loss-aware mask adaptation (LMA) for N:M sparsity, and (2) oscillation-aware block freezing (OBF) to stabilize frequently-oscillating N:M blocks. Together, these novel components enable BAME to significantly reduce training overhead while achieving state-of-the-art performance in N:M sparse network training.  We appreciate your comments and hope the above discussion clears up any misconceptions regarding our contribution in introducing BAME.\n\n**Q2**: The analysis of SAD is too coarse-grained in the method part.\n\n**A2**:  We sincerely appreciate this insightful suggestion. Following your valuable comment, we conducted additional analysis by tracking the average SAD variation per layer throughout training, with the results for sparsifying ResNet-50 at 1:16 pattern listed below (due to length limit, we random select some layers to show).\n\n| Layer | 1    | 7    | 18   | 22   | 39   | 42   | 45   | 46   |\n| ----- | ---- | ---- | ---- | ---- | ---- | ---- | ---- | ---- |\n| SAD   | 23.1 | 22.9 | 18.5 | 15.3 | 12.5 | 5.7  | 3.8  | 3.9  |\n\nThe above results reveal that deeper layers indeed exhibit significantly higher SAD values compared to shallower ones. Inspired by your suggestion, we further implemented a preliminary improvement: we linearly scaled the OBF parameter from 0 to 0.5 across network layers (shallow to deep). This modification demonstrates promising results on ImageNet with ResNet-50 under 1:16 sparsity, as shown in the following table:\n\n| Method        | N:M Pattern | Top-1 Accuarcy |\n| ------------- | ----------- | -------------- |\n| OBF           | 1:16        | 72.0        |\n| Layerwise OBF | 1:16        | 72.3    |\n\nWhile such a heuristic modification shows potential, a more adaptive and elegant parameter allocation scheme could yield further performance gains, which we leave as a promising future work. We sincerely appreciate your expert suggestion in highlighting this valuable research avenue.\n\n**Q3**: It would be beneficial to perform additional experiments on tasks such as object detection to validate the effectiveness of BAME in other domains.\n\n**A3**: Following your professional suggestion, we further exploit the generalization ability of BAME on the object detection and instance segmentation tasks of the COCO benchmark[2]. The results are delineated as follows:\n\n- Results on object detection tasks with Faster-RCNN[3] as the backbone.\n\n| Model  | Method   | N:M  | mAP  |\n| ------ | -------- | ---- | ---- |\n| F-RCNN | Baseline | -    | 37.4 |\n| F-RCNN | SR-STE   | 2:4 | 38.2 |\n| F-RCNN | BAME     | 2:4 | 38.5 |\n\n- Results on instance segmentation tasks with Mask-RCNN[4] as the backbone.\n\n| Model  | Method   | N:M  | Box mAP | Mask mAP |\n| ------ | -------- | ---- | ------- | -------- |\n| M-RCNN | Baseline | -    | 38.2    | 34.7     |\n| M-RCNN | SR-STE   | 2:4 | 39      | 35.3     |\n| M-RCNN | BAME     | 2:4 | 39.2    | 35.4     |\n\nWe will incorporate these results into our final version. Thanks for the valuable suggestion again.\n\n[1] Microsoft coco: Common objects in context. In ECCV, 2024.\n\n[2] Faster r-cnn: Towards real-time object detection with region proposal networks. In NeurIPs, 2015.\n\n[3] Mask r-cnn. In ICCV, 2017.\n\n Your time and effort in reviewing our paper are genuinely appreciated. If there are any additional questions or points that require clarification, we would be more than delighted to engage in further discussions.", "claims_and_evidence": "Yes, the authors provide clear experiment results to support the theoretical foundations of LMA and OBF in BAME.", "methods_and_evaluation_criteria": "1. The proposed method effectively addresses the current inefficiency of N:M sparsity methods.\n2. The evaluation criteria align well with the image classification benchmarks, including CIFAR-10 and ImageNet.", "theoretical_claims": "I have reviewed all the theoretical analyses and found them to be sound and correct.", "experimental_designs_or_analyses": "The overall experimental design and analyses are well-structured and reasonable.", "relation_to_broader_scientific_literature": "This paper effectively sets the stage for the problem of N:M sparsity and the efficiency limitations of current approaches, leading logically into the proposed solution with BAME.", "essential_references_not_discussed": "From the best of my knowledge, there is no essential reference that needs further discussion.", "other_strengths_and_weaknesses": "Strengths:\n1. The motivation behind BAME is to maintain consistent sparsity during both forward and backward propagation is very attractive. This contrasts sharply with existing methods that rely on dense backward propagation for weight updates, which can be computationally expensive. This is a necessary prerequisite for efficient N:M training on future work.\n2. Extensive empirical evidence supports the effectiveness of the BAME method, particularly in keeping state-of-the-art performance while drastically reducing the training FLOPs.\n3. The mathematical analysis behind Loss-aware Mask Adaptation is solid. The entire paper is well-written and easy to understand.\n\nWeaknesses:\n1. The proposed metric in Loss-aware Mask Adaptation is somewhat common in traditional sparse methods, which limits the novelty of the paper to some extent. Nevertheless, it still make clear contribution to the N:M sparsity literature.\n2. The analysis of SAD is too coarse-grained in the method part. The authors did not explore the mask variations across different layers in detail, which could further optimize the OBF mechanism.", "other_comments_or_suggestions": "The authors only conducted experiments on image classification tasks. It would be beneficial to perform additional experiments on tasks such as object detection to validate the effectiveness of BAME in other domains.", "questions_for_authors": "As previously mentioned, does the variation in masks differ across different layers of the network? If so, could designing different OBF parameters for each layer based on these variations lead to better performance?", "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:

Meta-review: The paper introduces a novel approach called BAME (Block-Aware Mask Evolution) for training N:M sparse networks aiming to preserve the sparsity through both forward and backward propagation. The core BAME novelty is combining a loss-aware local pruning/regrowth scheme with an oscillation detection approach specifically tailored for N:M blocks. All reviews are very positive, recognize the novelty of the findings, and appreciate the presentation of the paper. Therefore, the AC suggests acceptance and encourages the authors to include the rebuttal discussion in the final version of the paper.
Decision: Accept (poster)