Channel Permutations for N:M Sparsity

1 Introduction . Deep Neural Networks ( DNNs ) excel at complex tasks such as image classification , object detection , and language modeling . Their success often comes at the cost of immense computational complexity , and growing model size generally results in even higher-quality results . To combat this everincreasing computational workload at inference , various approaches have been used to save effort ; for example , network pruning removes weights from the network , and quantization can replace floating-point operations with simpler integer operations . We focus on one type of network pruning that is accelerated in hardware : 2:4 structured sparsity [ 22 ] ( a particular form of the more general N : M sparsity ) , since it has been shown to maintain accuracy across a wide range of networks and tasks while being easy to save computation . While many networks do recover accuracy with a fine-tuning step , some small , parameter-efficient image-classification networks that were designed with economy of memory and computation in mind can not recover the accuracy lost due to pruning with this suggested workflow . Further , this workflow repeats the dense training all over again , which , while straightforward , can be costly if the lifetime of the deployment model is not long enough to amortize this extra training . We seek to address these shortcomings with the observation that the 2:4 pruning step , which forces two out of every four consecutive elements to be zero ( the `` 2:4 constraint '' ) , sometimes must prune a relatively large , and intuitively important , value . If different options were available to each group of four values , then more advantageous choices can be made , as illustrated in Figure 1 . When the example weight matrix ( top-left ) is pruned with 2:4 sparsity in its original order , it results in a sparse matrix ( top-right ) with total weight magnitude of 44.5 . By first changing the order of the columns in the matrix , though ( bottom-left ) , the resulting sparse matrix ( bottom-right ) has a total weight magnitude of 58.6 . It is simple to find this fruitful column permutation in such a small example , but 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the complexity of the problem explodes with problem size . Further , once a permutation is found , applying it to some layer of a DNN can result in a performance overhead at run-time [ 14 ] . In this work , we give solutions to both these problems and show several benefits of using permutations for N : M sparsity . While we focus on a particular variant of N : M structured sparsity for practical reasons ( acceleration with readily-available hardware [ 26 ] ) , the techniques and observations herein apply equally to all N and M. Similarly , our permutation search strategies can be used to increase or decrease any metric of interest , not just magnitude . We note that while weight magnitude is not proven to be the optimal pruning metric , and that increasing magnitude is not guaranteed to improve accuracy , we show it to work well for the cases studied herein . Our contributions include ( 1 ) a method of permuting input channels of weight matrices without changing the output of the network or incurring a performance overhead , ( 2 ) an efficient algorithm for generating unique permutations for N : M structured sparse matrices , ( 3 ) two methods to overcome a shortcoming of greedy permutation generation algorithms , ( 4 ) searching for and applying permutations to improve accuracy of 2:4 sparse networks , and ( 5 ) transparently transforming unstructured sparse to structured sparse matrices . The rest of this paper is as follows : Section 2 discusses previous work in the area and provides background information on matrix permutations , and Section 3 presents our method of to permuting weights of a neural network with no run-time overhead . In Section 4 , we present a quality metric , an intractable exhaustive solution to find a high-quality permutation , a greedy solution , and improvements to the baseline greedy algorithm . We share an ablation study , quality metric/accuracy correlation , accuracy improvements , and results of transforming unstructured to structured sparsity in Section 5 , along with a brief discussion of the runtime needed to search for permutations . Finally , we offer a discussion and areas for future work in Section 6 . 2 Related Work and Background . Network pruning . Removing weights from a network is used to reduce both the storage and bandwidth required for the network as well as the computational power required to deploy the network . The area of network pruning is multi-dimensional : one can prune to different amounts , with different sorts of structure or regularity , using weight magnitude or other metrics to choose which weights to make zero [ 23 , 16 ] , pruning gradually [ 6 , 33 ] or in one shot , keeping the network the same size or growing it to accommodate for the sparsity [ 5 ] , etc . Let us focus on structure and regularity ; this dimension is the most important for this work and is generally composable with the others . Early work in pruning DNNs used unstructured sparsity , removing individual weight values [ 6 ] , but research soon moved on to pruning at a coarser granularity [ 34 , 1 , 19 , 18 , 15 ] : at the cost of accuracy loss for the same amount of sparsity [ 21 , 38 ] , removing entire channels or filters results in a smaller , dense workload which is easy to accelerate . At the extreme , sparsity can be in the form of blocks [ 5 ] or entire layers [ 4 ] . Simply put , network pruning has been a balancing act between maintaining accuracy and inducing enough sparsity to impart a useful benefit , either in storage requirements or the runtime of the resulting model . N : M structured sparsity . Recently-released hardware accelerates matrix-multiply-accumulate instructions if one operand satisfies a 2:4 sparsity constraint [ 22 ] , a particular form of N : M sparsity in which at least N out of every M ( contiguous , aligned ) elements are zero . In particular , this hardware requires the sparsity to be applied to the input channel dimension , C , of convolutional layers : weights along this dimension are divided into groups of four weights , and two of them are forced to zero . A workflow to generate networks that satisfy this 2:4 constraint , using a form of Learning Rate Rewinding [ 28 ] , has been empirically verified to maintain accuracy across a wide range of networks and tasks ( with the exception of the aforementioned small , parameter-efficient networks ) . Hubara et al . report a technique for finding good transposable masks for use in both forwards and backwards passes [ 11 ] , as well as two techniques for improving accuracy when fine-tuning N : M sparse networks for inference , but they are orthogonal to this work . Though we focus on fine-tuning an existing dense network , there has been work in training N : M networks from scratch [ 37 ] , but results for 2:4 sparsity fall short of the dense baseline . Matrix permutations . Applying a permutation to a matrix changes the order of the individual rows or columns that make up the matrix , and can be specified as a vector P of indices , with no repetitions , from 0 to the size of the dimension ( -1 ) . These indices are used to permute the matrix by moving the column ( or row ) at index j to position P [ j ] . A matrix multiplication ’ s operands can be permuted in two ways without materially affecting the results : permuting the inner dimension or the outer dimension of the operands . Simultaneously permuting the inner dimension of both operands changes the order in which the partial sums of the dot-product are computed , but not the eventual sum in each position of the output matrix . An example of this appears in the right side of Figure 2 : Layer i ’ s input channel , or C , dimension , common to both operands , is permuted ( numbers 1 and 2 ) by permutation Pi , and the output is unaffected . On the other hand , permuting rows of the first operand ( number 3 on the left side of Figure 2 ) or the columns of the second operand , either outer dimension , affects the position of the corresponding rows or columns in the output matrix . For this reason , when rows are permuted in the weight matrix of a neural network , it is necessary to apply the inverse permutation to the output to restore the original order . Since we are concerned with the order of only the C dimension , which defines each 2:4 group , permutations will be applied to the inner dimension of GEMM operations in high-performance math libraries [ 3 ] . Herein , permutations will apply to columns of a matrix ’ s weights or channels of a convolution tensor ’ s weights interchangeably . Each layer ( or , as discussed in Section 3 , group of layers ) can have a unique permutation , Pi ( for layer i ) . These properties have been used to cluster large values together when pruning a network with group sparsity , such as blocks or vectors [ 14 ] . The approach in this prior work has two shortcomings : it ( 1 ) incurs a runtime performance penalty and ( 2 ) can not escape from poor local minima . Permutations have also been used to make unstructured sparsity easier to accelerate on systolic arrays [ 2 ] . 3 Applying Permutations . 3.1 Permutations Across Layers . To apply a desirable permutation to some layer i ’ s C dimension , we can permute its weights just once , offline . To avoid an explicit operation to apply the same permutation to its input activations , we exploit the property that permuting weights along the output channel , or K , dimension also permutes output activations along the K dimension , and the output channel dimension of one layer ’ s activations is the input channel dimension of those same activations when used as the input to the next layer . Consider permutation Pi desired by layer i : that layer ’ s weights will be permuted along Ci , as indicated by the ‘ 1 ’ on the right side of Figure 2 . In order to keep the output constant , its activations need the same permutation applied along Ci ( ‘ 2 ’ ) . We effect this input activation permutation with another offline step : permuting the previous layer ’ s weights along Ki−1 ( ‘ 3 ’ ) . By permuting the K dimension of layer i − 1 ’ s weights to reorder its output activation ’ s K dimension , we have prepared layer i ’ s input activations as required by Pi without an explicit permutation operation . Since each layer needs to permute its own weights along only one dimension and can “ absorb '' an orthogonal permutation in the other dimension without conflict , we incur no performance overhead at run-time . We note that this simplification could also be applied to past work to eliminate the reported performance overhead due to each layer ’ s desired permutation being applied in multiple dimensions [ 14 ] , potentially at the expense of final network accuracy . So , each layer may see two permutations : one that is for itself , applied along its C dimension , and one that is desired by a downstream layer ( or layers , below ) , applied along its K dimension . For example : if following layer i + 1 has its own desired C permutation , layer i would be responsible for permuting the activations by permuting its own weights ’ K dimension . This second permutation is orthogonal to and will not conflict with its own C dimension permutation . Similarly , layer i− 1 can apply a permutation to its C dimension without affecting the K permutation for layer i .

Summary: Recent nVidia chips support hardware acceleration of structured sparsity patterns in weight matrices — specifically 2:4 sparsity where consecutive groupings of 4 entries have half of their entries set to zero.  The naive way to exploit this acceleration in deep learning is to prune weight matrices so as to meet this sparsity requirement, but doing this naively can be suboptimal if some groups contain many “important” weights and other groups contain all zeros already — it would be better if the “important” weights had a better distribution across groups.  So this paper proposes a way to permute the weights such that the pruning yields less loss.  The main idea is that if we can permute the weight rows/cols appropriately, then we can force pruning to give something that has more weight magnitude. The authors introduce permutation efficacy (and show on Imagenet experiments that this makes sense and acts as a kind of proxy for quality) and present a (mostly) greedy algorithm for optimizing this measure.  Final experiments on a number of Imagenet trained architectures show accuracy improvements over the naive method.