Taisuke Yasuda, Kyriakos Axiotis, Gang Fu, MohammadHossein Bateni, Vahab Mirrokni

Neural network pruning is a technique for improving the efficiency and generalisation ability of neural networks by replacing dense parameter tensors with sparse approximations.

With traditional sparsification techniques, we are not guaranteed any kind of structure in the resulting sparse matrix, but most inference hardware is unable to efficiently utilise unstructured sparsity in computations, so while in principle fewer FLOPs are required, in practice these efficiency gains cannot be realised.

In contrast, block sparsification involves masking out contiguous blocks in parameter tensors, so that parallel computing primitives such as CUDA thread blocks can be mapped directly onto the resulting sparse block structure, leading to much more tangible performance gains.

As the title suggests, this paper introduces a framework for combining differentiable pruning methods with combinatorial optimisation algorithms, where the former method is used for determining important entries in weight matrices, and the latter for iteratively constructing the block sparse matrix based on these importance scores.

Using this approach, the authors show that a wide variety of differentiable pruning techniques can be viewed as nonconvex regularisers that generalise the group LASSO, and that a wide class of such regularisers have unique solutions that coincide with the solution of a corresponding group LASSO problem.

Furthermore, based on this approach, the authors propose a new block pruning algorithm called SequentialAttention++, which fuses the Sequential Attention differentiable pruning technique with an algorithm called ACDC, which alternates between dense and sparse training phases to sparsify the weight matrices while allowing the sparse support to vary throughout training in case it is chosen sub-optimally early in training.

In the authors’ experiments, SequentialAttention++ consistently outperforms a number of other sparsification methods on both the ImageNet and Criterio datasets in terms of both validation loss and accuracy, across sparsities ranging from 90% to 99%.

Moreover, SequentialAttention++ performs particularly well against to the baselines for large block sizes and high sparsities, and this is the regime in which hardware efficiency improvements are the most pronounced, making this approach very attractive for real-world scenarios where model size or inference latency are critical.