PaperlayerView exact revision on arXiv

Source-linked AI summary

Bi-directional Masks for Efficient N:M Sparse Training

Yuxin Zhang, Yiting Luo, Mingbao Lin, Yunshan Zhong, Jingjing Xie, Fei Chao, Rongrong Ji

arXiv:2302.06058v1cs.CV

TL;DR

N:M sparse training offers practical hardware acceleration, but conventional masks can leave backward propagation dense, while transposable masks constrain sparsity patterns. Bi-Mask uses separate forward and backward masks with row permutation, and the authors report better results than transposable masks and performance on par with or better than methods without backward acceleration.

  • Problem

    Vanilla N:M sparse training leaves backward propagation dense because transposition disrupts N:M blocks, while transposable masks enable backward acceleration at the cost of a constrained sparsity pattern.

    Source: §1 · p. 2, §3.2 · p. 4

  • Method

    Bi-Mask builds separate forward and backward N:M masks and permutes weight rows to increase eligible backward blocks and reduce the gradient gap.

    Source: §1 · p. 2

  • Results

    Bi-Mask performs better than transposable masks and on par with or better than methods without backward acceleration; it reaches 71.5% Top-1 accuracy for 1:16 sparse ResNet-50 on ImageNet, surpassing T-Mask by 5.3%.

    Source: §1 · p. 3, §6 · p. 8

  • Takeaways & Limitations

    Bi-Mask offers N:M sparse training with forward and backward acceleration without imposing a transposable constraint.

    Source: §1 · p. 2, §6 · p. 8

    Evaluation is limited to image classification, and random permutation does not always maximize the number of N:M blocks.

    Source: §5 · p. 8

Abstract

from arXiv · show

We focus on addressing the dense backward propagation issue for training efficiency of N:M fine-grained sparsity that preserves at most N out of M consecutive weights and achieves practical speedups supported by the N:M sparse tensor core. Therefore, we present a novel method of Bi-directional Masks (Bi-Mask) with its two central innovations in: 1) Separate sparse masks in the two directions of forward and backward propagation to obtain training acceleration. It disentangles the forward and backward weight sparsity and overcomes the very dense gradient computation. 2) An efficient weight row permutation method to maintain performance. It picks up the permutation candidate with the most eligible N:M weight blocks in the backward to minimize the gradient gap between traditional uni-directional masks and our bi-directional masks. Compared with existing uni-directional scenario that applies a transposable mask and enables backward acceleration, our Bi-Mask is experimentally demonstrated to be more superior in performance. Also, our Bi-Mask performs on par with or even better than methods that fail to achieve backward acceleration. Project of this paper is available at \url{https://github.com/zyxxmu/Bi-Mask}.

1. Introduction

N:M sparsity offers hardware acceleration while retaining fine-grained weight removal, but existing masks either leave backward propagation dense or constrain sparsity patterns. Bi-Mask separates forward and backward masks and adds row permutation; experiments report stronger performance than transposable masks and results on par with or better than methods without backward acceleration.

1. Introduction

  • N:M sparsity preserves at most N of every M consecutive weights and offers practical acceleration, but ASP requires dense training followed by pruning and sparse retraining.That training cost can limit its use when training resources are scarce.Source: §1 · p. 2
  • Vanilla N:M masks accelerate forward multiplication but leave backward propagation dense; transposable masks accelerate both directions but constrain masks to N:M patterns in rows and columns.The transposable requirement is associated with greater performance degradation, particularly at higher sparsity levels such as 1:8 and 1:16.Source: Figure 1 · p. 1, §1 · p. 2
  • Bi-Mask constructs separate row-N:M forward and column-N:M backward masks, then permutes weight rows to increase eligible backward N:M blocks and reduce the gradient gap.The method avoids imposing a transposable constraint on one mask.Source: §1 · pp. 2–3, Figure 2 · p. 2
  • 71.5% Top-1 accuracy for 1:16 sparse ResNet-50 on ImageNet surpasses T-Mask by 5.3%, while 2:4 Bi-Mask exceeds SR-STE by 0.4% Top-1 accuracy.The paper also reports performance on par with or better than vanilla N:M methods that do not accelerate backward propagation.Source: §1 · p. 3, Table 1 · p. 3

2. Related Work

Network-sparsity research spans fine-grained methods with flexible weight removal and coarse-grained methods designed for hardware, while N:M sparsity seeks both performance maintenance and practical acceleration. Sparse-training work dynamically prunes and revives weights; Bi-Mask differs from transposable masks and prior permutation work in its separate directional masks and row-permutation objective.

2.1. Network Sparsity

  • Fine-grained sparsity removes weights at arbitrary positions and can maintain performance at high sparsity, but unstructured formats yield limited hardware speedup; coarse-grained sparsity is hardware friendly but can degrade performance.Coarse-grained methods remove weight blocks or convolution filters, and may suffer severe drops above 50% sparsity.Source: §2.1 · p. 3

2.2. Sparse Training

  • Sparse training dynamically prunes and revives weights, while N:M sparsity preserves at most N of every M consecutive weights and uses sparse tensor cores for practical acceleration.Bi-Mask uses separate forward and backward masks rather than the transposable mask's single row-and-column constraint; its row permutation targets eligible backward blocks in networks trained from scratch, unlike input-channel permutation of pretrained CNNs.Source: §2.2 · p. 3, §2.1 · p. 3

3. Methodology

N:M sparsity constrains each block of M consecutive weights to at most N nonzeros, but a row-wise mask fails to accelerate backward propagation after weight transposition. Bi-Mask uses separate directional masks and row permutation to support acceleration in both directions while limiting gradient mismatch.

3.1. Revisiting N:M Sparse Training

  • N:M sparsity keeps at most N of every M consecutive row weights nonzero, enabling forward acceleration, while transposition makes the vanilla mask dense in backward propagation.The backward multiplication requires N:M blocks in rows, but the transposed mask has them in columns.Source: §3.1 · p. 3, §3.2 · p. 4

3.2. Rethinking the Transposable N:M Mask

  • Transposable masks enable backward acceleration but restrict mask flexibility and show severe performance degradation, particularly at 1:8 and 1:16.Figure 3 compares mask diversity and sparse ResNet-50 performance on ImageNet.Source: §3.2 · pp. 4–5, Figure 3 · p. 4

3.3. Bi-directional N:M Masks

  • Bi-Mask separates forward and backward sparsity: the forward mask is row-wise N:M, while the backward mask imposes N:M along columns.This design is intended to support acceleration in both directions without the transposable-mask constraint.Source: §3.3 · p. 5
  • The backward mask approximates the ideal gradient, but can eliminate gradients of some weights that are masked only in the forward direction, causing performance degradation.The method adapts the backward mask using magnitudes of forward-masked weights.Source: §3.3 · p. 5
  • Row permutation increases eligible N:M blocks for backward masking, reducing the chance of eliminating gradients of nonzero masked weights.At intervals, the method selects the best of K randomly generated permutation candidates by eligible-block count.Source: §3.3 · p. 5

4. Experimentation

Across CIFAR-10 and ImageNet models, Bi-Mask generally outperforms transposable masks while matching or exceeding methods without backward acceleration. The experiments also examine search efficiency and how permutation updating, binarization, and the backward mask affect accuracy.

4.1. Settings

  • Experiments cover CIFAR-10 ResNet-32 and MobileNet-V2, and ImageNet ResNet-18/50 and DeiT-small; comparisons include ASP and SR-STE without backward acceleration and T-Mask with it.Top-1 classification accuracy is reported on both datasets.Source: §4.1 · p. 6

4.2. Comparison on CIFAR-10

  • 94.78% Top-1 accuracy at 2:4 on CIFAR-10 ResNet-32 exceeds SR-STE by 0.10% and T-Mask by 0.26%, while Bi-Mask is comparable to vanilla N:M training.On MobileNet-V2, Bi-Mask reaches 94.28% at 1:4, versus 94.43% for the dense model and 93.81% for T-Mask.Source: §4.2 · p. 6, Table 2 · p. 6, Table 3 · p. 6

4.3. Comparison on ImageNet

  • On ImageNet ResNet-18 at 2:4, Bi-Mask gains 1.6% over T-Mask and 0.9% over ASP, which does not accelerate backward propagation..Source: §4.3 · p. 6, Table 4 · p. 7
  • 77.4% Top-1 accuracy at 2:4 and 77.5% at 4:8 are reported for Bi-Mask on ImageNet ResNet-50, which beats competitors across N:M cases with the same or better acceleration.At 1:4, Bi-Mask surpasses SR-STE by 0.3%.Source: §4.3 · p. 7, Table 5 · p. 7
  • 15.0 seconds is reported for Bi-Mask to find a permutation at 1:16 on ResNet-50, compared with about 278.4 seconds for T-Mask.The search-time comparison uses one NVIDIA Tesla A100 GPU.Source: §4.3 · p. 7, Table 6 · p. 7
  • On ImageNet DeiT-small, Bi-Mask consistently achieves the best Top-1 accuracy across N:M cases, improving over SR-STE by 1.7% and T-Mask by 5.9% at 2:4.Source: §4.3 · p. 7, Table 7 · p. 7

4.4. Performance Study

  • Updating the permutation every 100 training iterations gives the best tested accuracy; shorter intervals destabilize sparse training, while longer intervals leave the permutation outdated.Performance saturates at 100 or more permutation candidates, so a time-consuming greedy search is unnecessary.Source: §4.4 · pp. 7–8, Figure 4 · p. 8
  • Magnitude-based binarization better fits the forward mask, while random selection causes the largest accuracy drop; adding permutation after the backward mask raises ResNet-18 accuracy by 0.3% over the baseline.The backward mask alone enables backward acceleration with a 0.3% Top-1 drop; permutation increases eligible N:M blocks and reduces incorrect gradient elimination.Source: §4.4 · p. 8, Table 8 · p. 8, Table 9 · p. 8

5. Limitations

The evaluation covers only image classification, leaving effectiveness on other tasks unverified. Weight-update acceleration and a permutation method that always maximizes eligible N:M blocks remain future work.

5. Limitations

  • Bi-Mask is evaluated only on image classification, so its effectiveness on tasks such as object detection remains unverified.The authors identify evaluation on other tasks as future work.Source: §5 · p. 8
  • The work accelerates forward and backward propagation, but does not explore accelerating the weight-update phase.The authors leave weight-update acceleration for future work.Source: §5 · p. 8
  • Random permutation does not always maximize eligible N:M blocks, although the authors report no overall performance damage.They expect a better method could locate the best permutation.Source: §5 · p. 8

6. Conclusion

Bi-Mask builds separate masks for forward and backward propagation, using weight-row permutation and magnitude-based pruning for backward acceleration. Experiments report superiority over several state-of-the-art methods, including parity or better performance against methods that often lack backward acceleration.

6. Conclusion

  • Bi-Mask independently builds forward and backward masks, using weight-row permutation followed by magnitude-based pruning for backward acceleration.The method avoids imposing a transposable constraint on the N:M sparse mask.Source: §6 · p. 8
  • Experiments show Bi-Mask surpasses its competitor by a large margin under the same acceleration effects.The conclusion also reports superiority over several state-of-the-art methods.Source: §6 · p. 8
  • Bi-Mask performs on par with or better than off-the-shelf methods that often fail to achieve backward acceleration.Source: §6 · p. 8

Related papers

Loading 2302.06058v1…