[Paper Review] Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks

Itay Hubara, Brian Chmiel+4 more|arXiv (Cornell University)|Feb 16, 2021

Advanced Neural Network Applications25 citations

TL;DR

The paper introduces a transposable N:M fine-grained sparsity mask to accelerate both forward and backward passes, formulates mask optimization as a min-cost flow problem (with a fast 2-approximation), and demonstrates 2x training speedups with negligible accuracy loss across vision and language models, plus a method to adapt unstructured sparsity to N:M masks with AdaPrune.

ABSTRACT

Unstructured pruning reduces the memory footprint in deep neural networks (DNNs). Recently, researchers proposed different types of structural pruning intending to reduce also the computation complexity. In this work, we first suggest a new measure called mask-diversity which correlates with the expected accuracy of the different types of structural pruning. We focus on the recently suggested N:M fine-grained block sparsity mask, in which for each block of M weights, we have at least N zeros. While N:M fine-grained block sparsity allows acceleration in actual modern hardware, it can be used only to accelerate the inference phase. In order to allow for similar accelerations in the training phase, we suggest a novel transposable fine-grained sparsity mask, where the same mask can be used for both forward and backward passes. Our transposable mask guarantees that both the weight matrix and its transpose follow the same sparsity pattern; thus, the matrix multiplication required for passing the error backward can also be accelerated. We formulate the problem of finding the optimal transposable-mask as a minimum-cost flow problem. Additionally, to speed up the minimum-cost flow computation, we also introduce a fast linear-time approximation that can be used when the masks dynamically change during training. Our experiments suggest a 2x speed-up in the matrix multiplications with no accuracy degradation over vision and language models. Finally, to solve the problem of switching between different structure constraints, we suggest a method to convert a pre-trained model with unstructured sparsity to an N:M fine-grained block sparsity model with little to no training. A reference implementation can be found at https://github.com/papers-submission/structured_transposable_masks.

Motivation & Objective

• Rank and understand how different sparsity masks affect model accuracy via a new measure called mask-diversity (MD).
• Develop a transposable N:M sparsity mask that accelerates both forward and backward passes during training.
• Provide efficient algorithms to find optimal masks (min-cost flow and a fast 2-approximation) and enable mask switching without retraining (AdaPrune).
• Demonstrate 2x speedups in matrix multiplications with negligible accuracy degradation on vision and language tasks.

Proposed method

• Define and motivate mask-diversity (MD) as the number of valid masks at a given sparsity under structure constraints.
• Formulate the transposable N:M mask design as an optimization problem aiming to maximize the retained L1 norm within MxM blocks with N nonzero per row and per column.
• Reduce the optimal-mask problem to a min-cost flow instance for pretrained dense models.
• Introduce a fast 2-approximation greedy algorithm for the mask when training from scratch, with almost linear time in block size.
• Propose AdaPrune to transform unstructured sparsity into an N:M structured mask with minimal retraining.
• Provide an Adaprune-based calibration approach to switch sparsity structures with little accuracy loss.

Experimental results

Research questions

• RQ1 Can mask-diversity predict the impact of a sparsity structure on model accuracy?
• RQ2 Can we design a transposable N:M sparsity mask that accelerates both forward and backward passes?
• RQ3 What are efficient algorithms to find optimal or near-optimal transposable masks during training?
• RQ4 Can we convert unstructured sparsity to N:M sparsity with minimal retraining (AdaPrune)?

Key findings

• A 4:8 transposable mask achieves comparable accuracy to a 2:4 mask but enables training acceleration by using sparse tensor cores for both forward and backward passes.
• The proposed 2x speed-up in matrix multiplications is demonstrated with no accuracy degradation on vision and language models.
• An exact min-cost-flow formulation provides optimal transposable masks for pretrained dense models, while a 2-approximation algorithm runs in near-linear time and yields solutions within a factor of 2 of optimal.
• AdaPrune can transform unstructured sparse models into 4:8 transposable sparse models with less than 1% degradation in some settings, enabling hardware-friendly deployment without full retraining.
• Empirical results show 2x sparse-tensor-core utilization during training with comparable accuracy to prior methods that only accelerate forward passes.