Title: S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training
Venue: NeurIPS 2024
Accepted: yes
Keywords: efficient machine learning

transformers

2:4 sparsity

dynamic sparse training
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

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 pretraining methods ( e.g. 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 light of this, 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 surpasses previous 2:4 pre-training recipes and is comparable even with full parameter models. Our toolkit is available at https://github.com/huyz2023/2by4-pretrain .

1 Introduction

Large scale transformers have achieved many impressive results such as chatbots [43], 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 accelerate deep learning computation is sparsity. N:M sparsity [31] is a hardware-friendly sparsity pattern, where every group of M dimensions only has N non-zero entries. Nvidia Ampere GPUs can multiply a 2:4 sparse matrix with a dense matrix, twice as fast as multiplying two dense matrices.

While N:M sparsity has been successfully applied to accelerate inference [31, 40, 15, 35, 9], extending 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 at 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 as single-shot pruning [24], iterative magnitude pruning [13, 29], and RigL [11] cannot be directly applied to this scenario. Moreover, besides forward propagation, the matrix multiplications in back propagation must be sparsified as well, to provide reasonable training speedup.

Methods based on the straight-through estimator (STE) [52, 2] have shown promise towards solving the challenging problem of sparse pre-training. They maintain a dense weight, which is sparsified in each iteration for fast forward&backward computation, and the dense weight is then updated with STE gradients. In this way, connectivity patterns can be learned jointly with weights in an end-to-end

fashion with stochastic gradient optimizers. SR-STE [52] is such a method to train sparse networks from scratch, with a regularization term to stabilize the training. Several subsequent works [21, 51, 8] 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 a regularization strength hyper-parameter, which is hard to tune. 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 straight-through estimator (S-STE) as a solution. Cruically, S-STE introduces a 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-error between original dense weight vector and sparse weight 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-based optimizers can be readily used. Furthermore, S-STE does not introduce any hyper-parameter, so its practical adoption is easier than SR-STE.

We devise comprehensive pre-training 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 a neural network can be formalized as an optimization problem min w F ( w ) , where w ∈ R d is the parameter and F is a differentiable empirical risk function: F ( w ) = R n ( w ) = 1 n n ∑ f ( w ; ξ [ i ] ) . Here, f is the loss function, n is the size of data set D = { ξ [ i ] } n i =1 and ξ [ i ] is the i -th i =1 sample. The optimization can be solved with standard stochastic gradient method (SG) [4]. Suppose the network is initialized with w 1 , { α k } is a positive learning rate sequence, and ξ [ i k ] is randomly chosen from { ξ [ i ] } n i =1 . Then, iteratively we have w k +1 = w k -α k ∇ w k f ( w k ; ξ [ i k ] ) . As we consider pre-training tasks, w 1 is simply a random initialization.

The training of a sparse network involves optimizing the parameter w in a constrained space W ⊂ R d . For an N:M-sparse 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:

Here, S is a pruning function which converts a dense weight w to a sparse weight ˜ w ∈ W . One common choice is the hard-thresholding pruning function [52, 45]. For every block of four adjacent elements a = [ a 1 , ..., a M ] ⊤ ∈ R M in the weight vector w , the pruning function can be defined as

where t is N -th largest element in a . 1 This essentially performs magnitude-based pruning, by zeroing out the two smallest elements. The hard thresholding function can also be written as S h ( a ) = a ⊙ m h ( a ) , where m h ( a ) is a 0/1 mask vector, with ( m h ( a )) i = 1 if | a i | > t .

However, Eq. (1) cannot be directly optimized since the pruning function S is not differentiable. Particularly, the derivative of the hard-thresholding function S h is undefined on 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 f ≈ ∇ ˜ w f and therefore ∂S h ( a ) /∂ a ≈ I :

With the pruning function and STE, each iteration of sparse training involves: (1) prune the dense weight to get the sparse weight: ˜ w = S ( w ) ; (2) compute the loss and gradient with the sparse weight; and (3) update the dense weight with the gradient. Among these, step 2 is most time-consuming, and it can be accelerated with sparse tensor cores given ˜ w is N:M-sparse. Next, we will focus on the optimization aspects of sparse training and defer the discussion of computation details to Sec. 4.

3 The curse of discontinuity

Classical stochastic optimization theory [4] 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 is the discontinuity of the pruning function S h , as visualized in Fig. 2. For a discontinuous function, an arbitrarily small change in input a can cause an unbounded change of the output S h ( a ) . Such discontinuity appears on the boundary when the N -th and N +1 -th largest elements have same magnitude. For example, for a 1:2-sparse pruning function, S h (1 , 0 . 999) = (1 , 0) , but S h (0 . 999 , 1) = (0 , 1) , and the boundary is the line a 1 = a 2 .

̸

When S h is discontinuous, the parameter space R d can be partitioned into regions {W m | m ∈ M} , where M ⊂ { 0 , 1 } d is the space of 0/1 masks with N:M pattern, and all the parameters in each region w ∈ W m have the same mask m h ( w ) = m . The loss landscape F ( S h ( w )) = F ( m h ( w ) ⊙ w ) is continuous and differentiable within each region, where gradient-based algorithms can work well. However, when the optimization trajectory crosses the boundary: m h ( w k +1 ) = m h ( w k ) , the behavior 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.

3.1 Phenomenon 1: incorrect descending direction

Here, we run a gradient descent algorithm (without stochasticity) on a small dataset. For a dense model where F is differentiable, with Taylor's formula we should have

That is, the objective function will monotonically decrease in each iteration once the learning rate α k is sufficiently small. However, it is not the case for sparse training. In Fig. 2(d), we measure the distribution of the amount of descent (AoD) ∆ F k := F ( w k ) -F ( w k +1 ) for training a 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 ( w k , m h ( w k )) generated by the training algorithm. We compare the following two quantities: the AoD by updating both weight and mask ∆ F 1 = F ( w k ⊙ m k ) -F ( w k +1 ⊙ m k +1 ) and the AoD by only updating the weight ∆ F 2 = F ( w k ⊙ m k ) -F ( w k +1 ⊙ m k ) . In Fig. 1, we can observe ∆ F 2 is mostly positive due to the piecewise continuity of F . However, ∆ F 1 is frequently negative and very often even smaller than ∆ F 2 (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 of 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, in the sense that they fail to predict the AoD at each step. From Eq. (4), we should have F ( w k ) -F ( w k +1 ) ≈ ( ∇ w k F ( w k )) ⊤ ( w k -w k +1 ) , where the left hand side is the actual

̸

AoD, and the right hand side is the predicted AoD. We plot the actual AoD against predicted AoD for dense (Fig. 2(a)) and sparse training (Fig. 2(b)). 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 w ∈ W m h ( w ) . Once w k +1 -w k is large enough that m h ( w k +1 ) = m h ( w k ) , the function crosses a border of S h , and F will have a sudden change which is unpredictable by the gradient.

0.04

(a)

0'0

(b)

0.04

(c)

1.0

(P)

0.03

0.03

0.03

0.8

cumulative probability

0.02

0.02

0.02

0.6

hard-thresholding

actual

actual

dense

10°0

0.01

0.01

0.4

soft-thresholding

00°0

00'0

<：

000

0.2

0.01

0.01

0.00

0.01

0.02

0.03

0°0

-0.01

10'0-

00'0

10'0

0.02

0.03

0.04

0.01

0.01

0.00

0.01

0.02

0.03

0'0

0.0

0101500050-01-1-0

expected

expected

expected

△F

1e2

3.3 Phenomenon 3: oscillation

Oscillation is 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 a toy example.

Flip rate Flip rate is a simple metric to measure the stability of sparse training [20]: r k = ∥ m h ( w k ) ⊕ m h ( w k -1 ) ∥ 1 /d, where ⊕ indicates XOR operation. As Hu et al. [20] points out, taking the flip rate of the dense model as standard, they observe larger flip rate of hard-thresholding: when training transformers, the flip rate can stay at 6% in the entire training process. However, a healthy training process should have a 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 hard-thresholding on the smallest problem, we devise a simple toy problem that contains two parameters: min w 1 ,w 2 g ( w 1 , w 2 ) = ( w 1 -w 2 ) 2 . 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 α k that can make the dense model converge to global minima, STE with hard-thresholding fails to converge and it oscillates back and forth.

First, for the dense model, the global minima lies on the line w 1 = w 2 . Suppose we start from w 1 = [0 . 2 , 0 . 1] ⊤ , by taking α k = 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 w 1 = w 2 = 0 . By starting from w 1 = [0 . 2 , 0 . 1] ⊤ and taking α k = 0 . 25 , we invariably jumps between w 2 t +1 = [0 . 2 , 0 . 1] ⊤ and w 2 t = [0 . 1 , 0 . 2] ⊤ , and g 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 on hard-thresholding can be explained by 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.

3.4 Overcoming the curse of discontinuity

One way to mitigate discontinuity is sparse-refined straight-through estimator (SR-STE), which adds a sparse-refined regularization on the gradients [52]: min w F ( ˜ w ) + λ W 2 ∥ w ⊙ m ( w ) ∥ 2 2 . 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 hyper-parameter λ W . Second, the new regularization term leads to a competition between loss function and sparse regularization. Finally, the loss function is still discontinuous unless λ W →∞ .

From the above analysis, discontinuity causes optimization problems. It would be ideal to have a continuous pruning function, yet the iterates ( ˜ w k ) still need to be sparse during the entire training process.

(a)

(b)

(c)

(d)

6.00

6.00

-6.00

-3.00

0.00

3.00

6.00

3.00

-3.00

0.00

-6.00

-3.00

-6.00

0.00

6

3.00

00°9

3.00

0.00

-3.00

-6.00

N

2

4

4

6

-6

①2

2

4

6

6

1

4

6

6 4 2

2

4

6

6

4 Methodology

In this section we propose a training algorithm (smooth straight-through estimator, S-STE) that contains two main parts combined with STE: 2:4 specific softthresholding, and fixed weight rescaling . They together work as the sparsifying function described in Sec. 2: ˜ w = S ( w ) = βS soft ( w ) . Results in Fig. 2(c)(d) and 4(d) 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.

4.1 2:4 specific soft-thresholding S soft

Motivation for the design As discussed in Sec. 3, 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-thresholding function for traditional pruning in Vanderschueren and Vleeschouwer [45] as our 2:4 specific soft-thresholding. Given a vector a = [ a 1 , a 2 , a 3 , a 4 ] ⊤ ∈ R 4 , 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 1 , t 2 , t 3 , t 4 ] is an rearrangement of a ⊤ , s.t. | t 1 | ≤ | t 2 | ≤ | t 3 | ≤ | t 4 | . Then, the pruning function can be defined as

The plots of soft-thresholding is drawn in Fig. 3, showing S soft is continuous everywhere. Note that although we define S soft by a block a ∈ R 4 , S soft can be extended to arbitrary a ∈ R 4 t for t ≥ 1 , by doing block-wise pruning.

Theorem 4.1. S soft ( a ) is a continuous projection for a ∈ R d .

A detailed discussion of the proof can be found in Appendix A.1.

Table 1: Validation loss and test accuracy of S-STE with different γ on Transformerbase.

Choosing optimal threshold Theoretically, any real number in [ | t 2 | , | t 3 | ] can be used as a feasible threshold. This gives us infinite options and we describe it with an interpolation as t = γ | t 2 | +(1 -γ ) | t 3 | with γ ∈ [0 , 1] . The larger γ is, the closer t is to | t 3 | , and the smaller ∥ S soft ( a ) ∥ is. This may affect model's capacity. In order to maximize the retention of information, using a small γ is necessary. In our method we propose to set γ = 0 . Experimental results in Table 1 also show that the network has the best accuracy when γ = 0 , i.e. , t = | t 2 | .

4.2 Fixed weight rescaling β

Because S soft reduce the total magnitude of w , it is not a close simulation of dense weights. Like Vanderschueren and Vleeschouwer [45], we scale up S soft ( w ) in our method as S ( w ) = βS soft ( w ) , but we modify weight rescaling in their study to adapt to our approach. First, we use a per-tensor scale β rather than a per-channel β 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-square-error (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 a dynamic β for every iteration, we argue that this doesn't align with our approach. We explain our solutions in two parts.

2.5le-3

(a)

18

(b)

4.0.le-3

(c)

2.5.1e3

(d)

β=1

16

epoch=5

3.5

fixed β

dense

2.0

β=2

14

epoch=10

3.0

dynamic β

2.0

SR-STE

rate

1.5

1.0

β = 4

12

10

8

6

epoch=15

epoch=25

epoch=20

flip r

2.5

2.0

1.5

1.0

1.5

S-STE

0.5

守

epoch=30

1.0

0.5

08.0

0.2

0.4

step

0.6

0.8

1.0

1e5

1.2

2

0

0

10

layer

15

20

25

0.5

0.0

00

0.2

0

step

0.6

0.8

1.0

1e5

1.2

0.8.0

0.2

0

step

0.6

0.8

1.0

1e5

1.2

First, we find it interesting that β have a subtle correlation with flip rate: in a sparse model, larger β 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 ∇ w f ≈ ∇ ˜ w f is valid when ˜ w and w are close enough. However, if scale is too large then optimal, w and ˜ w are too far apart to guarantee this. Such a mismatch leads to incorrectness in descending directions, and thus unstability in optimization increase flip rate; see Fig. 4(a).

Second, we argue that dynamically computing scaling factor for each iteration leads to high flip rate in our training process. Fig. 4(b) shows the results dynamically changing β will make it increase with iterations, especially for later layers. Fig. 4(c) shows flip rate of this network, which has a significantly higher tail than the dense one. Considering high flip rate is harmful, we propose to compute scaling factor β only in the first iteration. After that, we use the same β in the rest of the training. Fig. 4(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 S soft ( w ) to have the same L1-norm as w : β = ∥ w ∥ 1 / ∥ S soft ( w ) ∥ 1 . However, we choose to minimize the MSE gap of S soft ( w ) and w . As [8] point out, sparsifying weights in the forward pass should minimize MSE rather than an unbiased estimation. In our method, to determine an optimal scale β , we need to minimize

Rearrange the terms and taking partial derivative of β , we choose β = w ⊤ S soft ( w ) / ∥ S soft ( w ) ∥ 2 . The comparison between no scaling, our minimizing MSE and keeping L1-norm can be found in Table 2. Result show that our method yields the best results in practice.

5 Other implementation skills

5.1 Minimum-variance unbiased estimation

To accelerate the backward propagation, Chmiel et al. [8] suggest using a minimum-variance unbiased estimator (MVUE). For every linear layer Z l = X l S ( W l ) ⊤ , there are two matrix multiplications of the backward pass in total: ∇ X l = ∇ Z l S ( W l ) and ∇ W l = ∇ ⊤ Z l X l , where X l is the input of the l -th layer, W l and Z l are the weight matrix and output activation. We conduct MVUE on both two matrix multiplications and compare their results: ∇ X l = ∇ Z l MVUE( S ( W l ) ⊤ ) ⊤ and ∇ W l = MVUE( ∇ ⊤ Z l ) X l . Specifically, we choose S ( W l ) and ∇ Z l 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 al. [8] that minimum loss of accuracy is guaranteed for MVUE on ∇ Z l . Secondly, using MVUE on S ( W l ) will make errors accumulate along the back propagation, and results in large standard deviation of gradient for the first few layers. Besides, results in Table 3 also show minimum loss of accuracy of MVUE on ∇ Z l while obvious accuracy loss on S ( W l ) . Thus, we choose to sparsify only ∇ Z l in the backward pass.

5.2 FP8 training

To further accelerate pre-training of networks, we utilize popular FP8 workflow in training. Similar to Transformer Engine 2 , we use FP8 e3m4 in forward pass and e5m2 in backward pass. Besides, we 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 can accelerate an additional 2x on 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.

For Transformer, we train Transformer-base models on WMT 14 En-De dataset [3] with fairseq [34] codebase and evaluate it with BLEU [36] scores. For DeiT, we pre-train Deit-small model for ImageNet-1K [10] classification task. For GPT-2, we pre-train GPT-2 124M, 350M and 774M models on OpenWebText [16] and evaluate it on GLUE [47] and SQuAD [39] benchmarks. We also compare our method with state-of-the-art 2:4 training methods (SR-STE [52], Bi-Mask [51] and STEP [28]). The pre-training and evaluation scripts are publicly available at https://github.com/thu-ml/ 2by4-pretrain-acc-examples .

Machine translation We first apply S-STE to train a 12-layer Transformer-base and compare it 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 4. Compared with SR-STE, our method improves by 0.3 and 0.5 on test set and validation set respectively, which is the closest to baseline. Besides, we improve by 0.6 compared to STEP on test set.

Image classification We further investigate the effectiveness of S-STE to train DeiT-tiny and DeiT-small on ImageNet-1k; see Table 5. 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% top1 accuracy improvement) and popular SR-STE method (2.8% top1 accuracy improvement).

Generative language models We compare S-STE with dense, normal SR-STE and SR-STE with dense fine-tuning [20] (SR-STE+DF) models on GLUE and SQuAD tasks. The SR-STE+DF models first use SR-STE to train a 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 to dense models. Results in Table 6 and 9 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.

Fine-tuning We illustrate the viability of S-STE for fine-tuning a pre-trained model, presenting a coherent workflow of accelerating both training and inference (dense fine-tuning cannot produce a sparse model for inference acceleration); see Table 7, 6.

Table 6: SQuAD and GLUE scores of different sizes and pre-training methods on GPT-2. We use 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 al. [20]. S-STE settings here include backward MVUE & FP8 training.

Ablation study In this part, We explore the effectiveness of S-STE, MVUE, and FP8 separately. We pre-train DeiT-small model on ImageNet-1K dataset for image classification. Combinations of these partitions in Table 8 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 acc1).

Acceleration For acceleration, we measure the acceleration ratio of a typical GPT-2 model using implementation from Hu et al. [20]. Note that on H100 GPUs, FP8 2:4-spMM 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 [17, 18, 12, 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, it is hard for these methods to work well on modern GPUs. Coarse-grained structured sparsity [49, 23, 19, 22] takes effect to acceleration, but since they often remove a whole chennel or a block, loss of information is non-negligible.

Fine-grained N:M sparsity for inference and pre-training Among all pruning techniques for pre-training, N:M sparsity is a promising approach towards accelerating large models, which is also known as fine-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, 1, 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 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 [30, 37, 44, 48, 32]. With e3m4 data format used in forward and e5m2 format [41] in backward, pre-trained 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 in a continuous manner. Experiments show that our method surpasses previous state-of-the-art methods on a 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, our current choice of continuous pruning function represents only one possible solution; alternative, smoother pruning functions may be necessary to achieve improved continuity and mitigate potential discontinuities.

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 National Key Research and Development Program of China (No. 2021ZD0110502) and NSFC Project (Nos. 62376131). J.Z is also supported by the XPlorer Prize.

References

Abhimanyu Rajeshkumar Bambhaniya, Amir Yazdanbakhsh, Suvinay Subramanian, Sheng-Chun Kao, Shivani Agrawal, Utku Evci, and Tushar Krishna. Progressive gradient flow for robust n:m sparsity training in transformers, 2024.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation, 2013.

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 Tamchyna. Findings of the 2014 workshop on statistical machine translation. In WMT@ACL , 2014. URL https://api.semanticscholar.org/CorpusID:15535376 .

Léon Bottou, Frank E. Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning, 2018.

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.

Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia 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 15834-15846. Curran Associates, Inc., 2020. URL https://proceedings.neurips. cc/paper_files/paper/2020/file/b6af2c9703f203a2794be03d443af2e3-Paper.pdf .

Xiaohan Chen, Yu Cheng, Shuohang Wang, Zhe Gan, Zhangyang Wang, and Jingjing Liu. Earlybert: Efficient bert training via early-bird lottery tickets, 2021.

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://openreview.net/forum?id=vuD2xEtxZcj .

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 to controlled text generation, 2020.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition , pages 248-255, 2009. doi: 10.1109/CVPR.2009.5206848.

Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. 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 2943-2952. PMLR, 13-18 Jul 2020. URL https://proceedings.mlr.press/v119/evci20a.html .

Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations , 2019. URL https://openreview. net/forum?id=rJl-b3RcF7 .

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode connectivity 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. URL https://proceedings.mlr.press/v119/ frankle20a.html .

Jonathan Frankle, Gintare Karolina Dziugaite, Daniel M. Roy, and Michael Carbin. Stabilizing the lottery ticket hypothesis, 2020.

Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot, 2023.

Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skylion007.github.io/ OpenWebTextCorpus , 2019.

Song Han, Jeff Pool, John Tran, and William J. Dally. Learning both weights and connections for efficient neural networks, 2015.

Song Han, Huizi Mao, and William J. Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding, 2016.

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.

Yuezhou Hu, Kang Zhao, Weiyu Huang, Jianfei Chen, and Jun Zhu. Accelerating transformer pre-training with 2:4 sparsity. In 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 19531-19543. PMLR, 21-27 Jul 2024. URL https://proceedings.mlr.press/v235/hu24r.html .

Itay Hubara, Brian Chmiel, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. Accelerated 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, Advances in Neural Information Processing Systems , volume 34, pages 21099-21111. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper_files/paper/2021/file/ b0490b85e92b64dbb5db76bf8fca6a82-Paper.pdf .

François Lagunas, Ella Charlaix, Victor Sanh, and Alexander Rush. Block pruning for faster transformers. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 10619-10629, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.emnlp-main.829. URL https://aclanthology.org/2021.emnlp-main.829 .

Mike Lasby, Anna Golubeva, Utku Evci, Mihai Nica, and Yani Ioannou. Dynamic sparse training with structured sparsity, 2023.

Namhoon Lee, Thalaiyasingam Ajanthan, and Philip Torr. SNIP: SINGLE-SHOT NETWORK PRUNING BASED ON CONNECTION SENSITIVITY. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id=B1VZqjAcYX .

Yuanzhi Li and Yang Yuan. Convergence analysis of two-layer neural networks with relu activation. 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 https://proceedings.neurips.cc/paper_files/paper/2017/file/ a96b65a721e561e1e3de768ac819ffbb-Paper.pdf .

Yujun Lin, Song Han, Huizi Mao, Yu Wang, and William J. Dally. Deep gradient compression: Reducing the communication bandwidth for distributed training, 2020.

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: A review on background, technology, limitations, and opportunities of large vision models, 2024.

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://proceedings.mlr.press/v202/lu23c.html .

Jaron Maene, Mingxiao Li, and Marie-Francine Moens. Towards understanding iterative magnitude pruning: Why lottery tickets win, 2021.

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.

Asit Mishra, Jorge Albericio Latorre, Jeff Pool, Darko Stosic, Dusan Stosic, Ganesh Venkatesh, Chong Yu, and Paulius Micikevicius. Accelerating sparse deep neural networks, 2021.

Badreddine Noune, Philip Jones, Daniel Justus, Dominic Masters, and Carlo Luschi. 8-bit numerical formats for deep neural networks, 2022.

OpenAI, 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-Shapiro, Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brockman, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, 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-Lopes, Jonathan Gordon, Morgan Grafstein, Scott Gray, 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 Jonn, Heewoo Jun, Tomer Kaftan, Łukasz 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, Łukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger,

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, Reiichiro Nakano, Rajeev Nayak, Arvind Neelakantan, Richard Ngo, Hyeonwoo Noh, Long Ouyang, Cullen O'Keefe, Jakub 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 Proehl, 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 Shieh, 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 Zoph. Gpt-4 technical report, 2024.

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.

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: //api.semanticscholar.org/CorpusID:10798336 .

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.1073135. URL https://aclanthology.org/P02-1040 .

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.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.

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 Carreras, editors, Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing , pages 2383-2392, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/v1/D16-1264. URL https://aclanthology.org/D16-1264 .

Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. A simple and effective pruning approach for large language models. arXiv preprint arXiv:2306.11695 , 2023.

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.semanticscholar.org/CorpusID:202779157 .

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention, 2021.

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 arXiv:2302.13971 , 2023.

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.

Antoine Vanderschueren and Christophe De Vleeschouwer. Are straight-through gradients and softthresholding all you need for sparse training?, 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V on 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 https://proceedings.neurips.cc/paper_files/ paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf .

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Tal Linzen, Grzegorz Chrupała, and Afra Alishahi, editors, Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP , pages 353-355, Brussels, Belgium, November 2018. Association for Computational Linguistics. doi: 10.18653/v1/W18-5446. URL https://aclanthology.org/W18-5446 .

Naigang Wang, Jungwook Choi, Daniel Brand, Chia-Yu Chen, and Kailash Gopalakrishnan. Training deep neural networks with 8-bit floating point numbers, 2018.

Lu Yin, Gen Li, Meng Fang, Li Shen, Tianjin Huang, Zhangyang "Atlas" Wang, Vlado Menkovski, Xiaolong Ma, Mykola Pechenizkiy, and Shiwei Liu. Dynamic sparsity is channel-level sparsity learner. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems , volume 36, pages 67993-68012. Curran Associates, Inc., 2023. URL https://proceedings.neurips.cc/paper_files/paper/2023/file/ d6d0e41e0b1ed38c76d13c9e417a8f1f-Paper-Conference.pdf .

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.

Yuxin Zhang, Yiting Luo, Mingbao Lin, Yunshan Zhong, Jingjing Xie, Fei Chao, and Rongrong Ji. Bidirectional 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 4148841497. PMLR, 23-29 Jul 2023. URL https://proceedings.mlr.press/v202/zhang23ae.html .

Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. Learning n:m fine-grained structured sparse neural networks from scratch. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id=K9bw7vqp_s .

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, TsangWei 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 Choromanski, Xi Chen, Yevgen Chebotar, Justice Carbajal, Noah Brown, Anthony Brohan, Montserrat Gonzalez Arenas, and Kehang Han. Rt-2: Vision-language-action models transfer web knowledge to robotic control. In Jie Tan, Marc 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 https://proceedings.mlr.press/v229/zitkovich23a.html .

A Appendix / supplemental material

A.1 Proof of Theorem 4.1

Proof. We prove this by demonstrating S soft is continuous on every 4-element block. Of note, suppose a = [ a 1 , a 2 , a 3 , a 4 ] ⊤ . Assume, without loss of generality, that | a 1 | ≤ | a 2 | ≤ | a 3 | ≤ | a 4 | . Our goal is to prove ∀ ϵ > 0 , ∃ δ > 0 , s.t. when | a ′ 1 -a 1 | < δ, | a ′ 2 -a 2 | < δ, | a ′ 3 -a 3 | < δ and | a ′ 4 -a 4 | < δ, | ( S soft ( a ′ )) i -a i | < ϵ , where a ′ = [ a ′ 1 , a ′ 2 , a ′ 3 , a ′ 4 ] ⊤ .

1) We start from the simplest case where | a 1 | < | a 2 | < | a 3 | < | a 4 | . Then we have

This order holds when

Thus,

The signs of a

We have

is unchanged when

Take δ ≤ 1 2 ϵ and this is done. It is similar to prove that sign( a ′ 3 )( | a ′ 3 |-| a ′ 2 | ) -sign( a 3 )( | a 3 |-| a 2 | ) ≤ ϵ using the same method.

2) We then consider the cases where there are two equivalents in a . If | a 1 | = | a 2 | < | a 3 | < | a 4 | or | a 1 | < | a 2 | < | a 3 | = | a 4 | , the proof should follow 1) as no flip happens. Thus we only consider the situation where | a 1 | < | a 2 | = | a 3 | < | a 4 | . Under these circumstances, a flip will happen on the second and third dimensions of a .

Without loss of generality we assume | a ′ 1 | < | a ′ 2 | ≤ | a ′ 3 | < | a ′ 4 | . Thus,

The proof of the fourth dimension is similar to 1), so we only focus on a 3 .

Take δ ≤ 1 2 ϵ 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 can be reduced to 1) or 2).

A.2 GLUE scores of GPT-2

See Table 9.

A.3

Acceleration

See Table 10.

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 [31]) is not satisfactory, which causes inconvenience on implementation. The peak FLOPS of 2:4-spMM is lower than theoretical GEMM FLOPS; see Table 11.

A.5 Broader Impact

S-STE can be used mainly to accelerate the pre-training stage of large-scale networks, like LLaMA [43] and GPT-4 [33]. Theoretically GEMMs of the FFN layer can be accelerated up to 4x faster than FP16 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 noncompliance 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.

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 in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer: [Yes]

Justification: Result can be reproduced according to 'methodology' and 'experiments' chapters.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to 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 a 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 a 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 it 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 to 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 a 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 submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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., to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

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 of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: For those experiments with large errors, we 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs 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 https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The authors have reviewed the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 to Appendix.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper doesn't release any new models, generators or datasets.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Please refer to the 'experiments' chapter.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

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 to README.md in our code submission.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model 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 as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or 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 nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1.\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.\n\nWEAKNESSES:\n1.\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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\nPlease address the questions in the weaknesses section.\n\nLIMITATIONS:\n/", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\n**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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nPlease refer to the weakness.\n\nLIMITATIONS:\nThere is no end-to-end acceleration results evaluated of current methods.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}