Learning Best Combination for Efficient N:M Sparsity

TL;DR

This paper introduces a novel method called Learning Best Combination (LBC) for optimizing N:M sparsity in neural networks, achieving better performance and efficiency by framing the problem as a combinatorial optimization and using a learnable scoring mechanism.

Contribution

The paper proposes an efficient divide-and-conquer approach with a learnable scoring system to optimize N:M sparsity during training, outperforming existing methods.

Findings

LBC outperforms existing N:M sparsity methods across various networks.

The scoring mechanism effectively models the importance of combination subsets.

Efficient optimization during normal training phase enhances practical applicability.

Abstract

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 project is released at \url{https://github.com/zyxxmu/LBC}.