Learning Best Combination for Efficient N:M Sparsity
By forcing at most N out of M consecutive weights to be non-zero, the recent
N:M network sparsity has received increasing attention for its two attractive
advantages: 1) Promising performance at a high sparsity. 2) Significant
speedups on NVIDIA A100 GPUs. Recent studies require an expensive pre-training
phase or a heavy dense-gradient computation. In this paper, we show that the
N:M learning can be naturally characterized as a combinatorial problem which
searches for the best combination candidate within a finite collection.
Motivated by this characteristic, we solve N:M sparsity in an efficient
divide-and-conquer manner. First, we divide the weight vector into
$C_{\text{M}}^{\text{N}}$ combination subsets of a fixed size N. Then, we
conquer the combinatorial problem by assigning each combination a learnable
score that is jointly optimized with its associate weights. We prove that the
introduced scoring mechanism can well model the relative importance between
combination subsets. And by gradually removing low-scored subsets, N:M
fine-grained sparsity can be efficiently optimized during the normal training
phase. Comprehensive experiments demonstrate that our learning best combination
(LBC) performs consistently better than off-the-shelf N:M sparsity methods
across various networks. Our code is released at
\url{https://github.com/zyxxmu/LBC}.