Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks

Unstructured pruning reduces the memory footprint in deep neural networks\n(DNNs). Recently, researchers proposed different types of structural pruning\nintending to reduce also the computation complexity. In this work, we first\nsuggest a new measure called mask-diversity which correlates with the expected\naccuracy of the different types of structural pruning. We focus on the recently\nsuggested N:M fine-grained block sparsity mask, in which for each block of M\nweights, we have at least N zeros. While N:M fine-grained block sparsity allows\nacceleration in actual modern hardware, it can be used only to accelerate the\ninference phase. In order to allow for similar accelerations in the training\nphase, we suggest a novel transposable fine-grained sparsity mask, where the\nsame mask can be used for both forward and backward passes. Our transposable\nmask guarantees that both the weight matrix and its transpose follow the same\nsparsity pattern; thus, the matrix multiplication required for passing the\nerror backward can also be accelerated. We formulate the problem of finding the\noptimal transposable-mask as a minimum-cost flow problem. Additionally, to\nspeed up the minimum-cost flow computation, we also introduce a fast\nlinear-time approximation that can be used when the masks dynamically change\nduring training. Our experiments suggest a 2x speed-up in the matrix\nmultiplications with no accuracy degradation over vision and language models.\nFinally, to solve the problem of switching between different structure\nconstraints, we suggest a method to convert a pre-trained model with\nunstructured sparsity to an N:M fine-grained block sparsity model with little\nto no training. A reference implementation can be found at\nhttps://github.com/papers-submission/structured_transposable_masks.\n