Minimum Variance Unbiased N:M Sparsity for the Neural Gradients

TL;DR

The paper introduces a principled MVUE-based framework for pruning neural gradients with N:M sparsity, showing that unbiased minimum-variance pruning outperforms traditional MSE-based approaches for gradients. It derives exact MVUE masks for 1:2 sparsity and an approximately optimal method for 2:4 sparsity, demonstrating that 1:2 pruning often suffices for training while 2:4 serves as a fallback. By combining MVUE-based gradient pruning with transposable sparse weight pruning, the work envisions accelerating all training GEMMs with minimal degradation, potentially achieving up to x2 speedups. It also discusses architectural and inference implications, including when pruning both matrices is beneficial primarily for bandwidth and how activation pruning with quantization can further boost inference efficiency.

Abstract

In deep learning, fine-grained N:M sparsity reduces the data footprint and bandwidth of a General Matrix multiply (GEMM) up to x2, and doubles throughput by skipping computation of zero values. So far, it was mainly only used to prune weights to accelerate the forward and backward phases. We examine how this method can be used also for the neural gradients (i.e., loss gradients with respect to the intermediate neural layer outputs). To this end, we first establish a tensor-level optimality criteria. Previous works aimed to minimize the mean-square-error (MSE) of each pruned block. We show that while minimization of the MSE works fine for pruning the weights and activations, it catastrophically fails for the neural gradients. Instead, we show that accurate pruning of the neural gradients requires an unbiased minimum-variance pruning mask. We design such specialized masks, and find that in most cases, 1:2 sparsity is sufficient for training, and 2:4 sparsity is usually enough when this is not the case. Further, we suggest combining several such methods together in order to potentially speed up training even more.

Figures (4)

• Figure 1: Fine-grained 1:2 Sparsity for blocks located on the first quarter of the unit circle. The blocks $[a_1,a_2]$ (represented by red dots) are sampled 100 times each, and then averaged (green dots) using one of three methods: (a) greedy is the traditional method that generates the block $[0,a_2]$ if $a_1\leq a_2$, or $[a_1,0]$ otherwise. In this method, all 100 samples are the same for each block, resulting in a biased average. (b) unbiased - each block $[a_1,a_2]$ is equally likely to be pruned to $[2a_1, 0]$ or $[0,2a_2]$. Although the average of the 100 samples is unbiased, it does not have minimum variance. (c) Our unbiased method with minimum variance (Equation \ref{['our_method_main']}), has a smaller spread here than in (b).
• Figure 2: Ratio between the variance (\ref{['variance']}) of the approx-MVUE 2:4 and MVUE 2:4 when scanning (step size 0.005 all possible values of a block $[a_1,a_2,a_3,1]$, where $0 \leq a_1,a_2,a_3 \leq 1$. Notice that ratio is bounded below 2. The maximum is achieved near the left edge, when $a_4 \textcolor{black}{\gg} \max(a_1,a_2,a_3)$.
• Figure 3: Top-1 accuracy on ResNet18 Cifar10 dataset with 4:8 fine-grained sparsity on the activations. ReLU + greedy refers to applying ReLU and then the structured sparsity, while greedy does not include the ReLU function. As can be seen, applying only N:M structured sparsity as the activation function leads to accuracy degradation.
• Figure 4: Top-1 validation accuracy of ResNet18 over ImageNet dataset with the proposed approx-MVUE 2:4 and the amortization of the random samples every 50 iteration to reduce the overhead. Notice that both methods achieved similar accuracy.