NxMTransformer: Semi-Structured Sparsification for Natural Language Understanding via ADMM

1 Introduction . Large-scale Transformer networks have achieved remarkable success for a wide variety of natural language tasks , including natural language inferencing , sentiment analysis , question answering , and others . The state-of-the-art of these NLP models employs a transfer learning paradigm which contains two stages : a semi-supervised pre-training stage that trains a masked language modeling on massive web text , followed by a fine-tuning stage where the pre-trained model is adapted to specific 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . downstream tasks with much smaller datasets . The size of these language models has dramatically increased in recent years ; even relatively small models [ 6 , 27 ] consist of hundreds of millions of parameters while larger models [ 23 , 25 , 24 ] stretch well into multi-billions . The large model size brings challenges for both deployment and training costs . While training a large-scale model often requires significant time even on large training clusters , the trained model also incurs significant challenges in deployment due to latency and capacity constraints . These challenges motivate techniques to compress and accelerate these models , even in datacenter environments where hardware resource limitations are at their smallest . These techniques include but are not limited to model quantization [ 35 , 28 , 3 ] , low rank decomposition [ 19 ] , knowledge distillation [ 11 , 27 ] , and model sparsification [ 4 , 9 ] . These compression techniques can often be combined to maximize performance gains [ 37 , 19 , 10 ] . Among different techniques , sparsification attempts to identify parameters that can be removed from the model without significantly compromising model accuracy . Sparsification techniques typically fall under two broad categories : unstructured and structured . Unstructured techniques will remove the individual parameters based on their importance ( e.g. , weight magnitude ) , which often yield the best accuracy but are unfriendly to modern hardware . Structured sparsification techniques remove parameters in groups ( e.g. , entire rows or columns ) , which result in models that retain their dense structure but can also add constraints that limit the expressiveness of the model . Recently , hardware manufacturers introduced support for NxM semi-structured sparsity to provide the benefits of both structured and unstructured sparsity . In NxM semi-structured sparsity , a model may preserve M parameters from each contiguous group of N original parameters . This relatively weak constraint on sparsification allows for sparse representations similar in flexibility to those of unstructured approaches but also permits efficient hardware implementation as well . Consider the 4x2 semi-structured sparsity implementation found on NVIDIA GPUs based on the Ampere architecture [ 1 ] . The Ampere architecture introduces a small set of multiplexers that select values from the input matrix corresponding to the retained values in the weight matrix [ 21 ] . The output of this operation remains compatible with the efficient Tensor Cores for dense matrix-matrix operations . While the Ampere GPUs are the first to market with this capability , the matrix multiplication accelerators within them are similar to those used by other accelerators [ 15 ] which should enable other vendors to provide support for this type of sparsity as well . To induce semi-structured sparsity , a small set of approaches have been offered . ASP [ 21 ] proposes training the dense network until convergence , using single-shot magnitude-based pruning to induce sparsity conformant to the NxM constraints , and repeating the original training procedure to recover accuracy . Zhou , et al . [ 36 ] uses sparse-refined straight-through estimator ( SR-STE ) to introduce the sparsity throughout the entire training process . Both of these techniques pose some challenges for large-scale pre-trained Transformer networks in particular . ASP can require a second costly sparse pre-train of the model and the single-shot magnitude-based pruning might hurt the knowledge transferrability to different downstream tasks . SR-STE on the other hand sparsifies from the random model initialization , which avoids the costly sparse retraining but also necessitates performing the pre-training process with only a sparse representation in mind . Since NxM sparse hardware is not yet ubiquitous , maintaining compatibility with a single dense pretrained representation is valuable so the costly pre-training process does not need to be performed for both sparse and dense networks . In this work , we propose to effectively induce NxM semi-structured sparsity for large-scale Transformer networks to leverage newly released Sparse Tensor Core hardware by making the following contributions . ( 1 ) We introduce a principled method , NxMTransformer ( See Figure 1 ) , to obtain Transformer networks with NxM semi-structured sparsity for different downstream tasks using Alternating Direction Method of Multipliers ( ADMM ) , a technique designed for large-scale non-convex optimization problems with constraints . Such a method allows us to alternate promoting the NxM sparsity of the network and optimizing the fine-tuning performance . ( 2 ) We conduct comprehensive experiments and demonstrate that NxMTransformer achieves 1.7 points higher accuracy than state-ofthe-art techniques to introduce NxM sparsity for natural language processing tasks . ( 3 ) We perform detailed analysis on our approach and shed light on how ADMM affects fine-tuning accuracy for downstream tasks . ( 4 ) Finally , we show that NxMTransformer is complimentary to alternative model compression techniques such as knowledge distillation and can further sparsify distilled models while preserving accuracy . 2 Background and Related Work . Compression Techniques for NLP Models . Due to the prominence of large-scale language models , there has been significant interest in compressing these models . Sparsification has been shown to be a promising approach to improving model inference speed and reducing memory cost earlier in computer vision [ 10 ] . Sparsifying pre-trained language models turns out to be a challenging task because the semi-supervised models tend to have less intrinsic sparsity than CNN-based models . Some prior studies try to understand and apply sparsification for large-scale language model compression . Chen , et al . [ 4 ] extend the Lottery Ticket Hypothesis [ 7 ] to pre-trained models , finding that winning tickets for the pre-training tasks do transfer universally to downstream tasks . Gordon , et al . [ 8 ] find that the level of sparsity discovered during the pretraining process represents most of the sparsity that can be discovered in the model . Guo , et al . [ 9 ] show that a proximal pruning strategy achieves higher accuracy than competing lasso regularization methods and iterative magnitude pruning . Most of these studies still focus on unstructured sparsity , which encounters difficulty to obtain large speedups on modern hardware . As a structured technique , Michel , et al . [ 20 ] observe that many transformer heads themselves are redundant and can be pruned with minimal accuracy loss . While a structured technique , this technique provides limited acceleration benefits since much of the compute time for short-to-medium sequence lengths is spent in the intermediate layers rather than the attention heads themselves . In addition to parameter sparsification , BERT models can also realize compression gains with model quantization [ 35 , 28 ] and knowledge distillation [ 27 , 29 , 30 , 14 ] . These techniques are complimentary to our proposed method and can be combined to achieve more effective overall compression , as in [ 10 ] . Semi-structured Sparsity . There have been few studies to induce NxM sparsity for DNN models . Among them , SR-STE [ 36 ] is the most comparable work to this paper . SR-STE is single stage compression technique that uses an extended Straight-through Estimator and a sparse-refined term to improve the induction of NxM sparsity in the network . SR-STE is demonstrated across image and neural machine translation tasks and shows improved accuracy over NVIDIA ’ s ASP [ 21 ] . Unlike this work , SR-STE is designed for training from random initialization , not a pretrained model . From a performance perspective , cuSPARSElt [ 21 ] ( the NVIDIA linear algebra library for NxM sparsity ) demonstrated 1.3X to 1.6X performance uplift for the matrix multiplications prevalent in BERT models using sparse NxM Tensor cores on Nvidia A100 GPUs . Higher speedups may be achieved by increasing batch size , sequence length , or model dimensionality with a maximum of 1.8X improvement for fp16 and 1.9X for int8 . Realized speedups are smaller than the theoretical 2X improvement of 50 % sparsity primarily due to memory bandwidth constraints . The performance provided by cuSPARSElt is a reasonable expectation for what can be achieved with NxMTransformer . Models including NxM sparsity may be deployed for serving using optimized frameworks like TensorRT [ 2 ] , which include optimized layer operators that take advantage of the hardware ’ s features . Moving forward , inference-driven machine learning compilers like TVM [ 5 ] can also introduce support for the semi-structured sparse Tensor operators using the same mechanisms as were used to provide support for dense Tensor cores , which use a similar API . ADMM for Neural Networks . Alternating Direction Method of Multipliers ( ADMM ) has been used in previous works primarily for compressing convolutional neural networks . Ye , et al . [ 34 ] explore an iterative approach to ADMM that shortcuts full ADMM convergence with a masked retraining operation . Ren , et al . [ 26 ] use ADMM for both unstructured sparsification and quantization of the model while combining the technique with a heuristic for determining whether speedup will be achieved given the achieved compression of a layer . Ma , et al . [ 18 ] use domain knowledge of CNN filter structure to perform a structured pruning operation for mobile-hardware efficient compression . While the above techniques all use ADMM as the sparsifying mechanism , none examine in-depth applicability to pre-trained language models for NxM sparsity . 3 Methodology . In this section , we formulate ADMM for NxM sparsity , describe the specific aspects of pre-trained models that define the optimization pipeline , and describe the high-level optimization schedule . 3.1 Problem Definition .. Consider adapting a pre-trained large-scale language model with L Transformer layers ( e.g. , BERT ) to natural language understanding tasks , such as sentiment analysis , entailment , question and answering , etc , where the training instances are inputs ( often text phrases ) and target pairs : { xi , yi } Ni=1 . Assume the collection of pre-trained model weights is W 0 = { W0i } Li=11 , the goal of NxMTransformer is to load W 0 and fine-tune it to W 0 such that { W 0i } Li=1 satisfies the constraints of at most M weight parameters having non-zero values out of N consecutive weights , while achieving similar performance in comparison to fine-tuning the task-specific objective function f ( { Wi } Li=1 ) ( e.g. , cross-entropy for classification ) using W0 but without constraints . 3.2 NxMTransformer . Different from conventional DNN training objectives , the above problem is non-convex with combinatorial constraints . Therefore , it can not be directly solved by gradient-based methods such as stochastic gradient descent . To address this issue , we adopt the alternating direction method of multipliers ( ADMM ) , which is a reliable method for large-scale constrained optimization ( e.g. , with combinatorial constraints ) . In particular , we modify the objective function of the NxM sparsity problem as min { Wi } f ( { Wi } Li=1 ) + LX i=1 gi ( Zi ) subject to Wi = Zi , i = 1 , . . . , L ( 1 ) where f ( · ) is the fine-tuning objective function , and gi ( · ) is an added penalty function and Zi are auxiliary variables . To apply ADMM , we define the penalty function as gi ( Wi ) = ⇢ 0 if Wi 2 Si 1 otherwise ( 2 ) 1For convenience ’ s sake , we will omit the notation of the bias since it is not relevant to the task of sparsification . where Si represents the constraint set Si = { Wi that have at most M weight parameters having non-zero values out of N consecutive weights } . Choice of Si . Not all weights in pre-trained Transformer models need to satisfy this NxM sparsity constraint . BERT and similar pretrained models typically consist of three major components : embeddings , Transformer blocks , and classifiers ( See Figure 1 ) . For NxM semi-structured sparsity , we solely consider weights in Transformer blocks . Take BERT as an example , each Transformer block consists of 6 fully connected sub-layers : the query Q , key K , value V layers , the attention output layer Attn. , and two feed-forward network FFN1 and FFN2 . Each of the fully connected layers can take advantage of NxM semi-structured sparsity ; furthermore , these layers constitute the vast majority of inference wall-time . Of the six fully connected layers , FFN1 and FFN2 are particularly important to sparsify , alone requiring more than half of the inference wall-time for a Transformer block . Note that attention head pruning techniques [ 20 ] are unable to sparsify FFN1 and FFN2 . The self-attention mechanism itself does not include any trained parameters and is unable to be sparsified using this technique . We exclude the embedding layer since the lookup operation associated with the embedding layers is incompatible for acceleration with Tensor Cores . The classifier is composed of fully connected layers . For a given task , the classifier weights are randomly initialized at the beginning of the fine-tuning process . We find that sparsifying these matrices using ADMM will unnecessarily harm accuracy . Since the execution of these layers is typically under 2 % of the inference wall-time , the runtime cost is minimal for doing so . Decomposing the minimization problem into sub-problems . Once we define Si , we apply the augmented Lagrangian , which decomposes equation 1 into two sub-problems on W and Z : Sub-problem 1 ( performance-promoting ) : min { Wi } f ( { Wi } Li=1 ) + LX i=1 ⇢ 2 kWi Zki +Uki k2F ( 3 ) Sub-problem 2 ( NxM sparsity-promoting ) : min { Zi } LX i=1 gi ( Zi ) + LX i=1 ⇢ 2 kWk+1i Zi +U k i k2F ( 4 ) The first sub-problem solves the performance promoting problem , which consists of two terms . The first term is the standard objective function for fine-tuning the task , and the second term is a L2 regularization term . The regularization target Zki Uki is dynamically updated , based on Uki , which is the dual variable ( i.e. , the Lagrange multiplier ) . Since the L2 term is convex , the complexity of solving sub-problem 1 ( e.g. , via ADAM [ 16 ] ) is the same as minimizing f ( · ) . The second subproblem solves the sparsity promoting problem . Since it optimizes the sparse constraints separately , it can be solved analytically as the solution of the Euclidean projection of Wk+1i + Uk onto our constraint . For the case of NxM semi-structured sparsity , this is accomplished by retaining the M largest values of the contiguous group of N values ( See Figure 1 ) , which can be solved in linear time . Finally , we need to update the dual variable U as Uki : = U k 1 i +W k i Zki to guarantee that the dual feasibility condition is satisfied in each ADMM iteration . Sparsity-inducing based fine-tuning . The typical ADMM pipeline fully trains a model to convergence before introducing ADMM [ 26 ] . This two-step process is necessary since the primary objective of ADMM is to optimize the existing network to conform to the constraints ; ADMM will only introduce small changes to parameters it retains in the model . However , for pre-trained language models , the primary purpose of the fine-tune is to adapt the classifiers for the specific downstream with minimal disturbance to the parameters of the pre-trained representation . As a result , we apply the three aforementioned steps ( i.e. , two sub-problems and the update of the dual variable ) while fine-tuning the model . In particular , we perform the three steps in an alternating manner , i.e. , performing some number of fine-tuning steps with Adam to solve the first sub-problem , solving the Euclidean projection for each weight matrix for the second sub-problem , and finally updating the auxiliary variable . This sequence will be referred to as one ADMM iteration . The optimization proceeds until the W and Z variables have converged , at which point we have a sparsified network compliant with our NxM constraint .

Summary: This paper introduces a new method called NxMTransformer to induce NxM semi-structured sparsity on pertained language models, for recent hardware dedicated for efficient NxM sparsity (e.g. newly released Sparse Tensor Core). NxMTransformer fromulates the NxM sparsity as a constrained optimization problem and use ADMM to alternately promote the NxM sparsity constraint and finetuning performance. Experiments on GLUE datasets show that the proposed approach can beat the best NxM sparsity method by 1.7 points.