PLATON: Pruning Large Transformer Models with Upper Confidence Bound of Weight Importance
Large Transformer-based models have exhibited superior performance in various
natural language processing and computer vision tasks. However, these models
contain enormous amounts of parameters, which restrict their deployment to
real-world applications. To reduce the model size, researchers prune these
models based on the weights' importance scores. However, such scores are
usually estimated on mini-batches during training, which incurs large
variability/uncertainty due to mini-batch sampling and complicated training
dynamics. As a result, some crucial weights could be pruned by commonly used
pruning methods because of such uncertainty, which makes training unstable and
hurts generalization. To resolve this issue, we propose PLATON, which captures
the uncertainty of importance scores by upper confidence bound (UCB) of
importance estimation. In particular, for the weights with low importance
scores but high uncertainty, PLATON tends to retain them and explores their
capacity. We conduct extensive experiments with several Transformer-based
models on natural language understanding, question answering and image
classification to validate the effectiveness of PLATON. Results demonstrate
that PLATON manifests notable improvement under different sparsity levels. Our
code is publicly available at https://github.com/QingruZhang/PLATON.