ADMM Based Semi-Structured Pattern Pruning Framework for Transformer

NLP (natural language processing) has achieved great success through the transformer model. However, the model has hundreds of millions or billions of parameters, which is a huge burden when deployed on personal computers or small-scale servers. To deal with it, we either make the model's weight matrix relatively sparser or compress the attention layer. Pattern pruning, one of the most important pruning methods, permits selecting and pruning a fixed number of parameters in each divided pattern block. However, the effect of pattern pruning is strictly limited by the sparsity within a region of weights in each layer. In this paper, we first introduced the Alternating Direction Method of Multipliers (ADMM) based pattern pruning framework to reshape the distribution of the activation map. Specifically, we propose to formulate the pattern pruning on the transformer as a constrained optimization and use ADMM to optimize the problem. In this way, the initial dense feature maps are transformed into rather regionally sparsified ones. Therefore, we can then achieve a higher compression ratio with better performance based on the pattern pruning method. Additionally, this paper provides theoretical derivations of the ADMM with local sparsity. Finally, we also extend the proposed ADMM-based framework with Sparse-Refined Straight-Through Estimator (SR-STE) to demonstrate its generalization and to avoid the gradient vanishing problem. We conduct extensive experiments on classification tasks over GLUE datasets. Significantly, we achieve a 50% compression ratio while maintaining an overall score of 80.1 on the GLUE dataset.