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 describes a ADMM-based approach to NxM sparsification in large-scale language models (e.g. BERT). Semi-structured sparsification, such as the 4:2 sparsity implementation of Ampere GPUs, allows for a best of both worlds: the weak constraint (picking M weights out of a contiguous N) enables sufficient expressivity to not overly degrade accuracy, while still allowing for an efficient hardware implementation. The authors compare against a prior approach (ASP) on a range of GLUE benchmarks (MNLI, SST-2, QNLI, CoLA, STS-B, MRPC, and RTE), outperforming ASP by 2.9 points on average for tasks with small datasets (and comparably on those with larger datasets). The authors also provide insight into the fine-tuning accuracy, and show that the approach is complementary to a (recent but not quite state-of-the-art) knowledge distillation method.