ADMM Based Semi-Structured Pattern Pruning Framework For Transformer

TL;DR

This paper introduces an ADMM-based pattern pruning framework for transformers that reshapes activation maps to enable higher compression ratios without significant loss in performance.

Contribution

It formulates pattern pruning as a constrained optimization problem solved by ADMM, improving sparsity and compression in transformer models.

Findings

Achieved 50% compression ratio on GLUE with 80.1 score

Reshaped activation maps for higher sparsity and efficiency

Extended framework with SR-STE to improve generalization

Abstract

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