S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training
Authors: Yuezhou Hu
Jun Zhu
Jianfei Chen
Venue: NeurIPS 2024 poster (2024)
Keywords: efficient machine learning
transformers
2:4 sparsity
dynamic sparse training
Primary area: optimization_for_deep_networks

Abstract: Training deep neural networks (DNNs) is costly. Fortunately, Nvidia Ampere and Hopper GPUs can accelerate matrix multiplications twice as fast as a dense equivalent by implementing 2:4 sparsity. However, previous STE-based 2:4 pre-training methods (\eg~STE with hard-thresholding, SR-STE) suffer from optimization difficulties because of discontinuous pruning function.
In this study, we comprehensively analyse the bottleneck of traditional N:M sparse training and recognize three drawbacks with discontinuity: incorrect descending direction, inability to predict the amount of descent and sparse mask oscillation. In the light of this statement, we propose S-STE, a simple yet powerful 2:4 training method that contains two parts: to continuously project weights to be 2:4 sparse, and to rescale sparse weights with a per-tensor fixed scaling factor. Besides, we adopt minimum-variance unbiased estimation for activation gradient and FP8 quantization for whole process. Results show that our method surpass previous 2:4 pre-training recipes and is comparable even with full parameter models.

S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training
Yuezhou Hu', Jun Zhu', Jianfei Chen Dept. of Comp. Sci. & Tech., Institute for AI, BNRist Center, Tsinghua-Bose Bosch Joint ML Center, THBI Lab, Tsinghua University.
huyz21@mails
wtw.:mmmmlsssstnittaaaaae tsinghua. edu. en, {dcszj pacszj,jiane stanfehc edu.
Abstract
Training deep neural networks (DNNs) is costly. Fortunately, Nvidia Ampere and Hopper GPUs can accelerate matrix multiplications twice as fast as a dense equivalent by implementing 2:4 sparsity. However, previous STE-based 2:4 pre- training methods (e.g. STE with tard-thresholding SR-STE) suffer from opti- mization difficulties because of discontinuous pruning function. In this study, comprchensively analyse the bottleneck traditional N:M sparse training and recognize three drawbacks with discontinuity: incorrect descending direction, in- ability to predict the amount of descent and sparse mask oscillation. In light this, propose S-STE, simple yet powerful 2:4 training method that contains two parts: to continuously project weights to be 2:4 sparse, and to rescale sparse weights with per-tensor fixed scaling factor. Besides, we adopt minimum-varianee unbiased estimation for activation gradient and FP8 quantization for whole pro- cess. Results show that our method surpasses previous 2:4 pre-training recipes and is comparable even with full parameter models. Our toolkit is available at https:///sthuu..co//yy/////////0 wwt.huyz2023/ pretrain
1 Introduction
Large scale transformers have achieved many impressive results such as chatbots [33 text-to-video generation (27], and robot manipulation [53). However, the pre-training of these models is extremely expensive, typically requiring thousands of GPUs to train for months [5]. One possible way to accelcrate deep learning computation is sparsity. N.M sparsity [I]] is serdwware-friendle sparsity pattern, where every group M dimensions only has N non-zero entries. Nvidia Ampere GPUs can multiply 2:4 sparse matrix with dense matrix, twice fast multiplying two dense matrices.
While N:M sparsity has been successfully applied to accelerate inference the acceleration to pre-training is highly challenging. To accelerate pre-training, the sparse model must be trained from scratch (random initialization), and the network must stay sparse all training iterations. To meet these requirements, the algorithm should be able to actively explore connectivity patterns within the constrained N:M parameter space. Therefore, popular pruning methods such single-shot pruning (24). iterative magnitude pruning [13 29], and RigL [IT] cannot be directly applied to this scenario. Moreover, besides forward propagation, the matrix multiplications in back propagation must be sparsified as well, provide reasonable training speedup.
35 g extending
have shown promise towards solving the challenging problem of sparse pre-training. They maintain dense weight, which is sparsified in each iteration for fast Forward&backwart computation, and the dense weight is then updated with STE gradients. In this way, connectivity patterns can be learned jointly with weights in end-to-end
Methods based on the straight-throup estimator (STE)
"Corresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

fashion with stochastic gradient optimizers. SR-STE [52 is such method to train sparse networks from scratch, with a regularization term stabilize the training. Several subsequent works [T] g
further accelerate back propagation with sparse computations, and Hu et al. [20] applied it for
pre-training language models. However, these sparse training methods still have an accuracy gap compared to dense training. Moreover, SR-STE introduces regularization strength
which is hard to tunc. Due to these limitations, N:M sparsity is not yet used to accelerate pre-training.
In this work, we study STE-based pre-training from the optimization perspective. We point out that STE-based pre-training defines a discontinuous loss function, which existing optimization theory and algorithms cannot handle. We reveal several intriguing phenomena highlighting the difficulty of discontinuous optimization, including incorrect descending direction, inability to predict the amount of descent, and oscillation. We sidestep the curse of discontinuity by proposing smooth sraighhtthrough estimator (S-STE) as solution. Cruically, S-STE introduces new pruning function, which uses a continuous projection function to prune weights to be 2:4 sparse, and scales all nonzero elements to minimize the mean-square-ens vector. The proposed 2:4 soft-thresholding function is continuous but can still generate N:M sparse weights at all times. In this way, the objective function is continuous, and gradient-basee optimizers can be readily used. Furthermore, S-STE does not introduce any parame,, so its practical adoption is easier than SR-STE.
between original dense weight vector and sparse weight
We devise comprehensive pre- vre-trainin experiments on S-STE, including WMT machine translation, GPT-2 pre-training and, DeiT image classification. Results show that our method surpass previous 2:4 pre-training recipes on a wide range of tasks.
2 Formulation of sparse pre-training
The training neural network can be formalized as an optimization problem minw F(w), where Rd is the parameter and F is differentiable empirical risk function: F(w) Ra(w)
Σ f(w: Here, f is the loss function, n is the size of data set D = {Egg)i and Eji is the i-th 1=1
sample. The optimization can be solved with standard stochastic gradient method (SG) 4. Suppose
the network is initialized with w1, {au } is positive learning rate sequence, and Etx,) is randomly
chosen from } 1 Then, iteratively we have W = W. 0g w (W E(x.)) As we consider pre-training tasks, Wi simply a random initialization.
The training of a sparse network involves optimizing the parameter w a constrained space W C Rd For an N---spe network, the parameter can only have N non-zero elements in each contiguous M dimensions.
Alternative to constrained optimization, we can solve the unconstrained problem:
min F(w) where w S(w).
Here, 5 is pruning function which converts a dense weight w to sparse weight E W. One common choice is the hard-thresholding pruning function [52 45] For every block of four adjacent E RRM in the weight vector w, the pruning function can be defined as
elements a
(a)( = [ai if la. a. <t''
for i M,
(2)
where t is N-th largest element in a This essentially performs magnitude- based pruning, by zeroing out the two smallest elements. The hard thresholding function can also be written as S, (a) = C mp,a), where (a) is a 0/1 mask vector, with (m)( = if [a:| >
However, Eq. cannot be directly optimized since the pruning function S is not differentiable. Particularly, the derivative of the tard-liresholding function Sn. is undefined boundary where the second largest and third largest element have the same magnitude. Therefore, straight-through estimator (STE) [52 is for training, by approximating w V and therefore aS, (a)/aa /da I:
W++1 = We a4 w. f( E(r)))
In this paper, when talking about large and small, we refer to the magnitude. For example, "second largest element" means the element with second largest absolute value.

---

With the pruning function and STE, each iteration of sparse training involves: (1) prune the dense weight to get the sparse weight: and (3) update the dense weight with the gradient. Among these, step 2 is most time-consuming, and can be accelerated with sparse tensor cores given v is N-M-sparse. Next, we will focus on the optimization aspects of sparse training and defer the discussion of computation details to Sec.
S(w); (2) compute the loss and gradient with the sparse weight;
3 The curse of discontinuity
Classical stochastic optimization theory guarantees the convergence for nonconvex and smooth (i.e,, differentiable with Lipschitz continuous gradients) objective F. It can be also extended to handle non-differentiable functions such as ReLU [25]. The real problem of STE-based sparse training the discontinuity of the pruning function Sh. as visualized in Fig. [ For a discontinuous function, an arbitrarily small change in input a can cause an unbounded change of the output Sp (a). Such discontinuity appears on the boundary when the N-th and N + 1-th largest elements have same magnitude. For example, for 1:2-sparse pruning function, S₁. (1, 0.999) = (1, O), but Sp, (0.999, 1) (0, 1), and the boundary is the line d1 a2-
When Sa is discontinuous, the parameter space Rd can be partitioned into regions {MmlIm E M}. where M C {0, 1]d is the space of 0/1 masks with N:M pattern, and all the pa- rameters in each region same mask my (w) = m. The loss landscape F(S, (w)) Fmm,, (w) @ w) is continuous and differentiable within each region, where gradient-based algorithms can work well. How- ever, when the optimization trajectory crosses the boundary: my mmx.wa (W&+1) A my (ws), the be- havior is unpredictable. We highlight several intriguing phenomena observed in optimizing such discontinuous objective. We study these phenomena in both a toy problem and real neural networks.
Wm have the
005
4.10
0.0
005
6:3
0.50
300
Figure 1: Scatter plot of AF; with AF2 and their distributions GPT-2 small 124M for iteration E [1, 6000).
3.1 Phenomenon 1: incorrect descending direction
Here, we run gradient descent algorithm (without stochasticity) on small dataset. For dense model where F is differentiable, with Taylor's formula we should have
F(WL) F(W++1) ~ wVw, F(w.)) T (W - W6+11 au V F( Fl 2 0.
(4)
That is, the objective function will monotonically decrease in each iteration once the learning rate ar is sufficiently small. However, it is not the case for sparse training. In Fig. dd),, we measure the distribution of the amount of descent (AoD) AF == F(Wx)) F(Wx+1) for training GPT-2 large 774M model with Eq. 2 3). across each iteration k. The results clearly shows that the objective frequently fails to descent.
We can take a closer look to the weight and mask sequence (Wk: my (Wx)) generated by the training algorithm. We compare the following two quantities: the AoD by updating both weight and mask A F,- = F(w; @ m2) F(W @ m and the AoD by only updating the weight A F2 = F(w O m;) F( O me). In Fig. we can observe A F2 is mostly positive due to the piece- wise continuity of However, A F1 is frequently negative and very often even smaller than A F2 (updating mask is worse than not updating). This indicates that the main problem is the discontinuity make it hard to estimate the correct descending direction m.
3.2 Phenomenon 2: inability to predict the amount of descent
Besides making mistakes in finding the correct descending direction, algorithms do not know that they make a mistake, the sense that they fail to predict the AoD at each step. From Eq. ( we should have F(W+) F(W+41) se (Vw, F(wx)) (Wk wk+1), where the left hand side is the actual
3

---

AoD, and the right hand side is the predicted AoD, We plot the actual AoD against predicted AoD for dense (Fig. a)) and sparse training (Fig. 2b)). While for dense training, the two quantities closely matches, for hard-thresholding the actual AoD is often lower for the predicted AoD, particularly when the predicted AoD is large. To understand this, note that Eq. (4) only holds for E Wma (w)- Once WK+1 W is large enough that ma (W&+1) A m4 (W&), the function crosses a border of Sh. and F will have sudden change which is unpredictable by the gradient.
(d
62
04
Inskdrrhoking with wef-ffededdive
341
081
E0
032
Figure 2: (a)-(c) shows scatter plots of the predicted and actual loss reduction of dense, hard- thresholding and S-STE with GPT-2 large 774M model for iteration k: E [1, [0,0, The diagonal line is for reference. (d) shows empirical cumulative distribution of their actual AoD for & [1, 6000].
3.3 Phenomenon 3: oscillation
Oscillation probably the most significant problem in STE-based sparse training. Here, we revisit
existing discussions about oscillation [52 28 20), and then illustrate this issue using toy example.
Flip rate Flip rate is simple metric to measure the stability of sparse training 20 [Imm, (W*) @ my (Wa-1)llı /d, where @ indicates XOR operation. As Hu et al. 2 points out, taking the flip rate of the dense model as standard, they observe larger flip rate of hhrd-thresholding: when training transformers, the flip rate can stay at 6% in the entire training process. However, a healthy training process should have large flip rate in the early stage to explore connectivity patterns, and the flip rate should decrease to zero in later stage for the optimization to converge. Hu et al. [20] describe this phenomenon as "flip rate explosion", which is harmful to sparse training.
An exemplar toy problem Modern deep neuron networks have billions of parameters and is not strictly convex. These non-ideal conditions make our analysis more difficult with sparse weights. To analyze the characteristics of STE and toy problem that contains two parameters: min w2 ( w²2 = (wi w2)² )² This may differ from the real DNN optimization problem, but can help us understand what happens in the process. We are going to show that while using a feasible og that can make the dense model converge to global minima, STE with theshhhhesn fails to converge and it oscillates back and forth.
on the smallest problem, we devise simple
First, for the dense model, the global minima lies on the line wi W2- Suppose we start from W1 = [0.2, 0.1]" by taking ag = 0.25 we can reach global minima in one step. On the other hand, if we are in 1:2 sparse situation, the global minima should be the point w11 = Ww2 By starting from w1 = [0.2, 0.1]T and taking ag = 0.25, we invariably jumps between w = [0.2, 0.1] and W2: = 0.2 and never decreases.
High flip rate is harmful, because there are frequent changes on the connection of neurons, which means that a number of previous optimization steps on the neuron is deprecated. That is fatal at the end of training (20). The reason of high flip rate discontinuity: as there are no gentle transitions on both sides of the border, the gradient on the boundary is inaccurate and is unable to indicate the right descending direction. This misalignment is easy to make the tuple a to oscillate back and forth near the boundary, and cause extremely higher flip rate than the dense model.
hard-thresholding can be explained by
3.4 Overcoming the curse of discontinuity
One way to mitigate discontinuity is sparse-refined straight- estimator (SR-STE), which adds sparse-refined regularization on the gradients [52]: minw F(w)= + Aw
2 ||w 3 m(w) 1. While

---

SR-STE works on a wide range of tasks and optimization algorithms [52 20,, it still has some issues. First, the performance is quite sensitive to the term leads to competition between loss function and sparse regularization. Finally, the loss function still discontinuous unless Aw + 0o.
An. Second, the new regularization
From the above analysis, discontinuity causes optimization problems. It would be ideal to have continuous pruning function, yet the iterates (wx) still need to be sparse during the entire training process.
61
60
640
480
Figure 3: Pruning function of hard-thresholding and soft-thresholding for 1:2-sparsity. (a)(b) show the outputs of hard- threeshhohddin and (c) show that of soft coftthreshollding A sudden jump exists in threshlissg if [a1| |az|. while thte-hodrshhnnggggg continuous the domain.
4 Methodology
In this section we propose a training algorithm (smooth straight-through estimator, S-STH) that contains two main parts combined with STE: 2:4 specific soft- thresholding, and fixed weight rescaling. They together work as the sparsifying function described in Sec. 2 f0 (w). Results in Fig. (d) and dd show that S-STE successfully overcome the three curses of discontinuity. Notably, flip rate curves of S-STE are surprisingly consistent with their dense counterparts, indicating that S-STE is more natural and feasible than SR-STE.
Table I: Validation loss and test accuracy of S-STE with different y on Transformer- base
w S(w)
Y
Val loss
Test BLEU
0
4.007
26.30
0.33
4.014
26.01
0.67
4.015
26.16
4.072
25.63
4.1 2:4 specific soft-thresholding Ssaft
Motivation for the design As discussed in Sec.
hard-thresholding suffer from the discontinuous problem near the boundary of taking a flip. When input vector changes continuously across the border, two of the four elements simultaneously jump between zeroes and none-zero values. In a continuous pruning function, we need to overcome this drawback and keep these two elements zero on both sides of the border. This means when a flip happens in a four-element block, at lease three of the target elements should be zeroed out simultaneously (except the largest one).
With the above analysis, we modify soft-thresholdng function for traditional pruning in Van- derschueren and Vleeschouwer 45
2:4 specific soft-thresholding. Given vector [a1, 42, a3 aa] E R+, S-STE picks out the largest two elements and meanwhile, subtracts the third largest element from weight magnitudes. Assume, without loss of generality, that [t,, t2, ts. tx] is an rearrangement of aT, s..1/44 < |62| < 1/al $ 141/- Then, the pruning function can be defined as
a; if a, (00)
(Seof((a)); = 0 af E (-t,t) where t = =|t/|-
(5)
+t if E (-00, -t]
The plots of soft-thresholding is drawn in Fig. 3 showing Stoft is continuous everywhere. Note that although we define Stoft by a block E R Rt,SSoof can be extended to arbitrary a E R2.4 for t 2 1. by doing block-wise pruning.
Theorem 4.1. (a) is a continuous projection for a E Rdd
A detailed discussion of the proof can be found in Appendix A.I
5

---

Choosing optimal threshold Theoretically, any real number in [[tz]. [ts] can be used as a feasible threshold. This gives infinite options and we describe it with an interpolation r))(ka with E [0, 1J. The larger is, the closer L is to |/z], and the smaller f is. This may affect model's capacity. In order maximize the retention of information, using a small y
7|21| + (1
necessary. In our method we propose to set 7 = 0. Experimental results in Table Talso show that the network has the best accuracy when = O, Le.. t = |t2//
4.2 Fixed weight rescaling |
Because Stoft reduce the total magnitude of w, is not a close simulation of dense weights. Like Vanderschueren and Vleeschouwer [45], we scale up Seoft(W in our method as S(w) SS but we modify weight rescaling in their study to adapt to our approach. First, we use a per-tensor scale 3 rather than per-channel /3 for simplicity. Besides, two important improvements are made: to compute scale factor only at the beginning of training, rather than to dynamically update scale factor during training, and to minimize the mean-squae mean-squure-erro (MSE) between original dense weights and sparse weights, rather than to keep the total magnitude of weights unchanged.
Freezing scaling factor As Vanderschueren and Vleeschouwer 45 use dynamic 3 for every iteration, we argue that this doesn't It align with our approach. We explain our solutions in two parts.
Suil
St.si
SSIE
33
103
Figure 4: (a) Flip rate curve over the training process with different 3 Transformer base. (b) Dynamically recalculated 3 at each layer on different epochs. Results show that frequently updating will cause it to be unexpectedly large. (c) Flip rate over the training process with fixed and dynamic 3 on Transformer-base. (d) Flip rate of dense, SR-STE and S-STE algorithm on Transformer-base.
First, we find it interesting that /3 have a subtle correlation with flip rate: in sparse model, larger B usually results in higher flip rate. The reason can be explained by the accuracy of gradients. As we use STE in the backward pass, the approximation VwJ & V is valid when w and are close enough. However, if scale too large then optimal, w and are too far apart to guarantee this. Such mismatch leads to incorrectness in descending directions, and thus unstability in optimization increase flip rate; see Fig. (a..
Second, we argue that dynamically computing scaling factor for each iteration leads to high flip rate in our training process. Fig. 4b) shows the results dynamically changing 3 will make it increase with iterations, especially for later layers. Fig. (c) shows flip rate of this network, which has significantly higher tail than the dense one. Considering high flip rate harmful, we propose to compute scaling factor / only in the first iteration. After that, we use the same 3 in the rest of the training. Fig. d) shows the flip rate of our fixed scaling S-STE, which perfectly aligns with the dense one.
Minimizing MSE Vanderschueren and Vleeschouwer 45] choose to scale up Ssoft (w) to have the same Ll-norm as w: 8 = || I (Seoff(w)))) However, we choose minimize the MSE gap of Seofe (w) and w. As [8] point out, sparsifying weights in the forward pass should minimize MSE rather than unbiased estimation. In method, to determine optimal scale B, we need to minimize
MSE
|| (SSoof(()))
-|wl|| 2 T Ssoft +|S::of((w]]] 32.
(6)
Rearrange the terms and taking partial derivative of B, we choose 3 w T Swofe(W) (w)/ Sanft (w))(((w)))
(w)]
The comparison between no scaling, our minimizing MSE and keeping I. -norm can be found in Table 2
Result show that our method yields the best results in practice.

---

Table 2: Experimental result of different on Transformer-bast base..
3 Recipe
Test BLEU
Val loss
Avg epoch loss
No scaling
25.28
4.044
4.670
Keeping L1-norm same [45
25.85
4.019
4.627
Minimizing MSE (S-STE)
26.3
4.007
4.605
Table 3: Results of different MVUE strategies on GPT-2 774M with 4000 steps. Sparsifying S S(W
introduces huge loss of accuracy while sparsifying Vz is acceptable with little loss.
S-STE MVUE(S(W))) MVUE(VA))
comment
loss
X AX A
dense
3.3948
x
SR-STE
3,4739
X
3.4333
3.4644
3.4773
X
3.4480
Other implementation skills
5.1 Minimum-variance unbiased estimation
To accelerate the backward propagation, Chmiel et al. 8 suggest using minimum-varte variancc unbiased estimator (MVUE). For every linear layer Z₁ = -xs) there are two matrix multiplications of the backward pass in total: Vx. = Vz, (Ww) and Vw, = V, 2, X1, where X₁ is the input of the (-th layer, W; and Z are the weight matrix and output activation. We conduct MVUE on both two matrix multiplications and compare their results: Vx, = Vz, MVUE(S(W..)
and Vw, MVUE( V Xx, Specifically, we choose S(W,) and Vz, because they both have built-in sparsity [26] However, we only choose to sparsify the latter one. Firstly, it is proven by Hu et al. 20 and Chmiel et S that minimum loss of accuracy is guaranteed for MVUE on Vz, Secondly, using MVUE S (WW) will make errors accumulate along the back propagation, and results in large standard deviation of gradient for the first few layers. Besides, results Table Jalso show minimum loss of accuracy of MVUE on Vz, while obvious accuracy loss on S(W,). Thus, choose to sparsify only Vz, in the backward pass.
5.2 FP8 training
To further accelerate pre-training of networks, we utilize popular FP8 workflow in training. Similar Transformer Engine P FP8 e3m4 in forward pass and e5m2 in backward pass. Besides, use per-tensor rescaling before casting to FP8 formats.
Theoretical acceleration of S-STE While 2:4 sparsity can accelerate GEMMs up to 2x faster, FP8 quantization accelerate an additional 2x this basis. Thus, the three GEMMs in Sec. 5.1 can be 4x, 2x, 4x faster. To sum up, we have theoretically 3x faster in forward and backward pass.
6 Experiments
We validate the feasibility of our proposed method S-STE on machine translation (Transformer [46 image classification (DeiT (42) and generative large language models (GPT-2 (38 series). For all models, we replace the two linear layers in the feed forward network of each transformer block with

---

S-STE. We keep the rest of the networks, the optimization algorithms as well as all hyperparameters the same as their dense counterparts.
Transformer, we train Transformer- models on WMT 14 En-De dataset 3 with fairseq codebase and evaluate with BLEU [36] scores. For DeiT, pre-train Deit-small model for ImageNet-IK [0] classification task. For GPT-2, we pre-train GPT-2 124M, 350M and 774M models on OpenWebText [[6] and evaluate GLUE 47 and SQuAD 39 benchmarks. We also compare
For
34
our method with state-of-the-art 2:4 training methods (SR-STE [52] Bi-Mask [31] and STEP (28). com/thu ttuutthuuaa mml
The pee- pre-training and evaluation scripts are publicly available https::/ 2by4- pry--aretrain- acc- examples
Machine translation We first apply S-STE to train a 12-layer Transformer-base and compare with SR-STE and STEP. Note that we use fairseq codebase with SacreBleu metric, whose baseline should be 26.5 (the result of our reproduction is 26.42). The results are shown in Table Compared with SR-STE, method improves by 0.3 and 0.5 on test set and validation set respectively, which the closest to baseline. Besides, we improve by 0.6 compared to STEP on test set.
Table 4: Experimental Results for Transformer-base En-De dataset.
Method
Avg epoch loss
Test BLEU
Val BLEU
Val loss
Dense
4.555
26.42
26.49
3,977
SR-STE
4.61
25.84
26.08
4.023
STEP 4.682
25.52
26.01
4.085
S-STE
4.617
26.11
26.53
4.011
Table 5; Experimental Results for DeiT-small on ImageNet-Ik, The Bi-Mask and SR-STE results are
from (51]
Model
Method
Test acel
Test acc5
Dense
72.2
91.1
DeiT-tiny
SR-STE 51
67.8
88.6
S-STE
68.5
88.9
Dense
79.9
95
SR-STE 51
75.7
DeiT-small
Bi-Mask 51
77.6
S-STE
78.5
94.4
Image classification We further investigate the effectiveness of S-STE train DeiT-tiny and DeiT-small on ImageNet-lk; see Table Results show S-STE also achieve the best performance among different methods, with only has 1.4% degradation from the dense model. Notably, S-STE surpasses SOTA 2:4 training method Bi-Mask on this task (0.9% topl accuracy improvement) and popular SR-STE method (2.8% topl accuracy improvement).
Generative language models We compare S-STE with dense, normal SR-STE and SR-STE with dense fine-tuning (00 (SR-STE+DF) models on GLUE and SQuAD tasks. The 3R-STE+DF models first use SR-STE to train 2:4 sparse model, and switch to dense training for the last 1/6 iters of pre-training (which stands for "dense fine-tune").. In downstream tasks it also use dense parameters to make predictions, similar dense models. Results in Table 6 and show that S-STE completely surpasses SR-STE on both tasks. Even for SR-STE+DF models, S-STE still have an advantage, with an improvement of 1.5 on GLUE average and 1.2/0.9 on SQuAD for GPT-2 774M.
We illustrate the viability of S-STE for fine-tuning pre-trained model, presenting coherent workflow of accelerating both training and inference (dense fine-tuning cannot produce sparse model for inference acceleration); see Table
Fine-tuning

---

Table SQuAD and GLUE scores of different sizes and pre-training methods on GPT-2. We 2:4 sparse weights to evaluate S-STE model, while dense parameters to evaluate the rest. Of note, SR-STE denotes the original SR-STE workflow (without backward MVUE), and "T-SR-STE+DF" denotes the combination of transposable SR-STE & backward MVUE & sparse-dense training workflow, proposed by Hu et 20 S-STE settings here include backward MVUE FP8 training.
Pre-train val loss
SQuAD
Params
Pre-train
Fine-tune
GLUE@ Avg
EM
F1
Dense
Dense
2.907
67.6
78.8
73.9 + 1.1
ESR-STE+DF 20
Dense
2.952
67.5
78.5
74.3
124M
T-SR-STE
Dense
3.076
66.3
77.2
72.6 0.2
SR-STE
Dense
2.982
66.2
77.5
S-STE
S-STE
2.984
68
78.8
74.1 0.4
Dense
Dense
2.618
73.2
83.6
76.3 0.1
T-SR-STE+DF 20
Dense
2.688
71.9
82.4
77.1
350M
T-SR-STE
Dense
2.718
72.3
82.6
76.3
SR-STE
Dense
2.690
72.0
82.4
76.8
S-STE
S-STE
2.713
72.2
82.7
76.9 0.6
Dense
Dense
2.493
74.3
84.9
76.2 0.4
774M
T-SR-STE+DF (0
Dense
2.564
74.3
84.6
77.1 + 0.4
S-STE
S-STE
2.547
75.5
85.5
78.6 0.8
Table 7: Different fine-tuning results on GLUE and SQuAD.
Model
Downstream task
Pre-train
Fine-tune
Avg score
GLUE GLUE
S-STE
Hard Barr--irrshooddng
73.9 + 0.6
S-STE
S-STE
74.1 + 0.4
GPT-2 124M
SQuAD
S-STE
Hard-thresholding
67.6/78.6 68/78.8
SQuAD
S-STE
S-STE
Ablation study In this part, We explore the effectiveness of S-STE, MVUE, and FP8 separately. We pre-train DeiT-small model on ImageNet- IK dataset for image classification. Combinations of these partitions in Table & show that: 1) FP8 training has little affect on pre-training accuracy (0.2% of acc1); 2) MVUE leads to minimal loss of performance (0.1% % of accl).
Table 8: Experimental result of S-STE (soft-trresoolding and weight rescaling), MVUE and FP8 training with Deil-small ImageNet-IK.
Soft- Weight rescaling
MVUE VZ)
FP8
Comment
Test acc1
Test acc5
thresholding
x
Dense
79.9
95
X
Dense; FP8
79.7
94.9
X
x
tard-thresholding
77.7
93.9
X
78.8
94.6
78.9
94.7
78.6
94.4
78.9
94.6
78.2
94.2
78.5
94.4

---

Acceleration For acceleration. we measure the acceleration ratio of a typical GPT-2 model using implementation from Hu et [20]. Note that on H100 GPUs, FP8 2:4 kernel turns out to be unsatisfying; see Appendix A.4 Consequently, we fall back to use RTX3090 GPUs with FP16 training. For inference, we achieve 1.53x speedup with FFN layer and 1.23x speedup with
the network; for pre-training, we achieve 1.32x speedup for FFN layer and 1.18x speedup for the network (Appendix(A.3)
7 Related work
Unstructured pruning and coarse-grained structured pruning Pruning is to remove redundant
weights from the dense model. Traditional one-shot pruning methods [7 T2
14 31
24] and
dynamic sparse training methods [11 6
7 [50] mostly target on unstructured sparsity. While most of them have acceleration effect on CPUs, is hard for these methods work well on modern GPUs.
Coarse-grained structured sparsity 49 23 19 22 takes effect to acceleration, but since they often remove whole chennel or a block, loss of information non-negligible.
Fine-grained N:M sparsity for inference and pre-training Among all pruning techniques for pre-training. N:M sparsity promising approach towards accelerating large models, which is also known as fin-grained structured sparsity. Nvidia demonstrates 2x theoretical speedup on its
Ampere GPUs with 2:4 sparsity for post-training [31] and inference (40 15 35 9.. To leverage this
property to accelerate pre-training, a number of approaches and their improvements are proposed [52 28 51 8 21 20). However, all these methods are based on a discontinuous pruning function that is hard to optimize and results in unsatisfactory accuracy, which we will discuss in this study,
FP8 quantization While 16-bit float tensors are widely used pre-training, FP8 where float numbers stored in 8 bits- - is popular quantization methods which theoretically accelerates GEMMs up to 4x faster than its fp32 counterparts and 2x faster than its FP16/BF16 counterparts [30 48 32). With e3m4 data format used in forward and e5m2 format [4] in backward, pre-trrined models can achieve minimum loss of accuracy while greatly boosting the efficiency of training.
8 Conclusions and future work
In this study we discuss the importance of pruning continuity in effective 2:4 sparse pre-training. We analyse the drawback of traditional hard-thresholding pruning function and its variation (SR-STE) and argue that the main limits being discontinuity. Based on our analysis and soft-thresholding for channel pruning. we propose S-STE, which prunes weights continuous manner. Experiments show that our method surpasses previous state-of-the the-art methods on wide range of tasks.
Our proposed S-STE approach primarily targets linear layers within FFN networks. Nevertheless, QKV projection layers necessitate further exploration to devise an effective dynamic sparse training strategy that harmonizes with attention mechanisms. Furthermore, pruning function represents only one possible solution; alternative, smoother pruning functions may be necessary to achieve improved continuity and mitigate potential discontinuities.
current choice of continuous
Acknowledgments and Disclosure of Funding
We would like to thank Ziteng Wang, Bingrui Li, Haocheng Xi, Kang Zhao and Jintao Zhang, Brian Chmiel and Daniel Soudry for valuable discussions. This work was supported by the Na- tional Key Research and Development Program of China (No. 20212D01 10502) and NSFC Project (Nos. 62376131). J.Z is also supported by the XPlorer Prize.
References
[1] Abhimanyu Rajeshkumar Bambhaniya, Amir Yazdanbakhsh, Suvinay Subramanian, Sheng- Bheng-Chun Kao, Shivani Agrawal, Utku Evci, and Tushar Krishna. Progressive gradient flow for robust n:m sparsity training in transformers, 2024.
[2] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating propagating gradients through stochastic neurons for conditional computation, 2013.
10

---

[3] Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling. Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ales Tamehyna. Findings the 2014 workshop on statistical machine translation. In WMT@ACL. 2014. URL attps://ap1.... ceemanticscooollo
org/CorpusID: 15535376
[4] Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning, 2018.
[5] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
[6] Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijja Liu, Yang Zhang, Zhangyang Wang. and Michael Carbin. The lottery ticket hypothesis for pre-trained bert networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems. volume 33, pages 19834-15846 Curran Associates, Inc., 2020. URL fttppssts: http::::///ccoennggss neerripppp cc/paper_ffile www.ffeeeappppeee wwtpaperr/22///11
Paper saper.pdf
[7] Xiaohan Chen, Yu Cheng, Shuohang Wang, Zhe Gan, Zhangyang Wang, and Jingjing Liu, Earlybert: Efficient bert training via early-bird lottery tickets, 2021.
[8] Brian Chmiel, Itay Hubara, Ron Banner, and Daniel Soudry. Minimum variance unbiased n:m sparsity for the neural gradients. In The Eleventh International Conference on Learning Representations, 2023, URL https://openrevi... ntt/torumm idd -ddDmmmdxddddDDDxxxxxx
[9] Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung, Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu. Plug and play language models: A simple approach controlled text generation, 2020.
[10] Jia Deng. Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248 255, 2009. doi: 10.1000/CVR..200955222
[11] Utku Evei, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Frich Elsen. Rigging the lottery: Making all tickets winners. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Muchine Learning, volume 119 of Proceedings of Muchine Leurning Research, pages 2943-2952. PMLR, 13-18 Jul 2020.
.mlr. ptwss/vv19/ hate////evcc00e www.eeeccaaaaaa..oooe html
[12] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations, 2019. URL |tt httpre:// net/forum woterrumidd--ee J1 b3RcF7
[13] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connec- tivity and the lottery ticket hypothesis. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3259 3269. PMLR, 13-18 Jul 2020. URLLattps:/pproceediun mlr. press/ v119/ Trankle20a. .tml]
[14] Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Stahilizing the lottery ticket hypothesis, 2020.
[15] Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned one-shot. 2023.
[16] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. OpenWebTextCorpus 2019.
http://8Skylion007..itaaa githoa.. 1o/
[17] Song Han, Jeff Pool, John Tran, and William J. Dally. Learning both weights and connections for efficient neural networks, 2015.
[18] Song Han, Huizi Mao, and William J. Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding, 2016.
[19] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. In The IEEE International Conference on Computer Vision (ICCV), Oct 2017.
[20] Yuezhou Hu, Kang Zhao, Weiyu Huang. Jianfei Chen, and Jun Zhu. Accelerating transformer pre-training with 2:4 sparsity. Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver. Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research. pages 19543. PMLR, 27 Jul 2024. URL
mlr press/ wwwess/v33.////ee html

---

[21] Itay Hubara, Brian Chmiel, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. Ac- celerated sparse neural training: A provable and efficient method to find n:m transposable masks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Ad- vances in Neural Information Processing Systems, volume 34, pages 21099-21111 21111. Curran Asso- ciates, Inc., 2021. URL b0490685e92664cbb5bd cb688 Paper.pdi
neur wttpccccpppprr..pppppppp Sicaaper__ilese/eeeeeee htwpl///ppppee//2/////a/
[22] François Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush. Block pruning for faster transformers.
In MariieFrancine Moens. Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih. editors, Proceedings of the 202/ Conference on Empirical Methods in Natural Language Processing, pages 10619-10629, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational I Linguisties. doi: main.829. URL.ttts::///// httpapppee/acaaannoloogggg/
10. 18653/v1/2221
2021 1000021.emmmlppmmmm emnlp- main. 829
[23] Mike Lasby, Anna Golubeva, Utku Evci, Mihai Nica, and Yani Ioannou. Dynamic sparse training with
structured sparsity, 2023.
[24] Namhoon Lee. Thalaiyasingam Ajanthan. and Philip Torr. SNIP: SINGLE NETWORK PRUNING
BASED ON CONNECTION SENSITIVITY. In International Conference on Learning Representations. 2019. URL httpps://openrevies wwwpaarriwww.eetfforreeeeeeeeeee htt.//orrunndd1i2jjaa.ao
[25] Yuanzhi Li and Yang Yuan. Convergence analysis of two-layer neural networks with relu activa- tion. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and Neural Information Processing Systems, volume 30. Curran Asso-
R. Garnett, cditors, Advances http://///cceig.....m neur wwweripp.cccppeeee ciates, Inc., 2017. URL httpss: htttps://proceet
19666ba721e561eee3de (tw.aa/deee/aaa111ff1bb-----
pdf
[26] Yujun Lin, Song Han, Huizi Mao, Yu Wang, and William Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training, 2020.
[27] Yixin Liu, Kai Zhang, Yuan Li, Zhiling Yan, Chujie Gao, Ruoxi Chen, Zhengqing Yuan, Yue Huang, Hanchi Sun, Jianfeng Gao, Lifang He, and Lichao Sun. Sora: review on background, technology, limitations, and opportunities of large vision models, 2024.
[28] 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 Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 22812-22824. PMLR, 23-29 Jul 2023. URL https: https://prroceedings mlr. prwss/ees//ee.. wttppppwsss/eeppppppppppp h/2///22222222e html
[29] Jaron Maene, Mingxiao Li, and Maric-Francine Moens. Towards understanding iterative magnitude pruning: Why lottery tickets win. 2021.
[30] Paulius Micikevicius, Dusan Stosic, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenthwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, Naveen Mellempudi, Stuart Oberman, Mohammad Shoeybi, Michael Siu, and Hao Wu. Fp8 formats for deep learning, 2022.
[31] Asit Mishra, Jorge Alhericia Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. Accelerating sparse deep neural networks, 2021.
[32] Badreddine Noune, Philip Jones, Daniel Justus, Dominic Masters, and Carlo Luschi. 8-bit numerical formats for deep neural networks, 2022.
[33] OpenAL, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin. Suchir Balaji, Valerie Balcom, Paul Baltescu, Haiming Bao, Mohammad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shairro. Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brockman, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai,, Rosie Camphell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael, Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen. Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux. Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gibson, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Gontijo-Lops Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hallacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse. Alan Hickey, Wade Hickey, Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang. Roger Jiang. Haozhun Jin, Denny Jin, Shino Jomoto, Billie Joun, Heewoo Jun, Tomer Kaftan, Lukasz Kaiser, Ali Kamali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo. Lukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger,
Jonathan Gordon, Morgan Grafstein, Scott Gray,
12

---

Vishal Kuo, Michael Lampe. Ikai Lan, Teddy Lee. Jan Leike, Jade Leung, Daniel Levy, Chak Ming Li, Rachel Lim. Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning, Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reichiro Nakano, Rajeev Nayak, Arvind Neelakantan, Richard Ngo, Hyeonwoo Noh, Long Ouyang. Cullen OeKeefe, Jakuh Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng. Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny. Michelle Pokrass, Vitchyr H. Pong. Tolly Powell, Alethea Power, Boris Power, Elizabeth Prochl, Raul Puri, Alec Radford. Jack Rae, Aditya Ramesh, Cameron Raymond. Francis Real, Kendra Rimbach, Carl Ross, Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Selsam, Kyla Sheppard, Toki Sherbakov, Jessica Shich, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens, Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky. Yang Song. Natalie Staudacher, Felipe Petroski Such, Natalie Summers, Ilya Sutskever, Jie Tang. Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Preston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vijayvergiya, Chelsea Voss, Carroll Wainwright, Justin Jay Wang. Alvin Wang. Ben Wang. Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng, Matt Wiethoff. Dave Willner. Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Workman, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret 7oph. Gpt-4 technical report, 2024.
[34] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019 Demonstrations, 2019.
[35] Manoj Kumar Pandey, Daniel P. Delorey, Qiuyi Duan, Lei Wang, Charles D. Knutson, Daniel Zappala, and Ryan Woodings. Ria: An rf interference avoidance algorithm for heterogeneous wireless networks. 2007 IEEE Wireless Communications and Networking Conference, pages 4051-4056, 2007. URL https
Vappissmanti wwwaanttccccclllaacaoooo www.:olaaarrrooooouaooouaaea : 10798336
[36] Kishore Papineni. Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Pierre Isabelle, Eugene Charniak, and Dekang Lin, editors, Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, pages 311-318, Philadelphia, Pennsylvania, USA, July 2002. Association for Computational Linguistics. doi: 10.3115/1073083. 1073333 URL Uttttttps::/ puattps:/aaclase
.r.. Pwogg/P02-110 1040
[37] Sergio P. Perez, Yan Zhang, James Briggs, Charlie Blake, Josh Levy-Kramer, Paul Balanca, Carlo Luschi, Stephen Barlow, and Andrew William Fitzgibbon. Training and inference of large language models using 8-bit floating point, 2023.
[38] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models unsupervised multitask learners. 2019.
[39] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Jian Su, Kevin Duh, and Xavier Careras, editors, Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2383-2392, Austin, Texas, November 2016. Association for Computational Linguisties. doi: 10. 18653/ 1264. URL https://aclanthhoogg . org/D16-- 1264
[40] Mingjic Sun, Zhuang Liu, Anna Bair, and Zico Kolter. A simple and effective pruning approach for large language models. arviv preprint arXiv:2306. 16695, 2023.
[41] Xiao Sun, Jungwook Choi, Chia-Yu Chen, Naigang Wang. Swagath Venkataramani, Vijayalakshmi Srinivasan, Xiaodong Cui, Wei Zhang, and K. Gopalakrishnan. Hybrid 8-bit floating point (hfp8) training and inference for deep neural networks. In Neural Information Processing Systems. 2019. URL https //api. sereiissemntiie i cscholar Serg/CorpussID. 202779167
[42] Hugo Touvron, Matthicu Cord, Matthijs Douze Francisco Massa, Alexandre Sablayrollcs, and Hervé Jégou. Training data-efficient image transformers & distillation through attention, 2021.
[43] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint ar0iv: 13971, 2023.
[44] Mart van Baalen. Andrey Kuzmin, Suparna S Nair, Yuwei Ren, Eric Mahurin, Chirag Patel, Sundar Subramanian, Sanghyuk Lee, Markus Nagel, Joseph Soriaga, and Tijmen Blankevoort. Fp8 versus int8 for efficient deep learning inference, 2023.
13

---

[45] Antoine Vanderschueren and Christophe De Vleeschouwer. Are straight-through gradients and soft- thresholding all you need for sparse training?, 2022.
[46] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems. volume 30. Curran Associates. Inc., 2017. URL httpttttts:::/ceeeeeeeee /ttp///eee//ffeee2eeae2eeeeeeela
cesurips.ccppppse.. files/ Paaeeeaaaaaaaaaaaaaaaeaaa .df
[47] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman, GLUE: A multi-task benchmark and analysis platform for natural language undcrstanding. In Tal Linzen, Grzegorz Chrupala, and Afra Alishahi, editors, Proceedings of the 2018 EMNLP Workshop Black- boxNLP: Analyzing and Interpreting Neural Networks for NLP. pages 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10. 10.18653/1/W18854444 URL https://aclaattoology Arg/W18-5444
[48] Naigang Wang, Jungwook Choi, Daniel Brand, Chia- Yu Chen, and Kailash Gopalakrishnan. Training deep neural networks with 8-bit floating point numbers, 2018.
[49]
Yin, Gen Li, Meng Fang. Shen, Tianjin Huang. Zhangyang "Atlas" Wang, Vlado Menkovski. Xiaolong Ma, Mykola Pechenizkiy, and Shiwei Liu. Dynamic sparsity is channel-level sparsity Iearner. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Ad- vances in Neural Information Processing Systems, volume 36, pages 67993-6808 Curran Asso- ciates, Inc., 2023. URL httpp::/ppooceeeiiiee www.e/ddiiinnnnnnnaneee..ee.e peeuripa.cc/papeee pttppperrreeeeppppeeeeeeeeeepeee (taper/2023/f/110 3800441200444333//11114411111111 Coaeeee
pdf
[50] Haoran You, Chaojian Li, Pengfei Xu, Yonggan Fu, Yue Wang, Xiaohan Chen, Richard G. Baraniuk, Zhangyang Wang. and Yingyan Lin. Drawing early-bird tickets: Towards more efficient training of deep networks, 2022.
[S1] Yuxin Zhang, Yiting Luo, Mingbao Lin, Yunshan Zhong, Jingjing Xie, Fei Chao, and Rongrong Ji. Bi- directional masks for efficient N:M sparse training. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 41488- 41497. PMLR, 23-29 Jul 2023. URL htt.hhttpp//// htt.:p:s/pproceedigg mlr, preesevv022 200002//aan22aae html
[52] Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. leamming n:m fine-grained structured sparse neural networks from scratch. In International Conference on Learning Representations, 2021. URL aattps://pperrevier
[53] Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, Quan Vuong, Vincent Vanhoucke, Huong Tran, Radu Soricut, Anikait Singh, Jaspiar Singh, Pierre Sermanet, Pannag R. Sanketi, Grecia Salazar, Michael S. Ryoo, Krista Reymann, Kanishka Rao, Karl Pertsch, Igor Mordatch, Henryk Michalewski, Yao Lu, Sergey Levine, Lisa Lee, Tsang- Wei Edward Lee, Isabel Leal, Yuheng Kuang, Dmitry Kalashnikov, Ryan Julian. Nikhil J. Joshi, Alex Irpan, Brian Ichter. Jasmine Hsu. Alexander Herzog, Karol Hausman, Keerthana Gopalakrishnan, Chuyuan Fu, Pete Florence, Chelsea Finn, Kumar Avinava Dubey, Danny Driess, Tianli Ding, Krzysztof Marcin Choro- manski, Xi Chen, Yevgen Chebotar, Justice Carbajal, Noah Brown, Anthony Brohan, Montserrat Gonzalez Arenas, and Kehang Han. Rt-2; Vision-language-actinn In Jie Tan, Mare Toussaint, and Kourosh Darvish, editors. Proceedings of The 7th Conference on Robot Learning, volume 229 of Proceedings of Machine Learning Research, pages 2165-2183, PMLR, 06-09 Nov 2023. URL hotp:http////// ftttp:://proceednngg mlr, wwwpppeeessssee w229//222222/// /2/99//ztt www.wowwwiiaa...ooooooooooo
models transfer weh knowledge robotic control.
14

---

A Appendix supplemental material
A.1 Proof of Theorem 4.1
Proof. We prove this by demonstrating Ssaft is continuous on every 4-element block. Of note, suppose a = [11, 02, 03, a4l Assume, without loss of generality, that laı I aaz
S |asl I laal- Our goal is prove Ve 0, 30 > O.S.I. when Ja ai <Ō, 122 02) 6, ldg 0al ô and |44 a1l < 6, (S
agl < e, where a' [a4 22g ag.. d'1]"
1) We start from the simplest case where |a1| |02| |asl < |as) Then we have
Stoft = [0, O, sign(as) - aa2)). sign(as)((aal - (az) ]]-
This order holds when
min |ai)) laz1 la2)); |aal [asl))-
Thus,
(a') [0, 0, sign(as) sig)). D0sig |d21)]-
The signs of a is unchanged when
min 1021 laal. laal}.
We have
sign(a') - |a4] sign(as) la22)
$ |a'al |ag] |ag| + |a2l
< 1 |44] |agl + |ab]
|02|
< 26
Take ò S. fe and this is done. is similar to prove that sign (aa) - -(ds) sigd)))sinna)) (22)> S using the same method.
2) We then consider the cases where there are two equivalents a. |1 |ag| lasl < |aa| or
|01l < |azl < |as] = laal. the proof should follow 1) as no flip happens. Thus we only consider the situation where |ai| < |02| = |asl |as]. Under these cireumstances, flip will happen on second and third dimensions of a.
(a) [0, 0, sign(as) |aa])]
Without loss of generality we assume |a|| |4y| S la'l < |o/]- Thus, (a') = [0, sign(())) (1 |agl), (ag [ab
The proof of the fourth dimension is similar to 1), so we only focus 0z- sign(a') (a²])
< |a |لسم|
< |a's - |as|
+ |a41| + |02|
< 26
Take 8 < e and this is done.
3) If there exists three or four equivalent in a, a flip will happen at the second and third dimension. Thus, these cases be reduced to 1) 2).
A.2
GLUE scores of GPT-2
See Table
A.3
Acceleration
See Table [0
15

---

Table Comparison between GLUE scores of different pre-train methods on GPT-2 models. This table is the elaboration of Table
Piece
Niched
LslA
MANN
UPNLI
OUF
2012
SIS-6
MNLI
STCHER
14.00
50.20 60
-0.38:11
1633 2:10
410 2113
333 m.1.4
GR-STE
18:3.8 3.8
T&T L.773 RUR
87118.0.3
805.8-02
035-03 1.0
91S
15/015 -
311 8-1.8
18:--81
54:2--0
16:29-111
41--
1291 --8
12:99-27
1X0M
ESS ST5
185-2-4
08.11 1112
1000--113 11.
4133 85341 1m
S-STII
T2. 0.6
56.0 1.7
64.63.02
60.0 1.3/6661111111
0.9
96.5.8.0.3
96.688.0.1
67.5
B./0053
8.1.9
03.73:0.4
68.6
8.3.4
13.0
182224
0111
16:0:21
710M
721--44
15:1--85 8.0
012 012-00:000000 U.
00.0000
1:715-01
962---11
414-12/003-11
18:13:3.9
Table 10: Pre-training acceleration ratio with different different batch size N, sequence length embedding dimension d and heads number on single FFN block and transformer block of GPT-2 with RTX 3090 GPUs.
N
d
h
FFN
GPT-2
4
2048
5120
40
1.31
1.18
16
2048
7168
56
1.32
1.18
Pre-train
8
2048
7168
56
1.33
1.17
2048
7168
56
1.31
1.17
4
2048
9216
72
1.31
1.18
16
2048
7168
56
1.54
1.23
Inference
8
2048
7168
56
1,46
1.15
A.4 Limitations
As we propose accuracy results of S-STE on several tasks, no actual acceleration result is given. While theoretically 2x faster results can be expected (FP8 quantization), the NVIDIA acceleration library (cuSPARSElt (11) not satisfactory, which causes inconvenience on implementation. The peak FLOPS of 2:4-spMM is lower than theoretical GEMM FLOPS; see Table I1
Table 11: Peak FLOPS of general matrix multiplications (GEMMs) and 2:4 sparse matrix multiplica- tions (2:--spMMs) on H100. The size we take to 16384 16384 16384.
GPU
FP8 Tensor Core
H100 PCIe 2:4-SPMM
3200 TFLOPS
H100 PCIe GEMM
1600 TFLOPS
Specifications
H100 SXM 2:4-SPMM
4000 TFLOPS
H100 SXM GEMM
2000 TFLOPS
Actual results
H100 SXM 2:4-sPMM
1900 TFLOPS
with cuSPARSEIt
H100 SXM GEMM
1500 TFLOPS
Table 12: GPU Hours of pre-training models on RTX 4090.
GPU Hours
GPT-2 124M
400
GPT-2 350M
900
GPT-2 774M
2500
Transformer-basee
30
DeiT-base
120
16

---

A.5
Broader Impact
S-STE can used mainly to accelerate the pre-training stage of large-scale networks, like LLaMA 43 and GPT-4 [3]. Theoretically GEMMs of the FFN layer can be accelerated up to 4x faster than FPI6 dense models, which would greatly reduce the electric power consumption of pre-training modern large-scale models. However, this method may also used for some models that is non- compliance with regulations and ethics, such as models that generate discriminatory contents.
A.6 Experiments compute resources
To replicate our experiments, we provide the estimated GPU hours of each setting; see Table 12
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Both the abstract and beginning part of introduction indicate the scope of the paper's contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in Appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
18

---

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

---

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

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Please refer to Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The authors have reviewed the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Please refer Appendix.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
21

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper doesn't 't release any new models, generators or datasets.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Please refer the "experiments" chapter.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
22

---

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


Reviews:
[{"review_id": "VxUQAxsJZf", "reviewer": "Reviewer_18hX", "summary": "This work studies an efficient method for pre-training, identifying three significant limitations in previous 2:4 sparse pre-training approaches: incorrect descent direction, the inability to predict the extent of descent, and oscillations in the sparse mask. Subsequently, the authors introduce a novel training methodology that integrates a continuous weight projection technique with a rescaling strategy, thereby enhancing pre-training efficiency and achieving performance comparable to conventional full-parameter training.", "questions": "Please refer to the weakness.", "limitations": "There is no end-to-end acceleration results evaluated of current methods.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The structure of this work is well organized. It begins by highlighting the current challenges in 2:4 sparse pre-training, where the suboptimal performance largely stems from the discontinuity of the pruning function. This issue is then explored through various examples, followed by a demonstration of the proposed methods.\n2. The experiments conducted encompass a variety of model sizes and datasets.\n3. The motivation behind this work is well-founded and clearly articulated.", "weaknesses": "- The acceleration of S-STE is demonstrated solely in terms of theoretical gains. Seems like S-STE will introduce extra computation cost, which might sacrify part of the acceleration. Studying the end-to-end acceleration would further enhance the quality of this work.\n- Is the discontinuity issue discussed in Section 3 associated with the learning rate? A larger learning rate results in more substantial weight updates, potentially causing more frequent flipping. Additionally, is this issue still a big problem during fine-tuning scenarios, where the learning rate is generally lower?\n- The improvements of S-STE are modest, e.g., in Table 8, showing a 0.3 improvement in the F1 score for GPT-2 models with 124M and 350M parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 18hX,\n\nThank you for recognizing the potential and effectiveness of our work and for providing detailed constructive comments. Below, we address each point raised.\n\n**Weakness 1:** The acceleration of S-STE is demonstrated solely in terms of theoretical gains. Seems like S-STE will introduce extra computation cost, which might sacrify part of the acceleration. Studying the end-to-end acceleration would further enhance the quality of this work.\n\n**Reply:** It it true that S-STE will introduce extra computation cost, but be argue the impact won't be significant:\n\n1. In a FFN block, the computational complexity of a matrix multiplication with shape $M\\times N\\times K$ is $O(MNK)$. Other operations like activation function, pruning function and weight rescaling are considered **element-wise**, with computation cost bounded to $O(MN)$ or $O(NK)$. Compared to the $O(MNK)$ parts of a FFN block, this cost is negligible.\n2. As micro-batching technique is applied most of the time, the cost of the pruning and weight rescaling operation will be further reduced to $\\frac{1}{m}$, where $m$ denotes the accumulation steps over micro batches.\n\nTo further respond to your doubts about real-world acceleration, we report end-to-end acceleration as follows. For experiment setting, we choose different sizes of GPT-2 models and test acceleration with FP16 weights. **For inference, we achieve 1.53x speedup with FFN layer and 1.23x speedup with the network; for pre-training, we achieve 1.32x speedup for FFN layer and 1.18x speedup for the network.** This will be updated in our paper.\n\n\n\n*Table.* Pre-training acceleration ratio with different different batch size $N$, sequence length $n$, embedding dimension $d$ and heads number $h$ on GPT-2 with RTX 3090 GPUs.\n\n|N|n|d|h|acceleration@FFN|acceleration@GPT-2|\n|-|-|-|-|-|-|\n|4|2048|5120|40|1.309071284|1.176265882|\n|16|2048|7168|56|1.317673412|1.18020863|\n|8|2048|7168|56|1.325846831|1.173059355|\n|4|2048|7168|56|1.308463658|1.171455338|\n|4|2048|9216|72|1.311344165|1.176620318|\n\n*Table.* Inference acceleration ratio with different different batch size $N$, sequence length $n$, embedding dimension $d$ and heads number $h$ on GPT-2 with RTX 3090 GPUs.\n\n|N|n|d|h|acceleration@FFN|acceleration@GPT-2|\n|-|-|-|-|-|-|\n|16|2048|7168|56|1.536392435|1.233632|\n|8|2048|7168|56|1.464448312|1.149633|\n\n\n\nTo further investigate how we reach \\~1.2x speedup, we profile our code and break down the time costs as shown in the table below.\n\n\n\n*Table.* Time costs of each part of our network and the dense model in one iteration per layer. $m$ denotes the accumulation steps over micro batches. Our method is evaluated on GPT-2, with batch size 16, sequence length 1024, embedding dimension 1024 and heads number 16.\n\n|||||dense (ms/exec)|sparse (ms/exec)|acceleration ratio|frequency (exec/iter)|\n|-|-|-|-|-|-|-|-|\n|ffn|linear|fwd|GEMM|12173.8|7305.8|1.67|1|\n|||bwd|GEMM|23295|18688|1.25|1|\n||||mvue+prune|0|171.4|-|1|\n||||total|23295|18859.4|1.63|1|\n|||**total**||**35468.8**|**21558**|**1.24**|1|\n||others [2]|fwd||167|118.2|-|1|\n|||bwd||65.5|20|-|1|\n|||total||232.5|138.2|-|1|\n||total|fwd||12340.8|7424|1.66|1|\n|||bwd||23360.5|18879.4|1.24|1|\n|||total||35701.3|26303.4|1.36|1|\n|others [3]||fwd||6874.3|7090.6|-|1|\n|||bwd||13920.7|14117.5|-|1|\n|||total||20795|21208|-|1|\n|total||fwd||19215.1|14514.5|1.32|1|\n|||bwd||37281.2|32,996.9|1.13|1|\n|||**total**||**56496.3**|**47511.4**|**1.19**|1|\n|prune weight||||0|320.3|-|$\\frac{1}{m}$|\n\n\n\nDue to the reason that 1) previous work [1] achieves similar acceleration ratio on the same settings and 2) we only accelerate two matrix multiplications for pre-training linear layer while previous work [1] accelerate all three multiplications, we believe that the acceleration is reasonable.\n\nBased on the results above, we believe the overheads of continuous weight pruning function is negligible. According to the time cost table above, the cost of continuous weight pruning function per iteration is\n$$\n320.3 \\times \\frac{1}{m} = \\frac{320.3}{m} ms.\n$$\nCompared to other parts ($47511.4ms$ for the whole iteration), this is indeed negligible.\n\nIt is worth noting that the acceleration we achieve is made on RTX 3090 GPUs with FP16 data type. As we try our best to achieve real acceleration effect on H100 GPUs with popular FP8 precision, the acceleration test fail because FP8 2:4-spMM don't even meet dense baseline; see table below. We are in contact with NVIDIA to address this issue and hopefully will get reasonable results in the future.\n\n*Table.* Peak FLOPS of general matrix multiplications (GEMMs) and 2:4 sparse matrix multiplications (2:4-spMMs) on H100.\n\n||GPU|FP8 Tensor Core|\n|-|-|-|\n|Specifications|H100 PCIE 2:4-spMM|3200 TFLOPS|\n||H100 PCIE GEMM|1600 TFLOPS|\n||H100 SXM 2:4-spMM|4000 TFLOPS|\n||H100 SXM GEMM|2000 TFLOPS|\n|Actual results with cuSPARSElt|H100 SXM 2:4-spMM|1900 TFLOPS|\n||H100 SXM GEMM|1500 TFLOPS|", "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": "AFnt7Q4pHK", "reviewer": "Reviewer_JQrs", "summary": "The authors address the challenge of pre-training models, with 2:4 sparsity, to high quality (ideally matching a densely-trained model).  Building on existing methods, the authors point out three issues with discontinuous pruning functions: gradients move in the wrong direction, weight updates do not match expectation, and values can oscillate between being masked and not-masked.  Demonstrating these shortcomings and analyzing the source leads the authors to propose two modifications to the traditional straight-through estimator approach: 2:4-specific soft thresholding and fixed weight rescaling.  The former provides a continuous pruning function, and the latter compensates for the reduced weight magnitude caused by the soft thresholding.  Targeted experiments show that the proposed method does not suffer from the three shortcomings of traditional approaches.  In full training experiments, the authors also show that applying MVUE (an existing technique) to data gradients and using FP8 representations do not interfere with S-STE's success.  Machine translation, image classification, and generative tasks are used for testing various models with 2:4 sparsity applied to the MLPs of transformer blocks and the data indicates that the resulting models are superior to baselines, and in several cases comparable in quality to the dense baselines.", "questions": "- If a lower flip rate is better in Figure 4, is it concerning that SR-STE's flip rate is lower than S-STE's?  (If not, why not?)\n\n- Would S-STE apply to all N:M configurations by subtracting the N+1th largest entry in soft-thresholding, or is it truly 2:4- (and, given Figure 3, 1:2-) specific?", "limitations": "The discussion of the limitations in the appendix is appreciated, but it seems to suggest that sparsity should give a theoretical 4x increase in throughput, but I believe this should be 2x.  (FP8 can theoretically give another 2x.)  Also, the sizes of the GEMMs that gave these results should be listed.  Finally, I'd point out that the method has only been tested on two out of four linear layers in Transformer networks, not other types of networks, including covolutional networks.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "**Originality**\nWhile the individual advancements, soft-thresholding and weight rescaling, are rooted in prior work, I think the combination and the modifications made by the authors are highly original.  \n\n**Quality**\nThe experiments performed are precisely what are needed, and no more.  Each claim is supported with one clear figure or table so that the reader immediately sees that the authors considered the various angles and have shown why the angle they have selected is the best.  The breadth of full training experiments shows that the proposed method applies to image classification, machine translation, and text generation with particular success in the latter two.  Further, they combined their method with two largely orthogonal compression techniques: FP8 data representations for training and sparsifying another tensor (data gradients).  This makes the results even more compelling.\n\n**Clarity**\nI found the submissions to be easy to read in general and with an appropriate amount of detailing previous and related work.  The few issues I had were minor and listed below, and they did not detract from the well-organized paper.\n\n**Significance**\nThis work is of high importance.  It has done a great deal to close the gap between dense and sparse training for a constrained type of sparsity that offers practical acceleration in readily available hardware.  If it continues to be successful in broader experiments in the wild, then it could be widely adopted, saving significant resources in training new networks and reduce the latency of performing large-scale experiments.", "weaknesses": "**Originality**\nN/A\n\n**Quality**\nA missing experiment is one that would show the relative importance of scaling factor *beta*.  It would be a good addition to the ablation study presented in Table 7.\n\n**Clarity**\n\nTable 2 was confusing at first.  It wasn't obvious that the first row, without S-STE, was just dense training, rather than a different DST baseline.  Once I understood this, its low loss value not being the \"winner\" made sense.  The second confusing bit was that the second row had the next-lowest loss, but was also not **emphasized**.  Again, I realized that this is because it would not lead to acceleration in the backwards pass.  This second point could be clarified by adding \"... and accelerates the backwards pass\" to the caption.\n\nIn line 256, the authors say they \"choose to sparsify del(Z) only in the backwards pass,\" which sounds like in the forwards pass, they do not sparsify del(Z).  Clearly, they don't, because this term is not involved in the forwards pass, but it caught me off guard - I think \"choose to sparsify only del(Z) in the backwards pass,\" as opposed to sparsifying both del(Z) and S(W), is more clear.\n\nIn line 331, I think the authors mean to say that FP8 quantization accelerates GEMMs up to 2x faster than their 16b counterparts.  (As detailed in Section 5.2, the upper bound of speedup from dense BF16/FP16 to sparse FP8 is indeed 4x.)\n\nI noticed a few typos (there may be more lurking):\n\n- Line 13: \"our method surpass\" -> surpasses\n\n- Line 174: \"contains two main partitions\" -> parts\n\n- Line 201: \"the closer *t* is from |*t3*|\" -> \"the closer *t* is to |*t3*|\"\n\n- Line 263: \"forwrad\" -> \"forward\"\n\n- Line 326: \"leaverage\" -> \"leverage\"\n\n**Significance**\nTransformer blocks typically have two more weighted linear layers that were left dense in this work: the QKV projection before attention and the attention output projection.  Thus, there is potentially more memory and computation savings available, but without experiments, it is unknown if model quality will remain high if they were also made sparse.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer JQrs,\n\nThank you for the acknowledgment of the potential and effectiveness of our work and the detailed constructive comments. Below we provide a point-to-point response to all comments.\n\n**Weakness 1:** **Quality** A missing experiment is one that would show the relative importance of scaling factor *beta*. It would be a good addition to the ablation study presented in Table 7.\n\n**Reply:** We appreciate this insightful suggestion and add content to the ablation study as below. We add more experiments and redraw the ablation table. Results show that the effect of weight rescaling is not obvious on computer vision tasks like DeiT-small, but is significant for language models like Transformer-base. The reason behind this is due to the difference of their tasks: classification tasks are usually easier to complete than generation tasks, and the change may not be well reflected in the accuracy on simpler tasks.\n\n*Table.* Experimental result of S-STE (soft-thresholding and weight rescaling), MVUE and FP8 training with DeiT-small on ImageNet-1K.\n\n|soft-thresholding|weight rescaling|$\\operatorname{MVUE}(\\nabla_{\\mathbf{Z}}^\\top)$|FP8|comment|test acc1|test acc5|\n|-|-|-|-|-|-|-|\n|-|-|×|×|dense|79.9|95|\n|-|-|×|√|dense; FP8|79.7|94.9|\n|×|×|×|×|hard-thresholding|77.7|93.9|\n|√|√|×|×||78.8|94.6|\n|√|×|×|×||78.9|94.7|\n|√|√|×|√||78.6|94.4|\n|√|√|√|×||78.9|94.6|\n|√|×|√|×||78.2|94.2|\n|**√**|**√**|**√**|**√**||**78.5**|**94.4**|\n\nBesides, we'd like to kindly point out that another control experiment done on Transformer-base with WMT 14 En-De has already been presented; see Table 3 in the paper. To further clarify this, we expand this table to another ablation study, which presents the results with Transformer-base settings; see table below.\n\n*Table.* Experimental result of S-STE (soft-thresholding and weight rescaling), MVUE and FP8 training with Transformer-base on WMT 14 En-De.\n\n| soft-thresholding|weight rescaling|$\\operatorname{MVUE}(\\nabla_{\\mathbf{Z}}^\\top)$ | FP8 | comment | test BLEU | validation loss | average epoch loss |\n|-|-|-|-|-|-|-|-|\n|-|-|×|×|dense|26.42|3.977| 4.555 |\n|×|×|×|×|hard-thresholding|25.65|4.088 |4.686 |\n|√|×|×|×||25.28|4.044|4.67 |\n|√|√|×|×||26.3|4.007|4.605 |\n|√|√|√|×||25.93|4.01|4.602 |\n|**√**|**√**|**√**|**√**||**26.11**|**4.011**|**4.61**|\n\n\n\nThe new ablation study as well as the supplemental experimental data of the first one will be updated in our paper.\n\n**Weakness 2:** **Clarity** Table 2 was confusing at first. It wasn't obvious that the first row, without S-STE, was just dense training, rather than a different DST baseline. Once I understood this, its low loss value not being the \"winner\" made sense. The second confusing bit was that the second row had the next-lowest loss, but was also not **emphasized**. Again, I realized that this is because it would not lead to acceleration in the backwards pass. This second point could be clarified by adding \"... and accelerates the backwards pass\" to the caption.\n\n**Reply:** We realize the problem and redraw the table as follows. Besides, we will consider putting Table 2 beside Sec. 5, where it should have been. The previous deviation of format is due to insufficient layout.\n\n*Table.* Results of different MVUE strategies on GPT-2 774M with 4000 steps. Sparsifying $S(\\mathbf{W})^\\top$ and $\\nabla_{\\mathbf{Z}}^\\top$ can accelerate the two matrix multiplications of backward pass respectively. However, accuracy loss introduced by those sparse matrices are different.\n\n| S-STE | $\\operatorname{MVUE}(S(\\mathbf{W})^\\top)$ | $\\operatorname{MVUE}(\\nabla_{\\mathbf{Z}}^\\top)$ | comment | loss |\n|-|-|-|-|-|\n|-|×|×|dense|3.3948|\n|-|×|×|SR-STE|3.4739|\n|√|×|×||3.4333|\n|√|√|×||3.4644|\n|√|√|√||3.4773|\n|√|×|√||**3.448**|\n\n**Weakness 3:** **Clarity** In line 256, the authors say they \"choose to sparsify del(Z) only in the backwards pass,\" which sounds like in the forwards pass, they do not sparsify del(Z). Clearly, they don't, because this term is not involved in the forwards pass, but it caught me off guard - I think \"choose to sparsify only del(Z) in the backwards pass,\" as opposed to sparsifying both del(Z) and S(W), is more clear.\n\n**Reply:** The latter expression is definitely less confusing. Thanks for reminding us! We will update this in our paper.\n\n**Weakness 4:** **Clarity** In line 331, I think the authors mean to say that FP8 quantization accelerates GEMMs up to 2x faster than their 16b counterparts. (As detailed in Section 5.2, the upper bound of speedup from dense BF16/FP16 to sparse FP8 is indeed 4x.)\n\n**Reply:** Yes, that's exactly what we mean. In the updated version of our paper, this line will be replaced with \n\n> While 16-bit float tensors are widely used in pre-training, FP8 – where float numbers stored in 8 bits – is a popular quantization methods which theoretically accelerates GEMMs up to 4x faster than its fp32 counterparts and 2x faster than its FP16/BF16 counterparts.\n\n**Weakness 5:** **Typos** I noticed a few typos (there may be more lurking).\n\n**Reply:** We apologize for the oversight and would like to thank you for pointing them out. All the clerical errors that you mention and some other typos are corrected in the updated version.\n\n**Weakness 6:** **Significance** Transformer blocks typically have two more weighted linear layers that were left dense in this work: the QKV projection before attention and the attention output projection. Thus, there is potentially more memory and computation savings available, but without experiments, it is unknown if model quality will remain high if they were also made sparse.\n\n**Reply:** That's a very good question. We have done experiments with the four QKV projection layers, and the output is not satisfying. Sparsifying QKV projections may require attention specific methods, designed in conjunction with attention mechanism. We decide to leave them to future work.", "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": "Lc7M4s44XW", "reviewer": "Reviewer_73MA", "summary": "This paper presents a framework to circumvent the common challenges associated with STE-based 2:4 pre-training due to pruning function discontinuity. In particular, their framework addresses 3 aspects of this behavior: descent direction, amount of descent, and sparse mask oscillation.", "questions": "Please address the questions in the weaknesses section.", "limitations": "/", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* This paper provides a comprehensive analysis of the limitations of traditional pruning techniques, with a particular emphasis on how their framework addresses these challenges. I believe there is a novelty in the exploration of continuous pruning strategies.\n* Additionally, each component is comprehensively validated – for example, the problem of mask oscillation seems to be adequately documented in Section 3.3 with a viable exemplar to demonstrate impact.", "weaknesses": "* This work is centered on the potential of continuous pruning schemes for enabling faster sparse pre-training. However, I do see a strong resemblance with conventional soft mask approaches to pruning schemes. In particular, I wonder if the authors considered typical soft-pruning works in their comparisons, and if not, what was the reasoning for that experimental setting choice.\n* I found the experimental results to be quite sparse in model selection and competitive method benchmarking. Predominantly this method benchmarks against SR-STE and the dense framework, however, there isn’t enough context as to why other pruning-based methods are not included in scope. Further, there seem to be a few different tasks ablated however comprehensive details on why each application and/or model was selected are missing. For example, if demonstrating the efficacy of a sparse pre-training method it, would be beneficial to see the scaling effect on large models, say ViT-B/L for ImageNet-1K or the SWIN architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 73MA,\n\nThank you for the acknowledgment of the potential and effectiveness of our work and the detailed constructive comments. Below we provide a point-by-point response to all comments.\n\n**Weakness 1:** This work is centered on the potential of continuous pruning schemes for enabling faster sparse pre-training. However, I do see a strong resemblance with conventional soft mask approaches to pruning schemes. In particular, I wonder if the authors considered typical soft-pruning works in their comparisons, and if not, what was the reasoning for that experimental setting choice.\n\n**Reply:** We are not considering typical soft-pruning methods as our comparisons. These methods have different mission backgrounds and purposes from ours, and is meaningless to be used as contrast.\n\nIn \"2:4 pre-training\" context, we are talking about practically accelerating pre-training via 2:4 property [1]. In the light of this purpose, traditional dynamic sparse training methods such as RigL [2], Early-Bird Tickets [3,4], Sanity-Checking [5], Decaying Pruning [6] which can't possibly accelerate training shouldn't be considered. Besides, post-training methods such as  Deep Compression [7] which has nothing to do with pre-training phase shouldn't be considered too. The main experimental setting we choose has already covered the full spectrum of 2:4 related work, including hard-thresholding [8], SR-STE [8], STEP [9] and Bi-Mask [10]. Other 2:4 pre-training works, such as MVUE [11] and T-mask [12] focus on different questions and is also meaningless to be used as contrast.\n\n**Weakness 2:** I found the experimental results to be quite sparse in model selection and competitive method benchmarking. Predominantly this method benchmarks against SR-STE and the dense framework, however, there isn’t enough context as to why other pruning-based methods are not included in scope.\n\n**Reply:** For model selection, there are three reasons to explain why we choose these settings.\n\n1. **We are solving a very difficult task** and pre-training is quite consuming for academic research group. Different from post-training research which can leverage state-of-the-art model architectures like Llama 3, Phi-3 and Qwen2, pre-training where we **train a model from scratch** takes lots of hardware resources and thousands of GPU hours. Thus the model sizes are likely to be smaller than most of the other post-training pruning studies.\n\n2. As pre-training is very difficult, 2:4 pre-training is even more tough because we need to **jointly optimize activated masks and their weight** values. **The difficulty of this problem leads to insufficient baseline works from the past.** That's the reason we focus on this task; **we not only reveal the difficulty of the task, but also set up new SOTA baselines**. In other words, we are trying to **push the frontiers of relatively unpopular 2:4 pre-train research**.\n\n3. The models we use are chosen from a wide range of most representative works in deep learning with the acceptable size and available settings, spanning from natural language processing to computer vision tasks. For NLP tasks, we choose classical Transformer-base and GPT-2; for CV tasks, we choose popular Vision Transformer. We believe **these representative works can set up new SOTA baselines for future study to follow**.\n\nFor competitive method benchmarking and other pruning-based methods, please refer to Weakness 1.\n\n**Weakness 3:** Further, there seem to be a few different tasks ablated however comprehensive details on why each application and/or model was selected are missing. For example, if demonstrating the efficacy of a sparse pre-training method it, would be beneficial to see the scaling effect on large models, say ViT-B/L for ImageNet-1K or the SWIN architectures.\n\n**Reply:** For model selection, please refer to Weakness 2 and Weakness 1. For the scaling effect of large models, we pre-train different scales of GPT-2 and Vision Transformer. We list the result as follows and hopefully it solves your doubts. Basically, the performance get closer to or surpass dense baseline as model size increases.\n\n\n\n*Table.* Experimental results for DeiT. DeiT-base fail to due to infinity loss on SR-STE and S-STE.\n\n|Size|Method|Acc@1|Acc@5|\n|-|-|-|-|\n| DeiT-tiny |Original|72.2|91.1|\n||SR-STE|67.8|88.6|\n||**S-STE**| **68.5** |**88.9**|\n|DeiT-small| Original |79.9|94.5|\n||SR-STE [10]|75.7|-|\n||Bi-Mask [10]|77.6|-|\n||**S-STE**| **78.5** |**94.4**|\n\n*Table.* SQuAD scores of different sizes and pre-training methods on GPT-2. Similar to Table 5, we use 2:4 sparse weights to evaluate S-STE model, while dense parameters to evaluate the rest. Note that the results in Table 8 are close to the dense baselines, which means the room to improve is already small. Other experiments in Table 5 can show a larger improvement than SR-STE.\n\n|Params|Pre-training method|Fine-tuning method|EM|F1|\n|-|-|-|-|-|\n|124M|Dense| Dense|67.6|78.8|\n||Transposable SR-STE+Dense|Dense|67.5|78.5|\n||SR-STE|Dense|66.2|77.5|\n||**S-STE**| **S-STE**|**68**|**78.8**|\n|350M|Dense| Dense | 73.2 | 83.6 |\n||Transposable SR-STE+Dense|Dense|71.9|82.4|\n||SR-STE|Dense|72.0|82.4|\n||**S-STE**|**S-STE**|**72.2**|**82.7** |\n|774M|Dense|Dense|74.3|84.9|\n||Transposable SR-STE+Dense|Dense|74.3|84.6|\n||**S-STE**|**S-STE**|**75.5**|**85.5**|", "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": "iyu5zEk1xY", "reviewer": "Reviewer_WZj6", "summary": "The paper discusses three drawbacks of traditional N:M sparse training and suggests using soft-thresholding over hard-thresholding for 2:4 sparse pre-training. It introduces the idea of rescaling sparse weights with a fixed scaling factor per tensor. Results from experiments in machine translation, image classification, and large generative language models demonstrate its effectiveness.", "questions": "I have the following comments, which may improve the quality of the paper.\n\n1.\tThe discussion of formula (4) should use very rigorous mathematical derivation, however, the discussion of related contents is very imprecise.\n\n2.\tIn Theorem 4.1, it is better to provide the definition of continuous projection for a vector. Additionally, it is important to address why hard-thresholding and other algorithms fail to meet this property, and experimentally demonstrate the performance improvement that can result from adhering to this property.\n\n3.\tIn Table 1, is the \\gamma choice dependent on the \\beta setting? Does it establish the beta before selecting \\gamma?\n\n4.\tWhat is the reason for simulating dense weights? Why perform fixed weight rescaling \\beta?\n\n5.\tIn Table 4, it is better to separate “dense” and use it as a baseline. The best result in the second column is marked incorrectly.\n\n6.\tThere are many methods for N:M sparse pre-training as shown in the Related Work section. The paper needs to make a detailed comparison with these methods.\n\n7.\tThe core of this paper is the pre-training acceleration. However, it is far from enough to discuss the possible acceleration in theory, and it is necessary to carry out sufficient experiments to compare the acceleration effects of various methods.\n\n8.\tAll the figures in the paper are too small to see clearly.\n\n9.\tMost references lack the names of conferences or journals.", "limitations": "Not applicable.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.\tThe idea is simple and clear. The proposed approach is very easy to understand and implement.\n\n2.\tThe paper discusses several important issues for N:M sparse pre-training. These discussions are very useful for practice.\n\n3.\tThe experiments verify the selection of some important parameters, and also show the effectiveness of the proposed approach on several tasks.", "weaknesses": "1.\tMost of the techniques are discussed in other tasks, and the paper is only applied to N:M sparse pre-training, which is not innovative enough. So the contributions are not enough for NeurIPS. \n\n2.\tIt is very important to discuss the drawbacks of existing N:M sparse pre-training, but these discussions have not led to a new algorithm. There is insufficient inevitability between these discussions and the proposed algorithm.\n\n3.\tThe experiments are weak. There is a lack of important comparison methods and real training acceleration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer WZj6,\n\nThank you for the acknowledgment of the potential and effectiveness of our work and the detailed constructive comments. Below we provide a point-to-point response to all comments.\n\n**Weakness 1:** Most of the techniques are discussed in other tasks, and the paper is only applied to N:M sparse pre-training, which is not innovative enough. So the contributions are not enough for NeurIPS.\n\n**Reply:** We argue that the contributions are adequate:\n\n1. Although most of the techniques we discuss is rooted in prior work, we propose to **innovatively implementing soft-thresholding on 2:4 pre-training tasks**. Besides, we make a couple of **highly original 2:4 specific modifications** to utilize those methods (choosing optimal threshold, freezing scaling factor, minimizing MSE).  These are the **technological contributions** of our study.\n2. More importantly, we are trying to **push the frontiers of relatively unpopular 2:4 pre-train research**: on the one hand, as pre-training is very difficult, 2:4 pre-training is even more tough because we need to **jointly optimize activated masks and their weight** values. In this way, **our work reveal the difficulty of the problem**. On the other hand, prior work on 2:4 pre-training topic mostly focus on optimization skills, however **they never consider continuity as a main problem**. Our study tries to get to the essence of the question and we provide a deeper understanding of it. We **not only set up new SOTA baselines for future study to follow, but also provide an important and new research insight for future work to explore**.\n\n**Weakness 2:** It is very important to discuss the drawbacks of existing N:M sparse pre-training, but these discussions have not led to a new algorithm. There is insufficient inevitability between these discussions and the proposed algorithm.\n\n**Reply:** We believe that inevitability is quite clear:\n\n1. The most important drawback we see from previous work is discontinuity. To elaborate this, we observe three phenomena:  incorrect descending direction, inability to predict the amount of descent and oscillation (Sec. 3).\n1. To overcome the drawback of previous pruning function, we introduce 2:4 specific soft-thresholding pruning function to S-STE and theoretically prove its continuity.\n1. To complete the circle, we show that **S-STE successfully overcome the discontinuity phenomena mentioned above** via experiment; see Sec. 4, Fig. 2(c)(d) and 4(d).", "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:

Meta-review: Almost all of the reviewers agreed that the work was well-motivated, well-written and had significant novelty. Reviewer's concerns primarily centred on the lack of real-world timings for the proposed method, increased computational requirements of the proposed S-STE method, the modest gains in accuracy shown, lack of ablation over the scaling hyperparameter $\beta$, a similarity with soft-mask approaches for pruning, and a lack of comparison with pruning methods. 

Overall it appears that many of these concerns appear to have been addressed in the rebuttal by the authors. The post-rebuttal discussion amongst all the reviewers (except one) was also productive and further clarified I believe that that the most significant of each reviewer's concerns were addressed by the authors. As such, with the feedback from the reviewers/rebuttal incorporated, I believe this paper will be a relevant and significant contribution.
Decision: Accept (poster)