Exposing and Exploiting Fine-Grained Block Structures for Fast and Accurate Sparse Training
Abstract
Sparse training is a popular technique to reduce the overhead of training large models.
Although previous work has shown promising results for nonstructured sparse models, it is still unclear whether a sparse model with structural constraints can be trained from scratch to high accuracy.
In this work, we study the dynamic sparse training for a class of sparse models with shufﬂed block structures.
Compared to nonstructured models, such ﬁne-grained structured models are more hardware-friendly and can effectively accelerate the training process.
We propose an algorithm that keeps adapting the sparse model while maintaining the active parameters in shufﬂed blocks.
We conduct experiments on a variety of networks and datasets and obtain positive results.
In particular, on ImageNet, we achieve dense accuracy for ResNet50 and ResNet18 at 0.5 sparsity.
On CIFAR10/100, we show that dense accuracy can be recovered at 0.6 sparsity for various models.
At higher sparsity, our algorithm can still match the accuracy of nonstructured sparse training in most cases, while reducing the training time by up to 5x due to the ﬁne-grained block structures in the models.
Introduction
As large neural networks keep advancing the state-of-the-art in many machine learning tasks, training the models has become increasingly expensive.
It is reported that training GPT-3 costs an estimated $12 million in computational resources [41].
To reduce the cost, there is a growing interest in developing sparse training algorithms [30, 2, 8, 18, 26, 27].
By computing on a small subset of parameters throughout the training process, these algorithms aim to reduce the memory and computing power consumption for large-scale neural network training.
The existing sparse training methods can be categorized as either static or dynamic.
While training a static sparse model from scratch is appealing, it usually requires nontrivial initialization to achieve good accuracy [9, 23, 22].
In contrast, dynamic sparse training uses simple random initialization.
In every few iterations, it drops a small number of ‘unimportant’ weights and replaces them with some new parameters.
By exploring a larger set of parameters than the sparse model itself, dynamic sparse training, in general, achieves higher accuracy than static sparse training [14].
Most of the existing dynamic sparse training methods focus on nonstructured models.
That is, the active parameters are scattered in the model without any constraints.
Although they can retain high accuracy with much fewer parameters, the nonstructured sparse models do not necessarily accelerate the training process because the irregular computation suffers from poor data locality and low parallelism [13, 39].
A common way to improve the efﬁciency of sparse models is by imposing structures on model parameters.
However, current structured models have only been shown attainable by pruning and mostly beneﬁcial to inference [13, 39, 28, 1].
Whether structured models can be obtained by sparse training and in turn accelerate the training process is still an open problem.
0  1drop  0  1grow       0 1 2shuffle    0 1 2blocking  0
Before
After active zero removed added
Some recent work has explored ﬁne-grained structures in sparse models for efﬁcient training.
For example, Hubara et al.
[17] propose a method to ﬁnd N:M transposable sparsity that can beneﬁt both the forward and backward pass.
However, their method relies on specialized hardware (sparse tensor cores).
Chen et al.
[3] test the lottery ticket hypothesis [9] for sparse models with shufﬂed block structures.
Although they show promising results that such semi-structured models can be trained to high accuracy from scratch, how to generate a winning ticket without dense training remains unclear.
In this work, we take a step forward and investigate the existence of ﬁne-grained block structures in sparse models with dynamic sparse training.
Our idea is illustrated in Fig.
We ﬁrst initialize a sparse model and group the parameters in shufﬂed blocks.
Since the shufﬂed blocks and their transposes can be efﬁciently processed by hardware, our method accelerates both the forward and backward pass.
In every few iterations, we adapt the model by dropping the smallest weights and adding back the same number of new parameters.
Different from previous work which selects new parameters based on dense gradients [8] or dense weights [18], we select blocks of new parameters directly based on the input value and output gradient of each layer, making our algorithm purely sparse.
The adapted model is shufﬂed and blocked again to ensure that the active weights are always organized in shufﬂed blocks.
We test our algorithm on a variety of networks and datasets and answer the following questions:
Can a randomly initialized sparse model with ﬁne-grained block structures be trained to high accuracy?
We train a ResNet18 and a ResNet50 of 0 .
5 sparsity with parameters constrained in shufﬂed blocks of size 16 ⇥ 16 on the ImageNet dataset, and we recover the accuracy of dense models.
When the sparsity increases to 0 .
75 , the accuracy drops slightly but still matches the accuracy of nonstructured models with dynamic sparse training.
The result is further validated with various models on the CIFAR10/100 dataset.
Does the model adaptation of dynamic sparse training help improve the accuracy of sparse models with ﬁned-grained block structures?
We train sparse models with parameters in shufﬂed blocks with and without model adaptation.
For ResNet18 and ResNet50 on ImageNet, the dynamic approach achieves more than 1% higher accuracy than static training.
For different models on CIFAR10/100, dynamic training also achieves higher accuracy than static training, and the improvement is more noticeable when the sparsity is higher.
What is the tradeoff between accuracy and hardware efﬁciency with different block sizes?
We train sparse models with different block sizes on CIFAR10/100.
As expected, models with smaller block structures achieve higher accuracy.
We ﬁnd that the accuracy difference is small for block size 4, 8, 16 and is more noticeable for block size 32.
The sparse operations on blocks of size 16 ⇥ 16 can be more than 3x faster than on blocks of size 4 ⇥ 4 , but the performance improvement is small when the block size increases to 32.
This suggests that block size 16 can achieve a good tradeoff between accuracy and efﬁciency.
To the best of our knowledge, our work is the ﬁrst to apply shufﬂed blocking of model parameters to dynamic sparse training.
We believe that our ﬁndings can inspire future research on designing hardware-efﬁcient sparse training algorithms.
Related
Neural network pruning.
The conﬂict between the ever-increasing size of neural networks and the limited computing capacity of hardware has inspired a lot of research on neural network pruning.
Early works include [11, 10, 42, 35], in which iterative and heuristic pruning methods with limited and nonstructured sparsity were proposed.
These works focus on reducing the size of neural networks but do not consider their performance carefully.
The performance issue with nonstructured pruning is later identiﬁed in [40], where the authors propose to prune the entire input or output channels so that they can be easily accelerated on CPU and GPU.
The downside of these structured pruning methods is that they suffer from notable accuracy loss when sparsity increases.
To combine the beneﬁts of structured and nonstructured pruning, hybrid pruning strategies have been proposed.
Ma et al.
propose pattern-based kernel pruning [29] for convolutional neural networks.
The convolution kernels can only be pruned to one of several pre-deﬁned patterns so that the pruned model contains some regular structures.
Zhu et al.
propose vector-wise pruning [45].
They divide the weight matrix into ﬁxed-size vectors and keep the top-K elements in each vector.
In addition to vector-wise pruning, the authors also adapt the tensor core architecture on GPUs to sparse tensor core in order to support the efﬁcient execution of their vector-wise pruned models.
Rumi et al.
propose the ﬁrst ﬁne-grained block-structured pruning method for convolutional neural networks [36].
They reorder the convolution kernel matrix with a hypergraph partitioning procedure and group the nonzero weights into a number of shufﬂed blocks.
The authors also provide a CUDA implementation of the shufﬂed-block-based convolution operation, and show its performance advantage over nonstructured sparse models and dense models on GPU.
Recently, the shufﬂed blocking technique is extended for exploiting dense tensor cores on GPU to further accelerate the computation [16].
Some structured pruning methods can also reduce the computation in each iteration of the training process.
However, these methods usually require much more training iterations to recover accuracy.
For example,
Hrank [24] prunes the convolution layers one-by-one and needs to ﬁne-tune the model for 30 epochs after every layer is pruned.
ABCPruner [25] performs end-to-end ﬁne-tuning, but it needs to run the entire training process for multiple cycles.
ResRep [7] avoids ﬁne-tuning, but it needs to add extra "compactor" parameters for selecting the channels, resulting in extra computation in each iteration.
Moreover, these structured pruning methods require a pre-trained model as input.
Sparse training.
Pruning generates sparse models that are fast for inference; however, it does not accelerate the training process.
In fact, pruning often needs to be conducted iteratively and increases the training overhead.
As neural networks keep growing larger and training becomes prohibitively expensive, people start to be interested in exploiting sparsity to reduce the training cost.
Mocanu et al.
propose the ﬁrst dynamic sparse training algorithm, Sparse Evolutionary Training (SET), which starts training with a randomly initialized sparse model, prunes weights with small magnitude and adds back weights at random periodically [30].
Following SET, Bellec et al.
propose Deep Rewiring which augments traditional stochastic gradient descent with a random walk in parameter space [2].
Mostafa et al.
propose Dynamic Sparse Reparameterization (DSR) which allows the parameter budget to shift between different layers [31].
Dettmers et al.
propose Sparse Networks from Scratch (SNFS) which uses the momentum of each parameter as the criterion for growing weights [5].
Evci et al.
point out that all these algorithms can be expressed with a drop-and-grow framework and propose RigL which uses the gradient magnitude as the grow criterion [8].
An issue with RigL is that it needs to compute the full gradient every few iterations, which may not be affordable for large models.
Top-kast addresses this issue by maintaining all the weights but computing gradients on a small subset of weights in each iteration.
Although these sparse training algorithms show promising results, they still have an accuracy gap with pruning methods.
Recently, researchers show that this gap can be closed by using a small batchsize and training for more iterations [26] or by interleaving sparse training with dense training phases [34].
Besides dynamic sparse training, there has been a line of research on static sparse training.
Lee et al.
propose Single-Shot Network Pruning (SNIP) which ﬁnds an initial mask with one-shot pruning and keeps the sparse model ﬁxed during training [23].
In a follow-up paper, they reﬁne their initialization method from a signal propagation perspective [22].
Frankle et al.
propose the Lottery Ticket Hypothesis (LTH) which states that a sparse model can be trained to match the accuracy of the dense model if a good initialization is provided [9].
While appealing, the effectiveness of LTH in reducing the training cost is still unclear as obtaining the "winning ticket" is usually nontrivial [44, 9].
Most of the existing sparse training algorithms (either static or dynamic) focus on nonstructured sparse models.
Although they reduce the number of ﬂoating-point operations, these models cannot run efﬁciently on hardware due to their irregular memory access patterns and limited parallelism.
This puts the practicality of sparse training algorithms into doubt.
Inspired by the hybrid pruning methods, researchers have worked on extracting ﬁne-grained structures in sparse models to accelerate training.
Hubara et al.
propose a method to ﬁnd N:M transposable sparsity [17].
Their method accelerates both the forward and backward pass in training and can be used in either a static or dynamic fashion.
However, the N:M sparsity relies on specialized hardware (sparse tensor cores on GPU) to deliver good performance and thus is not generally applicable to other hardware or even older GPUs.
Chen et al.
[3] adopt the shufﬂed blocking technique from [36] and show that a sparse model with shufﬂed block structures can be trained to dense accuracy at a nontrivial sparsity ratio.
However, they need to pre-train a dense model to obtain an initialization for the sparse training.
Dietrich et al.
[6] propose a dynamic sparse training algorithm with block sparsity.
However, as we show in §5 and experiments, block sparsity without row/column shufﬂing has a very low mask diversity and achieves low model accuracy.
Preliminaries
Notation.
Throughout the paper, we use W 2 R R ⇥ C to denote the parameters in a layer.
Here, R is the output channel size; C is the input channel size (if it is a linear layer), or the input channel size times the convolution kernel size (if it is a convolutional layer).
We use In 2 R C ⇥ B to denote the input and Out 2 R R ⇥ B to denote the output, where B is the batch size.
The gradient w.r.t.
W is dW and has the same dimension as W .
The element at row i and column j of dW is denoted as dW ij .
The gradient w.r.t the output is dOut , and the gradient w.r.t.
the input is dIn .
We use || · || to denote the ` 1 norm of a matrix and use | · | to denote the absolute value of each element in a matrix or vector.
[Figure]
Model adaptation
Input : Network: f W , Dataset: D , Sparsity distribution:
S =
{ s 1 , . . .
, s l } , Update schedule:  T , T end W
Randomly sparsify model parameters using S ; for each training step t do Sample a batch B t ⇠ D ; L t =
P i 2 B t L ( f W ( x i ) , y i ) ; if t mod  T =
= 0 and t < T end then for each layer l do k
Decide drop number;
I drop
Select k weights to drop;
I grow
Select k new weights to grow; W
Update active weights using I drop and
I grow ; else W = W  ↵ r W L t ;
Dynamic sparse training.
Fig.
3 illustrates the idea of dynamic sparse training.
The algorithm starts with a randomly initialized sparse model.
It adapts the model in every  T iterations until iteration
T end .
The model adaptation involves two operations: drop and grow .
The drop operation discards some weights in the model, and the grow operation adds some new weights back.
Different sparse training algorithms use different criteria for the two operations.
For example, RigL [8] drops the weights with the smallest magnitudes and grows the same amount of weights with the largest gradients.
Top-kast [18] maintains a superset of active weights and drops and grows active weights according to their magnitudes.
With the periodic model adaptation, dynamic sparse training algorithms can explore more parameters and achieve higher accuracy than static sparse training [27, 26].
Computations in sparse neural networks.
Many operations in a dense neural network can be expressed as general matrix multiplications (GEMM).
In a sparse model, the computations are sparse matrix multiplications.
For example, the training of a linear layer involves three matrix multiplications: 1) in the forward pass, the input is multiplied with the sparse weight matrix to get the output (i.e., Out = W · In ); 2) in the backward pass, the gradient of loss w.r.t.
the input is computed by multiplying the transpose of the sparse weight matrix with the gradient w.r.t.
the
Input : W [ S 1 ][ S 2 ], rowidx [ S 1 ], colidx [ S 2 ], In [ C ][ B ]
Output: Out [ R ][ B ] In
[ S 2 ][ B ] =
In [ colidx ][ B ] ;
Out
[ S 1 ][ B ] = 0 ; for i = 0 to S 1  1 do for j = 0 to S 1 do for k = 0 to B  1 do
Out
[ i ][ k ] += W
[ i ][ j ] ⇤
In 0 [ j ][ k ] ;
Out [ rowidx ][ B ] =
Out 0 ; input : rowidx [ S 1 ], colidx [ S 2 ], dOut [ R ][ B ], In [ C ][ B ] output: dW [ S 1 ]
[ S 2 ]
In 0 [ S 2 ][ B ] =
In [ colidx ][ B ] ; dOut
[ S 1 ][ B ] = dOut [ rowidx ][ B ] ; dW 0 [ S 1 ][ S 2 ] = 0 ; for i = 0 to S 1  1 do for j = 0 to S 1 do for k = 0 to B  1 do dW
[ i ][ j ] += dOut
[ i ][ k ] ⇤
In 0 [ j ][ k ] ; dW [ rowidx ][ colidx ] = dW 0 ; output (i.e., dIn = W T · dOut ); 3) the gradient of loss w.r.t.
the weight is computed by multiplying the gradient w.r.t.
the output with the transpose of input (i.e., dW = dOut · In T ).
The ﬁrst two operations are sparse-matrix dense-matrix multiplication (SpMM).
The last operation is a sampled dense-dense matrix multiplication (SDDMM) because only the gradients for the active weights need to be computed.
Sparse matrix multiplication with shufﬂed blocks.
SpMM and SDDMM were ﬁrst studied for efﬁcient execution on GPU in the scientiﬁc computing domain [15, 19].
They show that the performance bottleneck of the two operations is the irregular access to the input dense matrix, and they propose adaptive tiling and reordering techniques to improve the memory access efﬁciency.
While they target large matrices with high sparsity, the reordering and tiling ideas have inspired shufﬂed blocking techniques to accelerate SpMM and SDDMM for smaller matrices with relatively lower sparsity in sparse neural networks [36, 16, 3].
The idea is illustrated in Fig.
The nonzeros in the sparse matrix are organized into small blocks.
Algorithm 2 shows the implementation of the forward pass for a shufﬂed block of parameters.
Suppose the row indices of the block are rowidx and column indices are colidx .
The program ﬁrst loads In [ colidx ][ B ] into cache In 0 and allocates a cache Out 0 for Out [ rowidx ][ B ] .
[Figure]
Since the values of In [ colidx ][ B ] and
Out [ rowidx ][ B ] are repeatedly accessed in the computation and accessing cache is much faster than main memory, the memory access overhead is reduced.
Similarly, Algorithm 3 shows the computation of the gradients for a block of parameters.
The program ﬁrst loads In [ colidx ][ B ] and dOut [ rowidx ][ B ] into cache and reuses the data in cache throughout the computation.
The larger the block size, the more data reuse the program has.
The computation can be further accelerated by exploiting the matrix-multiply units in emerging hardware (e.g., tensor cores on Nvidia GPU)
This is why structured sparse models can deliver better performance than nonstructured models that have even fewer parameters.
Proposed Algorithm
As explained in the previous section, for sparse models to run efﬁciently, it is important to have all the active weights in blocks.
Previous work ensures this through pruning [36, 16] or initialization [3].
We now show that sparse models with such ﬁne-grained block structures can also be obtained by dynamic sparse training.
Fig.
1 illustrates the model adaptation steps in our algorithm.
This section details each of the steps.
Initialization.
The algorithm starts with a randomly initialized nonstructured model.
We adopt the
ERK sparsity distribution [30] to determine the sparsity for each layer.
Then, a shufﬂed blocking procedure is invoked to reshape the sparse model into shufﬂed blocks.
Shufﬂed blocking.
Suppose N is the number of active weights in a layer and the block size is set to S 1 ⇥ S 2 .
The shufﬂed blocking procedure aims to put as many active weights into d N/ ( S 1 S 2 ) e blocks.
As shown in Fig.
4, it ﬁrst reorders the rows of the sparse matrix so that similar rows are close to each other.
The similarity between two rows is deﬁned as the Jaccard similarity of the nonzero columns in the two rows (i.e., the number of common columns divided by the total number of distinct columns).
Intuitively, two rows with more common columns have a larger similarity and should be put close to each other.
With the deﬁnition of similarity between rows, the shufﬂing of rows can be achieved by a simple clustering procedure.
We use the AgglomerativeClustering from sklearn [33] for this task.
Since row reordering is performed repeatedly during the training process, we want to keep its overhead as small as possible.
We observe that most rows have small similarities with other rows, and reordering them does not improve much of the block density.
Therefore, we only cluster the rows that have more than q nonzeros in common columns with other rows, where q ( 0 . 2 in our experiments) is a hyperparameter that controls the clustering overhead.
After the rows are reordered, we divide them into groups of size S 1 .
In Fig.
4, S 1 is set to 2, and thus the four rows are divided into two groups: {0,2} and {1,3}.
Within each row group, we select S 2 columns to form a block of size S 1 ⇥ S 2 .
To ensure large weights are preserved, we always select the column with the largest ` 1 norm in each row group and return the block with the largest ` 1 norm among all row groups.
This procedure continues until we ﬁnd enough blocks.
[Figure]
Block-aware drop criterion.
The sparse model is adapted in every  T iterations.
For each layer, we compute the number of weights to be dropped in iteration t with a cosine decaying function:
[Equation]
where ↵ is the initial drop ratio and N is the number of active weights in the layer.
Following previous dynamic sparse training algorithms [8, 26, 27, 18], we select the k weights with the smallest magnitudes (i.e. ArgTopK (  | W | , k ) ) for removal.
Before removing them, we check if these weights make up at least half of the elements in either a row or a column of a block.
More speciﬁcally, in each shufﬂed block B , we have the weights to be dropped as B \ ArgTopK (  | W | , k ) .
For a weight w 2 B \ ArgTopK (  | W | , k ) , the elements in the same row of
B is B [ rowidx ( w ) , :] , and the elements in the same column of B is B
[: , colidx ( w )] .
We remove w if and only if B [ rowidx ( w ) , :] \ ArgTopK (  | W | , k ) has at least S 1 / 2 elements or B [: , colidx ( w )] \ ArgTopK (  | W | , k ) has at least S 2 / 2 elements.
If a weight in ArgTopK (  | W | , k ) does not have enough other weights to be dropped in the same row/column of a block, the weight will not be dropped.
This is because even if we drop it, the weight is very likely to be added back after the blocking procedure.
However, because the number of weights to be added in the grow phase is determined by the number of dropped weights, this ineffective dropping of some small weights will result in more new weights being added in the grow phase and may cause more important weights to be discarded by the following shufﬂed blocking procedure.
Our block-aware drop criterion reduces the negative impact of such ineffective drops.
Intuitively, when the small weights are concentrated in a few rows or columns, we perform more model adaptation.
When the small weights are scattered in many rows and columns, we perform less aggressive adaptation.
Fig.shows an example of the block-aware drop criterion.
Suppose there are ﬁve small weights (marked with diagonal stripes) in a block.
Instead of dropping all the ﬁve weights, we only drop three of them that are in a row or column with at least two small weights (marked in blue).
The two weights with only themselves to be dropped in their rows and columns (marked in orange) will remain active.
In our experiments, we ﬁnd that this block-aware drop criterion is critical for achieving good accuracy for shufﬂed-block dynamic sparse training (See Appendix A.2).
Block-wise grow criterion.
The main novelty of our method lies in how we grow the weights.
Different from previous work that computes the gradients for all weights and adds the weights with the largest gradients[8, 26, 27], we use dOut and In to estimate the importance of weights.
Speciﬁcally, since dW = dOut · In T , dW ij = h dOut i , In j i , we have || dW ij ||  || dOut
i || · ||
In j || where dOut i represents row
i of dOut and In j represents row j of the input.
This suggests that the weights connecting the input channels of larger magnitudes to the output channels of larger gradients are of greater importance.
Our grow criterion is designed based on this intuition.
Suppose
X represents a group of S x rows in dOut and Y represents a group of S y rows in In .
It is easy to see that XY T is a shufﬂed block of size S x ⇥ S y in dW .
Our goal is to ﬁnd the blocks of weights with
[Table]
the largest || XY T || .
Without computing the actual gradients, we use || X || · || Y T || as an estimate of gradient magnitude.
We ﬁrst sort rows of dOut and In according to their norms in descending order and divide sorted dOut into groups of S x rows and sorted In into groups of S y .
Then, we compute ||
X || · || Y T || for the shufﬂed blocks and select the blocks with largest || X ||
· || Y T
|| .
For a dOut of size x ⇥ n and an input In of size y ⇥ n , the FLOPS for computing the norms of rows of dOut and
In is xn
+ yn , and the FLOPS for computing the sparse gradients is no more than kn where k is the number of dropped weights in Formula (1).
So the total FLOPS with our proposed grow criterion is xn
+ yn
+ kn .
In comparison, the FLOPS for computing dense gradient is xyn .
Since k is small, our block-wise grow criterion has a much smaller time complexity than dense gradient computation.
2 1 0    bbbbbbbbcurrent frontier !"
#$ %& !
Fig.
6 shows an example of the selection procedure with S x =and S y = 1 .
Suppose the sorted rows of input are 2, 1, 0, 3, and the sorted rows of dOut are 0, 2, 3, 1.
The selection starts from the block in the top left corner ( b 0 ) as it has the largest || X || · || Y T || .
With b 0 selected, the two blocks ( b 1 and b 4 ) adjacent to b 0 are added to the frontier, and the larger one between the two will be selected in the next step.
The selected block may overlap with the existing active blocks.
Only the non-overlapping parameters are considered new parameters.
We keep selecting new blocks and updating the frontier until new parameters exceed the drop number.
In the example of
Fig.
1, a block of new parameters ([0,2] and [2,2]) are added.
Finally, we reorganize the new model into shufﬂed blocks.
To prevent the new parameters from being discarded by the blocking procedure, we set their value to the maximum value of the active weights.
After the new blocks are formed, we reset the value of the new parameters to zero so that they will not affect the output of the network.
Comparison with Transposable N:M Sparsity
Our work is closely related to the line of research that explores ﬁne-grained N:M sparsity in neural networks [43, 17, 38].
While most work focuses on accelerating inference, a recent paper shows that transposable N:M sparsity can be found in sparse models to accelerate both the forward and backward pass in training [17].
This section compares transposable N:M sparsity with shufﬂed block sparsity.
In [17], the authors propose a measure called mask diversity to quantify how much a speciﬁc structure constrains the model.
Speciﬁcally, mask diversity is the total number of distinct sparse masks that satisfy certain sparsity and dimension constraints.
They show a positive correlation between the mask diversity and accuracy.
We follow their arguments and compute the mask diversity of shufﬂed block sparsity.
For a matrix of size R ⇥ C , suppose R and C are divisible by the block size S 1 and S 2 , and the sparsity is s .
Since the rows are reordered and divided into disjoint groups, the number of ways to choose row groups is R ! /
( S 1 !)
R/S 1 .
Once the row groups are ﬁxed, we need to select a total of
RC (1  s ) /S
2 columns from all row groups to form the desired number of blocks.
This leads to the following formula for the mask diversity of shufﬂed block sparsity:
[Equation]
In comparison, the mask diversity of nonshufﬂed block sparsity is simply choosing RC (1  s ) /
[Equation]
S 1 S
[Equation]
blocks from the total RC/ ( S 1 S 2 ) blocks.
Table 1 lists the mask diversity of shufﬂed block sparsity
[Figure] [Figure]
u
[Figure]
c y
[Figure]
[Figure]
along with other sparsity constraints for an 8 ⇥ 8 matrix and a 16 ⇥ 16 matrix.
The formula for other sparsity constraints can be found in [17].
We can see that nonstructured sparsity has the highest diversity, which explains its superior accuracy to structured sparsity.
When the block size is small, shufﬂed block sparsity has a mask diversity close to nonstructured sparsity.
For example, for the 8 ⇥ 8 matrix, when S 1 = S 2 = R/ 4 = 2 , the mask diversity is of the same order of magnitude as nonstructured sparsity.
When the block size increases, the mask diversity decreases.
This illustrates a tradeoff between hardware efﬁciency and model accuracy for shufﬂed block sparsity.
Compared to
N:M sparsity which only beneﬁts the forward pass, shufﬂed block sparsity beneﬁts both the forward and backward pass and yet has a higher mask diversity when S 1 = S 2  4 and M = 2 .
Compared to
N:M transposable sparsity, shufﬂed block sparsity has a higher mask diversity when S 1 = S 2 and M  4 .
Moreover, N:M sparsity relies on sparse-matrix-multiply units in emerging hardware (sparse tensor cores on Nvidia Ampere GPU) to deliver good performance, whereas shufﬂed blocking accelerates the computation on traditional computing units and can exploit dense-matrix-multiply units which are more broadly supported by hardware (e.g. dense tensor cores on GPU, AMX on Intel CPU, MMA on IBM Power10 processor).
As the current Nvidia GPU only supports sparse tensor cores with
N = 2 , M = 4 and dense tensor cores with S 1 = S 2 = 4 , shufﬂed blocking is expected to be a more ﬂexible and more efﬁcient way to accelerate sparse training.
We can also see that shufﬂed block sparsity has a much higher mask diversity than nonshufﬂed block sparsity.
This is because the shufﬂing of rows and columns signiﬁcantly increases the number of ways to form nonzero blocks.
Experiments
We test our algorithm with ResNet18 and ResNet50 [12] on the ImageNet dataset [4], and
WideResNet22-2, ResNet18 and VGG16 [37] on CIFAR10 [20] and CIFAR100 dataset [21].
For our
[Table]
[Figure] [Figure]
ImageNet, the models are trained for 100 epochs with batch size 256.
The learning rate schedule starts with a linear warm up reaching its maximum value of 0.1 at epoch 5 which is then dropped by a factor of 10 at epochs 30, 70, 90.
For CIFAR10/100, the models are trained for 160 epochs with batch size 128.
The learning rate is initialized to 0.1 and is decreased by a factor of 10 at epoch 80, 120.
The standard SGD with momentum 0.9 is used as the optimizer.
For dynamic sparse training, the initial drop ratio ↵ is set to 0.3,  T is set to 1600 for ImageNet and 100 for CIFAR10/100, and
T end is set to 400000 for ImageNet and 75000 for CIFAR10/100.
Accuracy compared with other training methods.
Table 2 lists the top-1 accuracy of ResNet and ResNet50 trained with different methods on ImageNet.
Our dynamic shufﬂed block training (DSB) recovers the accuracy of the dense models at 0.5 sparsity.
In comparison, the structured sparse lottery ticket (SSLT) [3] has a slight accuracy drop for ResNet50 at 0.51 sparsity while requiring a dense pre-training for initializing the sparse mask.
The accuracy of DSB at sparsity 0.75 is also higher than that of SSLT at sparsity 0.7.
Fig.
7 shows the test accuracy of WideResNet22-2, ResNet and VGG16 trained with different methods on CIFAR10.
All the experiments are run ﬁve times.
The accuracies plotted in the ﬁgures are the average values.
The difference in accuracy between different runs is smaller than 0.05%.
We can see that our DSB recovers dense accuracy at 0.5 and 0.
sparsity for all three models.
In contrast, SSLT has a noticeable accuracy drop for WideResNet and ResNet18 (although they can achieve dense accuracy at lower sparsities [3]).
Compared with nonstructured dynamic sparse training (RigL [8]), our DSB has slightly lower accuracy at 0.5 and 0.
sparsity, and the accuracy gap widens as the sparsity gets higher.
Fig.
8 shows the test accuracy on
CIFAR100.
Again, our DSB recovers dense accuracy at 0.5 and 0.6 sparsity.
For WideResNet22and ResNet18, our DSB achieves slightly higher accuracy than SSLT at small sparsity (  0 . 7 ) and almost the same accuracy at higher sparsity ( > 0 . 7 ).
For VGG16, SSLT achieves surprisingly higher accuracy than both dense training and our DSB, but they require a dense pre-training for initialization.
For all the test cases, nonshufﬂed block sparsity (DNSB) achieves apparently lower accuracies than other methods, which is consistent with the mask diversity shown in Table 1.
Importance of model adaptation.
The results in Table 2 and Fig.
7 and 8 also show that the model adaptation is critical to the accuracy of sparse training.
For ResNet50 and ResNet18 on ImageNet, without model adaptation, static shufﬂed block training (SSB) cannot recover dense accuracy at 0.5 sparsity.
Its accuracy at 0.75 sparsity is more than 1% lower than DSB.
For WideResNet22-2,
ResNet18, and VGG16 on CIFAR10/100, SSB also has lower accuracy than DSB, and the accuracy gap is larger at higher sparsity.
Tradeoff between accuracy and efﬁciency with different block sizes.
To show the effects of block size on accuracy, we run DSB with different block sizes.
The results are shown in Fig.
DSB-4,8, achieve almost the same accuracy when sparsity  0 . 8 .
DSB-4 shows a slight advantage over
DSB-8,16 at 0.9 sparsity.
When the block size increases to 32, we observe an apparent accuracy drop, and the drop is more signiﬁcant at higher sparsity.
To show the effects of block size on performance, we test the implementation of shufﬂed block convolution from [36] on an Nvidia RTX3090 GPU.
Fig.
10 shows the average speedups of shufﬂed block convolution with different block sizes over dense convolution.
We can see that when the block size is small, the sparse convolution has almost no performance advantage over dense convolution, even at a high sparsity.
As the block size increases, the shufﬂed block convolution achieves better performance.
Notably, when the block size increases from 8 to 16, the sparse computation achieves the largest marginal speedups.
When the block size increases from 16 to 32, the performance improvement is small.
This is because the convolution operation with 16 ⇥ 16 blocks already saturates the cache on the GPU platform.
The results suggest that block size 16 achieves a good tradeoff between the accuracy and efﬁciency of DSB.
[Figure]
Overhead of shufﬂed blocking.
Although shufﬂed blocking is an expensive operation, it is performed only once in every  T iterations.
Fig.
11 shows the overhead of shufﬂed blocking.
The execution time of shufﬂed blocking accounts for 1% to 45% of total training time.
We can see that, as the sparsity increases, the overhead of shufﬂed blocking decreases.
As the block size increases, the number of blocks decreases, and the overhead also decreases.
When block size is set to 16, the shufﬂed blocking procedure takes about 10% of the total training time.
According to the speedups in Fig.
10, the overhead can be justiﬁed by the performance improvements brought by sparse computation.
Conclusion
We propose a dynamic sparse training algorithm that extracts and exploits ﬁne-grained block structures in sparse models.
Our algorithm is designed based on the drop-and-grow model adaptation framework and features a block-aware drop criterion and a block-wise grow criterion.
We show that a randomly initialized sparse model with shufﬂed block structures can be trained to high accuracy, and the training process can be effectively accelerated.