Pruning Redundant Mappings in Transformer Models via Spectral-Normalized Identity Prior (2020)

Abstract: Traditional (unstructured) pruning methods for a Transformer model focus on regularizing the individual weights by penalizing them toward zero. In this work, we explore spectral-normalized identity priors (SNIP), a structured pruning approach which penalizes an entire residual module in a Transformer model toward an identity mapping. Our method identifies and discards unimportant non-linear mappings in the residual connections by applying a thresholding operator on the function norm, and is applicable to any structured module including a single attention head, an entire attention blocks, or a feed-forward subnetwork. Furthermore, we introduce spectral normalization to stabilize the distribution of the post-activation values of the Transformer layers, further improving the pruning effectiveness of the proposed methodology. We conduct experiments with BERT on 5 GLUE benchmark tasks to demonstrate that SNIP achieves effective pruning results while maintaining comparable performance. Specifically, we improve the performance over the state-of-the-art by 0.5 to 1.0% on average at 50% compression ratio.

## introduction natural language processing (nlp) has recently achieved great success by using the transformerbased pre-trained models @xcite @xcite . however, these models often consume considerable storage, memory bandwidth, and computational resource. to reduce the model size and increase the inference throughput, compression techniques such as knowledge distillation @xcite @xcite @xcite and weight pruning (guo knowledge distillation methods require the specification of a student network with a smaller architecture, which often has to be identified by a tedious sequence process of trial-and-error based decisions. by comparison, the iterative pruning methods gradually prune the redundant model weights or layers from the full-size model, and provide a full picture of the trade-off between the task performance and the model size with a single training process, as illustrated in figure 1 . this allows the iterative pruning methods to easily determine the most compact architecture given a required level of model performance. however, many of the existing pruning methods rely on classic regularizers that act on the individual weights by penalizing them to zero @xcite . as a result, the pruned model tends to maintain the same architecture as the original model despite the reduced parameter count, which does not practically lead to an improvement in inference latency @xcite . this leads to a question: is it possible to perform more structured pruning on a transformer model to modify its model architecture (e.g., reducing width and depth)? to this end, we notice that many previous works have suggested that the learned transformer models often have much redundancy @xcite @xcite . for example, @xcite ; @xcite found that most of the attention heads in a transformer model can be removed without significantly impacting accuracy, and @xcite found that while the earlier layers and the later layers in a bert model play clear roles in extracting either low-level or task-specific linguistic knowledge, the roles of the intermediate layers are less important. these observations have motivated the idea that transformer models may exhibit considerable structural redundancy -i.e., some layers can be removed during the training without harming the final performance. in this paper, we propose a structured pruning method that reduces the architecture of a transformer by directly targeting its sub-modules as a whole -for example, a single attention head, an attention module, a feed-forward subnetwork, etc. we take an approach we call spectral-normalized identity prior (snip), which imposes a functionlevel prior centered around the identity function on the aforementioned modules. specifically, we take the advantage of the residual blocks (f(x) + x) within a transformer layer and compress them to strict identity mappings @xcite by identifying the residual blocks whose nonlinear mapping's absolute values (|f(x)|) mostly fall below a threshold . with this strategy, the weights of the transformer model can still be under-regularized when using simple l 1 or l 2 based regularizers, which cause the distribution of the post-activation values prone to be noisy even after layer normalization @xcite . to address this issue, we further leverage spectral normalization @xcite to stabilize the distribution of the post-activation values by regularizing the largest singular value of the weight matrices. we use bert @xcite as a case study in this paper. across multuple tasks in the glue benchmark @xcite , snip improves the performance over the state-of-the-art by 0.5 to 1.0% on average at 50% compression ratio. 1 we also show that spectral normalization results in more sparse and regulated layer mappings during pre-training. we compare the remaining model components across the tasks at a fixed compression ratio in an ablation study, and show that the remaining components are similar but not identical. our contributions are three-fold: first, we introduce identity-inducing prior, a structured pruning approach that imposes identity-inducing regularization on the transformer mappings as a whole rather than its individual weights. second, we show that through a novel combination with the spectral normalization regularization, the resulting spectral-normalized identity prior (snip) leads to well-regularized weight distribution and sparse layer mappings in a bert model. finally, we conduct thorough experiments to validate the snip approach over 5 standard nlu tasks. our results suggest that different model components in a transformer play critically different roles across tasks, suggesting the importance of performing task-specific pruning to obtain an architecture that is the most suitable for the target task. ## related work pre-trained language model compression the major existing efforts to compress pre-trained language models such as bert include knowledge distillation @xcite and pruning @xcite . the knowledge distillation approach enables the transfer of knowledge from a large teacher model to a smaller student model. such attempts have been made to distill bert models, e.g., distil-bert @xcite , bert-pkd @xcite , distilled bilstm @xcite , tiny-bert @xcite , mobilebert @xcite , etc. all of these methods require carefully designing the student architecture. furthermore, to choose which intermediate results that the student model can learn from, e.g., the outputs of each layer, the attention maps, is still under discussion. similar to other pruning-based methods, our method can iteratively remove the least important weights or connections, explore the full spectrum of trade-offs, and find the best affordable architecture in one shot. many language representation model pruning methods focus on individual components of the weight matrices. for example, @xcite integrates reweighted l 1 minimization with a proximal algorithm to search sparsity patterns in the model; @xcite uses magnitude weight pruning, which compresses the model by removing weights close to 0; @xcite applies deterministic first-order weight pruning method where both weights with low and high values can be pruned. a very few works try structured weight pruning, e.g., @xcite proposes a structured pruning approach based on low-rank factorization and augmented lagrangian l 0 norm regularization. on the other hand, there also exist works that prune a coherent set of sub-modules in the transformer model. for example, @xcite and @xcite propose to prune individual attention heads either manually via head importance score, or automatically via a relaxed l 0 regularization. @xcite applies random pruning to the entire layers. in contrast, our method allows finer-grained structured pruning on transformer modules (i.e., both attention heads and feed-forward layers) and propose to improve the mathematical property of a transformer (i.e., lipschitz condition) for more effective pruning. other compression approaches include weight sharing @xcite , quantization @xcite and neural architecture search @xcite , but are not within the discussion of this paper. we refer interested readers to @xcite for further details. applications of spectral normalization spectral normalization is first proposed for generative adversarial network (gan) as a regularization technique to stabilize the discriminator training @xcite . it was later applied to improve the performance of the other types of generative neural networks @xcite , and was analyzed theoretically in the context of adversarial robustness and generalization @xcite . spectral normalization regularizes the lipschitz condition of the model mappings and is known to benefit model generalization under both the classic and the adversarial settings @xcite @xcite . in this paper, we will explore the benefit of spectral regularization for improving the effectiveness of pruning. ## methods in this section, we first briefly review the basic transformer layers in @xcite (3.1). we then introduce our identity prior into transformer's residual connections using threshold (3.2). in section 3.3, we give mathematical foundations to the spectral normalization and show how it could help with our identity prior. finally, to put it all together, we establish our structured iterative pruning methods for bert fine-tuning (3.4). ## conclusion in this work, we propose a structured pruning method for compressing transformer models, which prunes redundant mappings via spectralnormalized identity priors (snip). we achieve effective pruning results on bert fine-tuning while maintaining comparable performance. our work shows the importance of the mathematical properties of the transformer model (specifically, the lipschitz condition) on the effectiveness of pruning. additionally, we quantify task-specific trade-offs between model complexity and task performance, as well as the progressive improvement provided by multi-head attention (mha) and feedforward networks (ffn). our results show that applying pruning at the level of mappings instead of individual weights allows for better model compression, when combined with the appropriate regularization. this suggests that developing more global pruning strategies may be a fruitful avenue for future research. in the future, we plan to apply a similar approach to further reduce the width of transformer layers, i.e., the hidden dimension, to achieve an even higher compression ratio. we are also interested in jointly using the proposed approach with other compression methods.